[wgsl] Handle override-sized array params in ProgramToIR

We have special-case handling for this when emitting variable
declarations. Pull that out to a helper function so that we can use it
in the parameter case too.

Fixed: 449592100
Change-Id: I7c242dec479a3195720784d6df0ee64643ee2a34
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/265477
Commit-Queue: James Price <jrprice@google.com>
Reviewed-by: dan sinclair <dsinclair@chromium.org>
diff --git a/src/tint/lang/wgsl/reader/program_to_ir/ir_program_test.h b/src/tint/lang/wgsl/reader/program_to_ir/ir_program_test.h
index 9d7b76b..e0681bc 100644
--- a/src/tint/lang/wgsl/reader/program_to_ir/ir_program_test.h
+++ b/src/tint/lang/wgsl/reader/program_to_ir/ir_program_test.h
@@ -79,7 +79,10 @@
     /// @returns the generated module
     Result<core::ir::Module> Build(std::string wgsl) {
         Source::File file("test.wgsl", std::move(wgsl));
-        auto result = wgsl::reader::WgslToIR(&file);
+        wgsl::reader::Options options{
+            .allowed_features = wgsl::AllowedFeatures::Everything(),
+        };
+        auto result = wgsl::reader::WgslToIR(&file, options);
         if (result != Success) {
             return result.Failure();
         }
diff --git a/src/tint/lang/wgsl/reader/program_to_ir/program_to_ir.cc b/src/tint/lang/wgsl/reader/program_to_ir/program_to_ir.cc
index fc33710..068bdd9 100644
--- a/src/tint/lang/wgsl/reader/program_to_ir/program_to_ir.cc
+++ b/src/tint/lang/wgsl/reader/program_to_ir/program_to_ir.cc
@@ -328,7 +328,7 @@
         Vector<core::ir::FunctionParam*, 1> params;
         for (auto* p : ast_func->params) {
             const auto* param_sem = program_.Sem().Get(p)->As<sem::Parameter>();
-            auto* ty = param_sem->Type()->Clone(clone_ctx_.type_ctx);
+            auto* ty = RemapOverrideSizedArrayIfNeeded(param_sem->Type());
             auto* param = builder_.FunctionParam(p->name->symbol.NameView(), ty);
 
             for (auto* attr : p->attributes) {
@@ -1228,35 +1228,7 @@
             var,
             [&](const ast::Var* v) {
                 auto* ref = sem->Type()->As<core::type::Reference>();
-                const core::type::Type* store_ty = nullptr;
-
-                const auto* ary = ref->StoreType()->As<core::type::Array>();
-                // If the array has an override count
-                if (ary && !ary->Count()
-                                ->IsAnyOf<core::type::RuntimeArrayCount,
-                                          core::type::ConstantArrayCount>()) {
-                    core::ir::Value* count = tint::Switch(
-                        ary->Count(),  //
-                        [&](const sem::UnnamedOverrideArrayCount* u) {
-                            return EmitValueExpression(u->expr->Declaration());
-                        },
-                        [&](const sem::NamedOverrideArrayCount* n) {
-                            return scopes_.Get(n->variable->Declaration()->name->symbol);
-                        },
-                        TINT_ICE_ON_NO_MATCH);
-
-                    if (!count) {
-                        return;
-                    }
-
-                    auto* ary_count =
-                        builder_.ir.Types().Get<core::ir::type::ValueArrayCount>(count);
-                    store_ty = builder_.ir.Types().Get<core::type::Array>(
-                        ary->ElemType()->Clone(clone_ctx_.type_ctx), ary_count, ary->Size());
-                } else {
-                    store_ty = ref->StoreType()->Clone(clone_ctx_.type_ctx);
-                }
-
+                auto* store_ty = RemapOverrideSizedArrayIfNeeded(ref->StoreType());
                 auto* ty = builder_.ir.Types().Get<core::type::Pointer>(ref->AddressSpace(),
                                                                         store_ty, ref->Access());
 
@@ -1375,6 +1347,38 @@
         }
         TINT_UNREACHABLE();
     }
+
+    const core::type::Type* RemapOverrideSizedArrayIfNeeded(const core::type::Type* ty) {
+        // Check that we have an override-sized array, or a pointer to one.
+        const auto* ary = ty->UnwrapPtr()->As<core::type::Array>();
+        if (!ary ||
+            !ary->Count()
+                 ->IsAnyOf<sem::NamedOverrideArrayCount, sem::UnnamedOverrideArrayCount>()) {
+            return ty->Clone(clone_ctx_.type_ctx);
+        }
+
+        // If the array has an override count, we need to remap it to a value array count.
+        core::ir::Value* count = tint::Switch(
+            ary->Count(),  //
+            [&](const sem::UnnamedOverrideArrayCount* u) {
+                return EmitValueExpression(u->expr->Declaration());
+            },
+            [&](const sem::NamedOverrideArrayCount* n) {
+                return scopes_.Get(n->variable->Declaration()->name->symbol);
+            },
+            TINT_ICE_ON_NO_MATCH);
+
+        auto* ary_count = builder_.ir.Types().Get<core::ir::type::ValueArrayCount>(count);
+        const core::type::Type* remapped_ty = builder_.ir.Types().Get<core::type::Array>(
+            ary->ElemType()->Clone(clone_ctx_.type_ctx), ary_count, ary->Size());
+
+        // If the original type was a pointer, wrap the remapped array in a pointer too.
+        if (auto* ptr = ty->As<core::type::Pointer>()) {
+            remapped_ty = builder_.ir.Types().ptr(ptr->AddressSpace(), remapped_ty, ptr->Access());
+        }
+
+        return remapped_ty;
+    }
 };
 
 }  // namespace
diff --git a/src/tint/lang/wgsl/reader/program_to_ir/program_to_ir_test.cc b/src/tint/lang/wgsl/reader/program_to_ir/program_to_ir_test.cc
index 4c6bab6..05f86da 100644
--- a/src/tint/lang/wgsl/reader/program_to_ir/program_to_ir_test.cc
+++ b/src/tint/lang/wgsl/reader/program_to_ir/program_to_ir_test.cc
@@ -1523,5 +1523,41 @@
 )");
 }
 
+TEST_F(IR_FromProgramTest, OverrideSizedArrayInPointerParameter) {
+    auto* src = R"(
+override x = 1;
+
+var<workgroup> arr : array<u32, x>;
+
+fn a(p: ptr<workgroup, array<u32, x>>) {
+}
+
+fn b() {
+  a(&arr);
+}
+)";
+    auto res = Build(src);
+    ASSERT_EQ(res, Success);
+
+    auto m = res.Move();
+    EXPECT_EQ(Dis(m), R"($B1: {  # root
+  %x:i32 = override 1i @id(0)
+  %arr:ptr<workgroup, array<u32, %x>, read_write> = var undef
+}
+
+%a = func(%p:ptr<workgroup, array<u32, %x>, read_write>):void {
+  $B2: {
+    ret
+  }
+}
+%b = func():void {
+  $B3: {
+    %6:void = call %a, %arr
+    ret
+  }
+}
+)");
+}
+
 }  // namespace
 }  // namespace tint::wgsl::reader