Repository: gfx-rs/wgpu (language: rust)

Issue:
Naga: `Compose`s of vector `ZeroValue`s trips up constant evaluation
Take `binop()` in constant evaluator as an example. We call `eval_zero_value_and_splat()` on each operand to attempt to reduce each side to `Compose`s and `Literal`s. But if we have `vec3(vec2f(), 1.0)` this will result in a `Compose` with components `[ZeroValue, Literal]`.

Then we use `proc::flatten_compose()` to flatten nested `Compose`s. And it even works with `Splat`s. But it doesn't handle `ZeroValue`s. (And cannot, as it has an immutable reference to the expression arena and we'd need to add a new expression to remove for `ZeroValue`s.)

Because `binop()` handles Composes recursively, we _kind of_ do the right thing regardless in some circumstances:

```wgsl
var x = vec3(vec2i(), 0) + vec3(1);
```

gets evaluated to

```wgsl
var x: vec3<i32> = vec3<i32>(vec2<i32>(1i, 1i), 1i);
```

This will produce the right result at runtime, but ideally we'd have evaluated it down to a single compose, ie `vec3<i32>(1i, 1i, 1i)`.

However, if we make the right hand side more complex, things start to fall apart:

```wgsl
var x = vec3(vec2i(), 0) + vec3(0, 1, 2);
```

This should result in `vec3<i32>(0i, 1i, 2i)`, but we get `vec3<i32>(vec2<i32>(0i, 0i), 1i)`.

Let's walk through this.
* binop() gets called with lhs: `Compose(ZeroValue<vec2i>, Literal(0i))` and rhs: `Compose(Literal(0i), Literal(1i), Literal(2i))`.
* After calling `eval_zero_value_and_splat()` then `flatten_compose()` we have lhs: `[ZeroValue<vec2>, Literal(0)]` and rhs: `[Literal(0), Literal(1), Literal(2)]`.
* It `zip()`s these and recursively calls itself for each item. So first we are called with lhs: `ZeroValue<vec2i>`, rhs: `Literal(0)`. 
* After `eval_zero_value_and_splat()` this is now `Compose(Literal(0i), Literal(0i)` and `Literal(0i)`
* It thinks we are doing a Vector + Scalar which it can handle. But this is wrong: we're going to use the 0th item from the RHS to add to both the 0th and 1st items on the LHS.
* Then the next iteration from the `zip()` is `Literal(0)` + `Literal(1)`. Again this is wrong: we are adding the 1st item for the RHS to the 2nd LHS.

And if we have a compose containing a vector zero val on both sides of the equation, with the zero val in different locations, we can fail validation altogether:

```wgsl
var x = vec3(vec2i(), 2) + vec3(1, vec2i());
```

This should give `vec3<i32>(1i, 0i, 2i)`, but instead we get this:
```
error: Function [0] 'main' is invalid
  ┌─ in.wgsl:1:1
  │  
1 │ ╭ fn main() {
2 │ │     // var x = vec3(vec2i(), 0) + vec3(0, 1, 2);
3 │ │     var x = vec3(vec2i(), 2) + vec3(1, vec2i());
  │ │             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ naga::Expression [6]
  │ ╰────────────────────────────────────────────────^ naga::Function [0]
  │  
  = Expression [6] is invalid
  = Composing expects 3 components but 4 were given
```


Discussion / hints:
We can fix this by making `eval_zero_val_and_splat()` recursively call itself for each component when `expr` is a `Compose`

