Naga: Composes of vector ZeroValues trips up constant evaluation

[jamienicol]

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 Composes and Literals. 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 Composes. And it even works with Splats. But it doesn't handle ZeroValues. (And cannot, as it has an immutable reference to the expression arena and we'd need to add a new expression to remove for ZeroValues.)

Because binop() handles Composes recursively, we kind of do the right thing regardless in some circumstances:

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

gets evaluated to

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:

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:

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

[jamienicol]

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