tint/resolver: Materialize array index expression

Bug: tint:1504
Change-Id: Ib14f4aa8ed7ca9d4bc4cdf6f4acdfa7eec03b716
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/91961
Kokoro: Kokoro <noreply+kokoro@google.com>
Commit-Queue: Ben Clayton <bclayton@google.com>
Reviewed-by: Dan Sinclair <dsinclair@chromium.org>

src/tint/resolver/array_accessor_test.cc

@@ -156,19 +156,36 @@ TEST_F(ResolverIndexAccessorTest, Vector) {
     EXPECT_TRUE(ref->StoreType()->Is<sem::F32>());
 }
 
-TEST_F(ResolverIndexAccessorTest, Array) {
+TEST_F(ResolverIndexAccessorTest, Array_Literal_i32) {
-    auto* idx = Expr(2_i);
     Global("my_var", ty.array<f32, 3>(), ast::StorageClass::kPrivate);
-
+    auto* acc = IndexAccessor("my_var", 2_i);
-    auto* acc = IndexAccessor("my_var", idx);
     WrapInFunction(acc);
-
     EXPECT_TRUE(r()->Resolve()) << r()->error();
-
     ASSERT_NE(TypeOf(acc), nullptr);
-    ASSERT_TRUE(TypeOf(acc)->Is<sem::Reference>());
-
     auto* ref = TypeOf(acc)->As<sem::Reference>();
+    ASSERT_NE(ref, nullptr);
+    EXPECT_TRUE(ref->StoreType()->Is<sem::F32>());
+}
+
+TEST_F(ResolverIndexAccessorTest, Array_Literal_u32) {
+    Global("my_var", ty.array<f32, 3>(), ast::StorageClass::kPrivate);
+    auto* acc = IndexAccessor("my_var", 2_u);
+    WrapInFunction(acc);
+    EXPECT_TRUE(r()->Resolve()) << r()->error();
+    ASSERT_NE(TypeOf(acc), nullptr);
+    auto* ref = TypeOf(acc)->As<sem::Reference>();
+    ASSERT_NE(ref, nullptr);
+    EXPECT_TRUE(ref->StoreType()->Is<sem::F32>());
+}
+
+TEST_F(ResolverIndexAccessorTest, Array_Literal_AInt) {
+    Global("my_var", ty.array<f32, 3>(), ast::StorageClass::kPrivate);
+    auto* acc = IndexAccessor("my_var", 2_a);
+    WrapInFunction(acc);
+    EXPECT_TRUE(r()->Resolve()) << r()->error();
+    ASSERT_NE(TypeOf(acc), nullptr);
+    auto* ref = TypeOf(acc)->As<sem::Reference>();
+    ASSERT_NE(ref, nullptr);
     EXPECT_TRUE(ref->StoreType()->Is<sem::F32>());
 }
 
@@ -249,7 +266,7 @@ TEST_F(ResolverIndexAccessorTest, Array_Literal_I32) {
     EXPECT_TRUE(r()->Resolve()) << r()->error();
 }
 
-TEST_F(ResolverIndexAccessorTest, EXpr_Deref_FuncGoodParent) {
+TEST_F(ResolverIndexAccessorTest, Expr_Deref_FuncGoodParent) {
     // fn func(p: ptr<function, vec4<f32>>) -> f32 {
     //     let idx: u32 = u32();
     //     let x: f32 = (*p)[idx];
@@ -265,7 +282,7 @@ TEST_F(ResolverIndexAccessorTest, EXpr_Deref_FuncGoodParent) {
     EXPECT_TRUE(r()->Resolve()) << r()->error();
 }
 
-TEST_F(ResolverIndexAccessorTest, EXpr_Deref_FuncBadParent) {
+TEST_F(ResolverIndexAccessorTest, Expr_Deref_FuncBadParent) {
     // fn func(p: ptr<function, vec4<f32>>) -> f32 {
     //     let idx: u32 = u32();
     //     let x: f32 = *p[idx];

src/tint/resolver/materialize_test.cc

@@ -566,16 +566,16 @@ enum class Method {
     // let a = abstract_expr;
     kLet,
 
-    // min(abstract_expr, abstract_expr);
+    // min(abstract_expr, abstract_expr)
     kBuiltinArg,
 
-    // bitcast<f32>(abstract_expr);
+    // bitcast<f32>(abstract_expr)
     kBitcastF32Arg,
 
-    // bitcast<vec3<f32>>(abstract_expr);
+    // bitcast<vec3<f32>>(abstract_expr)
     kBitcastVec3F32Arg,
 
-    // array<i32, abstract_expr>();
+    // array<i32, abstract_expr>()
     kArrayLength,
 
     // switch (abstract_expr) {
@@ -587,7 +587,10 @@ enum class Method {
     // @workgroup_size(abstract_expr)
     // @stage(compute)
     // fn f() {}
-    kWorkgroupSize
+    kWorkgroupSize,
+
+    // arr[abstract_expr]
+    kIndex,
 };
 
 static std::ostream& operator<<(std::ostream& o, Method m) {
@@ -608,6 +611,8 @@ static std::ostream& operator<<(std::ostream& o, Method m) {
             return o << "switch";
         case Method::kWorkgroupSize:
             return o << "workgroup-size";
+        case Method::kIndex:
+            return o << "index";
     }
     return o << "<unknown>";
 }
@@ -692,6 +697,10 @@ TEST_P(MaterializeAbstractNumericToDefaultType, Test) {
             Func("f", {}, ty.void_(), {},
                  {WorkgroupSize(abstract_expr()), Stage(ast::PipelineStage::kCompute)});
             break;
+        case Method::kIndex:
+            Global("arr", ty.array<i32, 4>(), ast::StorageClass::kPrivate);
+            WrapInFunction(IndexAccessor("arr", abstract_expr()));
+            break;
     }
 
     auto check_types_and_values = [&](const sem::Expression* expr) {
@@ -756,6 +765,14 @@ constexpr Method kScalarMethods[] = {
     Method::kBitcastF32Arg,
 };
 
+/// Methods that support abstract-integer materialization
+/// Note: Doesn't contain kWorkgroupSize as @workgroup_size has tighter constraints on the range of
+///       allowed integer values.
+constexpr Method kAIntMethods[] = {
+    Method::kSwitch,
+    Method::kIndex,
+};
+
 /// Methods that support vector materialization
 constexpr Method kVectorMethods[] = {
     Method::kLet,
@@ -826,12 +843,13 @@ INSTANTIATE_TEST_SUITE_P(
                          Types<f32M, AFloatM>(AFloat(-kSubnormalF32), -kSubnormalF32),  //
                      })));
 
-INSTANTIATE_TEST_SUITE_P(MaterializeSwitch,
+INSTANTIATE_TEST_SUITE_P(MaterializeAInt,
                          MaterializeAbstractNumericToDefaultType,
                          testing::Combine(testing::Values(Expectation::kMaterialize),
-                                          testing::Values(Method::kSwitch),
+                                          testing::ValuesIn(kAIntMethods),
                                           testing::ValuesIn(std::vector<Data>{
                                               Types<i32, AInt>(0_a, 0.0),                        //
+                                              Types<i32, AInt>(10_a, 10.0),                      //
                                               Types<i32, AInt>(AInt(kHighestI32), kHighestI32),  //
                                               Types<i32, AInt>(AInt(kLowestI32), kLowestI32),    //
                                           })));
@@ -878,10 +896,10 @@ INSTANTIATE_TEST_SUITE_P(MatrixValueCannotBeRepresented,
                                               Types<f32M, AFloatM>(0.0_a, -kTooBigF32),  //
                                           })));
 
-INSTANTIATE_TEST_SUITE_P(SwitchValueCannotBeRepresented,
+INSTANTIATE_TEST_SUITE_P(AIntValueCannotBeRepresented,
                          MaterializeAbstractNumericToDefaultType,
                          testing::Combine(testing::Values(Expectation::kValueCannotBeRepresented),
-                                          testing::Values(Method::kSwitch),
+                                          testing::ValuesIn(kAIntMethods),
                                           testing::ValuesIn(std::vector<Data>{
                                               Types<i32, AInt>(0_a, kHighestI32 + 1),  //
                                               Types<i32, AInt>(0_a, kLowestI32 - 1),   //

src/tint/resolver/resolver.cc

@@ -1205,7 +1205,10 @@ bool Resolver::ShouldMaterializeArgument(const sem::Type* parameter_ty) const {
 }
 
 sem::Expression* Resolver::IndexAccessor(const ast::IndexAccessorExpression* expr) {
-    auto* idx = sem_.Get(expr->index);
+    auto* idx = Materialize(sem_.Get(expr->index));
+    if (!idx) {
+        return nullptr;
+    }
     auto* obj = sem_.Get(expr->object);
     auto* obj_raw_ty = obj->Type();
     auto* obj_ty = obj_raw_ty->UnwrapRef();