Fix patch:
diff --git a/naga/src/proc/constant_evaluator.rs b/naga/src/proc/constant_evaluator.rs
index fcb325f539..f4b58bc89a 100644
--- a/naga/src/proc/constant_evaluator.rs
+++ b/naga/src/proc/constant_evaluator.rs
@@ -1411,6 +1411,17 @@ impl<'a> ConstantEvaluator<'a> {
         mut expr: Handle<Expression>,
         span: Span,
     ) -> Result<Handle<Expression>, ConstantEvaluatorError> {
+        // If expr is a Compose expression, eliminate ZeroValue and Splat expressions for
+        // each of its components.
+        if let Expression::Compose { ty, ref components } = self.expressions[expr] {
+            let components = components
+                .clone()
+                .iter()
+                .map(|component| self.eval_zero_value_and_splat(*component, span))
+                .collect::<Result<_, _>>()?;
+            expr = self.register_evaluated_expr(Expression::Compose { ty, components }, span)?;
+        }
+
         // The result of the splat() for a Splat of a scalar ZeroValue is a
         // vector ZeroValue, so we must call eval_zero_value_impl() after
         // splat() in order to ensure we have no ZeroValues remaining.


Test patch:
diff --git a/naga/tests/in/const-exprs.wgsl b/naga/tests/in/const-exprs.wgsl
index ee9304ce45..5165c49ceb 100644
--- a/naga/tests/in/const-exprs.wgsl
+++ b/naga/tests/in/const-exprs.wgsl
@@ -87,3 +87,10 @@ fn compose_of_splat() {
 
 const add_vec = vec2(1.0f) + vec2(3.0f, 4.0f);
 const compare_vec = vec2(3.0f) == vec2(3.0f, 4.0f);
+
+// Ensure binary ops correctly flatten compositions of vector zero values
+fn compose_vector_zero_val_binop() {
+    var a = vec3(vec2i(), 0) + vec3(1);
+    var b = vec3(vec2i(), 0) + vec3(0, 1, 2);
+    var c = vec3(vec2i(), 2) + vec3(1, vec2i());
+}
diff --git a/naga/tests/out/glsl/const-exprs.main.Compute.glsl b/naga/tests/out/glsl/const-exprs.main.Compute.glsl
index 0b318a65e3..4b473bed7c 100644
--- a/naga/tests/out/glsl/const-exprs.main.Compute.glsl
+++ b/naga/tests/out/glsl/const-exprs.main.Compute.glsl
@@ -86,6 +86,13 @@ uint map_texture_kind(int texture_kind) {
     }
 }
 
+void compose_vector_zero_val_binop() {
+    ivec3 a = ivec3(1, 1, 1);
+    ivec3 b = ivec3(0, 1, 2);
+    ivec3 c = ivec3(1, 0, 2);
+    return;
+}
+
 void main() {
     swizzle_of_compose();
     index_of_compose();
diff --git a/naga/tests/out/hlsl/const-exprs.hlsl b/naga/tests/out/hlsl/const-exprs.hlsl
index aa2ba75ed6..9d62504f78 100644
--- a/naga/tests/out/hlsl/const-exprs.hlsl
+++ b/naga/tests/out/hlsl/const-exprs.hlsl
@@ -93,6 +93,15 @@ uint map_texture_kind(int texture_kind)
     }
 }
 
+void compose_vector_zero_val_binop()
+{
+    int3 a = int3(1, 1, 1);
+    int3 b = int3(0, 1, 2);
+    int3 c = int3(1, 0, 2);
+
+    return;
+}
+
 [numthreads(2, 3, 1)]
 void main()
 {
diff --git a/naga/tests/out/msl/const-exprs.msl b/naga/tests/out/msl/const-exprs.msl
index cb0959f72e..dc0c394868 100644
--- a/naga/tests/out/msl/const-exprs.msl
+++ b/naga/tests/out/msl/const-exprs.msl
@@ -93,6 +93,14 @@ uint map_texture_kind(
     }
 }
 
