[msl-writer] Handle non-struct entry point parameters

Add a sanitizing transform to collect location-decorated parameters
into a struct.

Bug: tint:510
Change-Id: I1e9bf829dac946e5fec0ecfe6da7e1ef9cebff8e
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/44322
Commit-Queue: James Price <jrprice@google.com>
Auto-Submit: James Price <jrprice@google.com>
Reviewed-by: Ben Clayton <bclayton@google.com>

src/transform/msl.cc

@@ -15,8 +15,10 @@
 #include "src/transform/msl.h"
 
 #include <utility>
+#include <vector>
 
 #include "src/program_builder.h"
+#include "src/semantic/variable.h"
 
 namespace tint {
 namespace transform {
@@ -266,10 +268,125 @@ Transform::Output Msl::Run(const Program* in) {
   ProgramBuilder out;
   CloneContext ctx(&out, in);
   RenameReservedKeywords(&ctx, kReservedKeywords);
+  HandleEntryPointIOTypes(ctx);
   ctx.Clone();
 
   return Output{Program(std::move(out))};
 }
 
+void Msl::HandleEntryPointIOTypes(CloneContext& ctx) const {
+  // Collect location-decorated entry point parameters into a struct.
+  // Insert function-scope const declarations to replace those parameters.
+  //
+  // Before:
+  // ```
+  // [[stage(fragment)]]
+  // fn frag_main([[builtin(frag_coord)]] coord : vec4<f32>,
+  //              [[location(1)]] loc1 : f32,
+  //              [[location(2)]] loc2 : vec4<u32>) -> void {
+  //   var col : f32 = (coord.x * loc1);
+  // }
+  // ```
+  //
+  // After:
+  // ```
+  // struct frag_main_in {
+  //   [[location(1)]] loc1 : f32;
+  //   [[location(2)]] loc2 : vec4<u32>
+  // };
+
+  // [[stage(fragment)]]
+  // fn frag_main([[builtin(frag_coord)]] coord : vec4<f32>,
+  //              in : frag_main_in) -> void {
+  //   const loc1 : f32 = in.loc1;
+  //   const loc2 : vec4<u32> = in.loc2;
+  //   var col : f32 = (coord.x * loc1);
+  // }
+  // ```
+
+  for (auto* func : ctx.src->AST().Functions()) {
+    if (!func->IsEntryPoint()) {
+      continue;
+    }
+
+    std::vector<ast::Variable*> worklist;
+    ast::StructMemberList struct_members;
+    ast::VariableList new_parameters;
+    ast::StatementList new_body;
+
+    // Find location-decorated parameters.
+    for (auto* param : func->params()) {
+      // TODO(jrprice): Handle structs (collate members into a single struct).
+      if (param->decorations().size() != 1) {
+        TINT_ICE(ctx.dst->Diagnostics()) << "Unsupported entry point parameter";
+      }
+
+      auto* deco = param->decorations()[0];
+      if (auto* builtin = deco->As<ast::BuiltinDecoration>()) {
+        // Keep any builtin-decorated parameters unchanged.
+        new_parameters.push_back(ctx.Clone(param));
+      } else if (auto* loc = deco->As<ast::LocationDecoration>()) {
+        // Create a struct member with the location decoration.
+        struct_members.push_back(
+            ctx.dst->Member(param->symbol().to_str(), ctx.Clone(param->type()),
+                            ast::DecorationList{ctx.Clone(loc)}));
+        worklist.push_back(param);
+      } else {
+        TINT_ICE(ctx.dst->Diagnostics())
+            << "Unsupported entry point parameter decoration";
+      }
+    }
+
+    if (worklist.empty()) {
+      // Nothing to do.
+      continue;
+    }
+
+    // Create a struct type to hold all of the user-defined input parameters.
+    auto* in_struct = ctx.dst->create<type::Struct>(
+        ctx.dst->Symbols().New(),
+        ctx.dst->create<ast::Struct>(struct_members, ast::DecorationList{}));
+    ctx.dst->AST().AddConstructedType(in_struct);
+
+    // Create a new function parameter using this struct type.
+    auto struct_param_symbol = ctx.dst->Symbols().New();
+    auto* struct_param =
+        ctx.dst->Var(struct_param_symbol, in_struct, ast::StorageClass::kNone);
+    new_parameters.push_back(struct_param);
+
+    // Replace the original parameters with function-scope constants.
+    for (auto* param : worklist) {
+      // Create a function-scope const to replace the parameter.
+      // Initialize it with the value extracted from the struct parameter.
+      auto func_const_symbol = ctx.dst->Symbols().New();
+      auto* func_const =
+          ctx.dst->Const(func_const_symbol, ctx.Clone(param->type()),
+                         ctx.dst->MemberAccessor(struct_param_symbol,
+                                                 param->symbol().to_str()));
+
+      new_body.push_back(ctx.dst->WrapInStatement(func_const));
+
+      // Replace all uses of the function parameter with the function const.
+      for (auto* user : ctx.src->Sem().Get(param)->Users()) {
+        ctx.Replace<ast::Expression>(user->Declaration(),
+                                     ctx.dst->Expr(func_const_symbol));
+      }
+    }
+
+    // Copy over the rest of the function body unchanged.
+    for (auto* stmt : func->body()->list()) {
+      new_body.push_back(ctx.Clone(stmt));
+    }
+
+    // Rewrite the function header with the new parameters.
+    auto* new_func = ctx.dst->create<ast::Function>(
+        func->source(), ctx.Clone(func->symbol()), new_parameters,
+        ctx.Clone(func->return_type()),
+        ctx.dst->create<ast::BlockStatement>(new_body),
+        ctx.Clone(func->decorations()));
+    ctx.Replace(func, new_func);
+  }
+}
+
 }  // namespace transform
 }  // namespace tint