+void compose_vector_zero_val_binop(
+) {
+    metal::int3 a = metal::int3(1, 1, 1);
+    metal::int3 b = metal::int3(0, 1, 2);
+    metal::int3 c = metal::int3(1, 0, 2);
+    return;
+}
+
 kernel void main_(
 ) {
     swizzle_of_compose();
diff --git a/naga/tests/out/spv/const-exprs.spvasm b/naga/tests/out/spv/const-exprs.spvasm
index afd9fe8499..63b256d5de 100644
--- a/naga/tests/out/spv/const-exprs.spvasm
+++ b/naga/tests/out/spv/const-exprs.spvasm
@@ -1,12 +1,12 @@
 ; SPIR-V
 ; Version: 1.1
 ; Generator: rspirv
-; Bound: 109
+; Bound: 120
 OpCapability Shader
 %1 = OpExtInstImport "GLSL.std.450"
 OpMemoryModel Logical GLSL450
-OpEntryPoint GLCompute %100 "main"
-OpExecutionMode %100 LocalSize 2 3 1
+OpEntryPoint GLCompute %111 "main"
+OpExecutionMode %111 LocalSize 2 3 1
 %2 = OpTypeVoid
 %3 = OpTypeInt 32 0
 %4 = OpTypeInt 32 1
@@ -16,137 +16,151 @@ OpExecutionMode %100 LocalSize 2 3 1
 %8 = OpTypeVector %6 2
 %10 = OpTypeBool
 %9 = OpTypeVector %10 2
-%11 = OpConstant  %3  2
-%12 = OpConstant  %4  3
-%13 = OpConstant  %4  4
-%14 = OpConstant  %4  8
-%15 = OpConstant  %6  3.141
-%16 = OpConstant  %6  6.282
-%17 = OpConstant  %6  0.44444445
-%18 = OpConstant  %6  0.0
-%19 = OpConstantComposite  %7  %17 %18 %18 %18
-%20 = OpConstant  %4  0
-%21 = OpConstant  %4  1
-%22 = OpConstant  %4  2
-%23 = OpConstant  %6  4.0
-%24 = OpConstant  %6  5.0
-%25 = OpConstantComposite  %8  %23 %24
-%26 = OpConstantTrue  %10
-%27 = OpConstantFalse  %10
-%28 = OpConstantComposite  %9  %26 %27
-%31 = OpTypeFunction %2
-%32 = OpConstantComposite  %5  %13 %12 %22 %21
-%34 = OpTypePointer Function %5
-%39 = OpTypePointer Function %4
-%43 = OpConstant  %4  6
-%48 = OpConstant  %4  30
-%49 = OpConstant  %4  70
-%52 = OpConstantNull  %4
-%54 = OpConstantNull  %4
-%57 = OpConstantNull  %5
-%68 = OpConstant  %4  -4
-%69 = OpConstantComposite  %5  %68 %68 %68 %68
-%78 = OpConstant  %6  1.0
-%79 = OpConstant  %6  2.0
-%80 = OpConstantComposite  %7  %79 %78 %78 %78
-%82 = OpTypePointer Function %7
-%87 = OpTypeFunction %3 %4
-%88 = OpConstant  %3  10
-%89 = OpConstant  %3  20
-%90 = OpConstant  %3  30
-%91 = OpConstant  %3  0
-%98 = OpConstantNull  %3
-%30 = OpFunction  %2  None %31
-%29 = OpLabel
-%33 = OpVariable  %34  Function %32
-OpBranch %35
-%35 = OpLabel
-OpReturn
-OpFunctionEnd
-%37 = OpFunction  %2  None %31
+%11 = OpTypeVector %4 3
+%12 = OpConstant  %3  2
+%13 = OpConstant  %4  3
+%14 = OpConstant  %4  4
+%15 = OpConstant  %4  8
+%16 = OpConstant  %6  3.141
+%17 = OpConstant  %6  6.282
+%18 = OpConstant  %6  0.44444445
+%19 = OpConstant  %6  0.0
+%20 = OpConstantComposite  %7  %18 %19 %19 %19
+%21 = OpConstant  %4  0
+%22 = OpConstant  %4  1
+%23 = OpConstant  %4  2
+%24 = OpConstant  %6  4.0
+%25 = OpConstant  %6  5.0
+%26 = OpConstantComposite  %8  %24 %25
+%27 = OpConstantTrue  %10
+%28 = OpConstantFalse  %10
+%29 = OpConstantComposite  %9  %27 %28
+%32 = OpTypeFunction %2
+%33 = OpConstantComposite  %5  %14 %13 %23 %22
+%35 = OpTypePointer Function %5
+%40 = OpTypePointer Function %4
+%44 = OpConstant  %4  6
+%49 = OpConstant  %4  30
+%50 = OpConstant  %4  70
+%53 = OpConstantNull  %4
+%55 = OpConstantNull  %4
+%58 = OpConstantNull  %5
+%69 = OpConstant  %4  -4
+%70 = OpConstantComposite  %5  %69 %69 %69 %69
+%79 = OpConstant  %6  1.0
+%80 = OpConstant  %6  2.0
+%81 = OpConstantComposite  %7  %80 %79 %79 %79
+%83 = OpTypePointer Function %7
+%88 = OpTypeFunction %3 %4
+%89 = OpConstant  %3  10
+%90 = OpConstant  %3  20
+%91 = OpConstant  %3  30
+%92 = OpConstant  %3  0
+%99 = OpConstantNull  %3
+%102 = OpConstantComposite  %11  %22 %22 %22
+%103 = OpConstantComposite  %11  %21 %22 %23
+%104 = OpConstantComposite  %11  %22 %21 %23
+%106 = OpTypePointer Function %11
+%31 = OpFunction  %2  None %32
+%30 = OpLabel
+%34 = OpVariable  %35  Function %33
+OpBranch %36
 %36 = OpLabel
-%38 = OpVariable  %39  Function %22
-OpBranch %40
-%40 = OpLabel
 OpReturn
 OpFunctionEnd
-%42 = OpFunction  %2  None %31
+%38 = OpFunction  %2  None %32
+%37 = OpLabel
+%39 = OpVariable  %40  Function %23
+OpBranch %41
 %41 = OpLabel
-%44 = OpVariable  %39  Function %43
-OpBranch %45
-%45 = OpLabel
 OpReturn
 OpFunctionEnd
-%47 = OpFunction  %2  None %31
+%43 = OpFunction  %2  None %32
+%42 = OpLabel
+%45 = OpVariable  %40  Function %44
+OpBranch %46
 %46 = OpLabel
-%56 = OpVariable  %34  Function %57
-%51 = OpVariable  %39  Function %52
-%55 = OpVariable  %39  Function %49
-%50 = OpVariable  %39  Function %48
-%53 = OpVariable  %39  Function %54
-OpBranch %58
-%58 = OpLabel
-%59 = OpLoad  %4  %50
-OpStore %51 %59
-%60 = OpLoad  %4  %51
-OpStore %53 %60
-%61 = OpLoad  %4  %50
-%62 = OpLoad  %4  %51
-%63 = OpLoad  %4  %53
-%64 = OpLoad  %4  %55
-%65 = OpCompositeConstruct  %5  %61 %62 %63 %64
-OpStore %56 %65
 OpReturn
 OpFunctionEnd
-%67 = OpFunction  %2  None %31
-%66 = OpLabel
-%70 = OpVariable  %34  Function %69
-OpBranch %71
-%71 = OpLabel
+%48 = OpFunction  %2  None %32
+%47 = OpLabel
+%57 = OpVariable  %35  Function %58
+%52 = OpVariable  %40  Function %53
+%56 = OpVariable  %40  Function %50
+%51 = OpVariable  %40  Function %49
+%54 = OpVariable  %40  Function %55
+OpBranch %59
+%59 = OpLabel
+%60 = OpLoad  %4  %51
+OpStore %52 %60
+%61 = OpLoad  %4  %52
+OpStore %54 %61
+%62 = OpLoad  %4  %51
+%63 = OpLoad  %4  %52
+%64 = OpLoad  %4  %54
+%65 = OpLoad  %4  %56
+%66 = OpCompositeConstruct  %5  %62 %63 %64 %65
+OpStore %57 %66
 OpReturn
 OpFunctionEnd
-%73 = OpFunction  %2  None %31
+%68 = OpFunction  %2  None %32
+%67 = OpLabel
+%71 = OpVariable  %35  Function %70
+OpBranch %72
 %72 = OpLabel
-%74 = OpVariable  %34  Function %69
-OpBranch %75
-%75 = OpLabel
 OpReturn
 OpFunctionEnd
-%77 = OpFunction  %2  None %31
+%74 = OpFunction  %2  None %32
+%73 = OpLabel
+%75 = OpVariable  %35  Function %70
+OpBranch %76
 %76 = OpLabel
-%81 = OpVariable  %82  Function %80
-OpBranch %83
-%83 = OpLabel
 OpReturn
 OpFunctionEnd
-%86 = OpFunction  %3  None %87
-%85 = OpFunctionParameter  %4
+%78 = OpFunction  %2  None %32
+%77 = OpLabel
+%82 = OpVariable  %83  Function %81
+OpBranch %84
 %84 = OpLabel
-OpBranch %92
-%92 = OpLabel
-OpSelectionMerge %93 None
-OpSwitch %85 %97 0 %94 1 %95 2 %96
-%94 = OpLabel
-OpReturnValue %88
+OpReturn
+OpFunctionEnd
+%87 = OpFunction  %3  None %88
+%86 = OpFunctionParameter  %4
+%85 = OpLabel
+OpBranch %93
+%93 = OpLabel
+OpSelectionMerge %94 None
+OpSwitch %86 %98 0 %95 1 %96 2 %97
 %95 = OpLabel
 OpReturnValue %89
 %96 = OpLabel
 OpReturnValue %90
 %97 = OpLabel
 OpReturnValue %91
-%93 = OpLabel
-OpReturnValue %98
+%98 = OpLabel
+OpReturnValue %92
+%94 = OpLabel
+OpReturnValue %99
+OpFunctionEnd
+%101 = OpFunction  %2  None %32
+%100 = OpLabel
+%105 = OpVariable  %106  Function %102
+%107 = OpVariable  %106  Function %103
+%108 = OpVariable  %106  Function %104
+OpBranch %109
+%109 = OpLabel
+OpReturn
 OpFunctionEnd
-%100 = OpFunction  %2  None %31
-%99 = OpLabel
-OpBranch %101
-%101 = OpLabel
-%102 = OpFunctionCall  %2  %30
-%103 = OpFunctionCall  %2  %37
-%104 = OpFunctionCall  %2  %42
-%105 = OpFunctionCall  %2  %47
-%106 = OpFunctionCall  %2  %67
-%107 = OpFunctionCall  %2  %73
-%108 = OpFunctionCall  %2  %77
+%111 = OpFunction  %2  None %32
+%110 = OpLabel
+OpBranch %112
+%112 = OpLabel
+%113 = OpFunctionCall  %2  %31
+%114 = OpFunctionCall  %2  %38
+%115 = OpFunctionCall  %2  %43
+%116 = OpFunctionCall  %2  %48
+%117 = OpFunctionCall  %2  %68
+%118 = OpFunctionCall  %2  %74
+%119 = OpFunctionCall  %2  %78
 OpReturn
 OpFunctionEnd
\ No newline at end of file
diff --git a/naga/tests/out/wgsl/const-exprs.wgsl b/naga/tests/out/wgsl/const-exprs.wgsl
index 411e835b5d..4649807eb8 100644
--- a/naga/tests/out/wgsl/const-exprs.wgsl
+++ b/naga/tests/out/wgsl/const-exprs.wgsl
@@ -85,6 +85,14 @@ fn map_texture_kind(texture_kind: i32) -> u32 {
     }
 }
 
+fn compose_vector_zero_val_binop() {
+    var a: vec3<i32> = vec3<i32>(1i, 1i, 1i);
+    var b: vec3<i32> = vec3<i32>(0i, 1i, 2i);
+    var c: vec3<i32> = vec3<i32>(1i, 0i, 2i);
+
+    return;
+}
+
 @compute @workgroup_size(2, 3, 1) 
 fn main() {
     swizzle_of_compose();