src/transform/msl.h

@@ -33,6 +33,9 @@ class Msl : public Transform {
   /// @param program the source program to transform
   /// @returns the transformation result
   Output Run(const Program* program) override;
+
+  /// Hoist location-decorated entry point parameters out to struct members.
+  void HandleEntryPointIOTypes(CloneContext& ctx) const;
 };
 
 }  // namespace transform

src/transform/msl_test.cc

@@ -326,6 +326,68 @@ INSTANTIATE_TEST_SUITE_P(MslReservedKeywordTest,
                              "vec",
                              "vertex"));
 
+using MslEntryPointIOTest = TransformTest;
+
+TEST_F(MslEntryPointIOTest, HandleEntryPointIOTypes_Parameters) {
+  auto* src = R"(
+[[stage(fragment)]]
+fn frag_main([[builtin(frag_coord)]] coord : vec4<f32>,
+             [[location(1)]] loc1 : f32,
+             [[location(2)]] loc2 : vec4<u32>) -> void {
+  var col : f32 = (coord.x * loc1);
+}
+)";
+
+  auto* expect = R"(
+struct tint_symbol_4 {
+  [[location(1)]]
+  tint_symbol_2 : f32;
+  [[location(2)]]
+  tint_symbol_3 : vec4<u32>;
+};
+
+[[stage(fragment)]]
+fn frag_main([[builtin(frag_coord)]] coord : vec4<f32>, tint_symbol_5 : tint_symbol_4) -> void {
+  const tint_symbol_6 : f32 = tint_symbol_5.tint_symbol_2;
+  const tint_symbol_7 : vec4<u32> = tint_symbol_5.tint_symbol_3;
+  var col : f32 = (coord.x * tint_symbol_6);
+}
+)";
+
+  auto got = Transform<Msl>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(MslEntryPointIOTest, HandleEntryPointIOTypes_Parameters_EmptyBody) {
+  auto* src = R"(
+[[stage(fragment)]]
+fn frag_main([[builtin(frag_coord)]] coord : vec4<f32>,
+             [[location(1)]] loc1 : f32,
+             [[location(2)]] loc2 : vec4<u32>) -> void {
+}
+)";
+
+  auto* expect = R"(
+struct tint_symbol_4 {
+  [[location(1)]]
+  tint_symbol_2 : f32;
+  [[location(2)]]
+  tint_symbol_3 : vec4<u32>;
+};
+
+[[stage(fragment)]]
+fn frag_main([[builtin(frag_coord)]] coord : vec4<f32>, tint_symbol_5 : tint_symbol_4) -> void {
+  const tint_symbol_6 : f32 = tint_symbol_5.tint_symbol_2;
+  const tint_symbol_7 : vec4<u32> = tint_symbol_5.tint_symbol_3;
+}
+)";
+
+  auto got = Transform<Msl>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
 }  // namespace
 }  // namespace transform
 }  // namespace tint