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

Add a sanitizing transform to collect input parameters into a
struct. HLSL does not allow non-struct entry-point parameters, so any
location- or builtin-decorated inputs have to be provided via a struct
instead.

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

src/transform/hlsl.cc

@@ -15,11 +15,13 @@
 #include "src/transform/hlsl.h"
 
 #include <utility>
+#include <vector>
 
 #include "src/ast/variable_decl_statement.h"
 #include "src/program_builder.h"
 #include "src/semantic/expression.h"
 #include "src/semantic/statement.h"
+#include "src/semantic/variable.h"
 
 namespace tint {
 namespace transform {
@@ -31,6 +33,7 @@ Transform::Output Hlsl::Run(const Program* in) {
   ProgramBuilder out;
   CloneContext ctx(&out, in);
   PromoteArrayInitializerToConstVar(ctx);
+  HandleEntryPointIOTypes(ctx);
   ctx.Clone();
   return Output{Program(std::move(out))};
 }
@@ -103,5 +106,133 @@ void Hlsl::PromoteArrayInitializerToConstVar(CloneContext& ctx) const {
   }
 }
 
+void Hlsl::HandleEntryPointIOTypes(CloneContext& ctx) const {
+  // Collect 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 {
+  //   [[builtin(frag_coord)]] coord : vec4<f32>;
+  //   [[location(1)]] loc1 : f32;
+  //   [[location(2)]] loc2 : vec4<u32>
+  // };
+
+  // [[stage(fragment)]]
+  // fn frag_main(in : frag_main_in) -> void {
+  //   const coord : vec4<f32> = in.coord;
+  //   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;
+    }
+
+    // Build a new structure to hold the non-struct input parameters.
+    ast::StructMemberList struct_members;
+    for (auto* param : func->params()) {
+      if (param->type()->Is<type::Struct>()) {
+        // Already a struct, nothing to do.
+        continue;
+      }
+
+      if (param->decorations().size() != 1) {
+        TINT_ICE(ctx.dst->Diagnostics()) << "Unsupported entry point parameter";
+      }
+
+      auto name = ctx.src->Symbols().NameFor(param->symbol());
+
+      auto* deco = param->decorations()[0];
+      if (auto* builtin = deco->As<ast::BuiltinDecoration>()) {
+        // Create a struct member with the builtin decoration.
+        struct_members.push_back(
+            ctx.dst->Member(name, ctx.Clone(param->type()),
+                            ast::DecorationList{ctx.Clone(builtin)}));
+      } else if (auto* loc = deco->As<ast::LocationDecoration>()) {
+        // Create a struct member with the location decoration.
+        struct_members.push_back(
+            ctx.dst->Member(name, ctx.Clone(param->type()),
+                            ast::DecorationList{ctx.Clone(loc)}));
+      } else {
+        TINT_ICE(ctx.dst->Diagnostics())
+            << "Unsupported entry point parameter decoration";
+      }
+    }
+
+    if (struct_members.empty()) {
+      // Nothing to do.
+      continue;
+    }
+
+    ast::VariableList new_parameters;
+    ast::StatementList new_body;
+
+    // Create a struct type to hold all of the non-struct 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 : func->params()) {
+      if (param->type()->Is<type::Struct>()) {
+        // Keep struct parameters unchanged.
+        new_parameters.push_back(ctx.Clone(param));
+        continue;
+      }
+
+      auto name = ctx.src->Symbols().NameFor(param->symbol());
+
+      // 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().Register(name);
+      auto* func_const =
+          ctx.dst->Const(func_const_symbol, ctx.Clone(param->type()),
+                         ctx.dst->MemberAccessor(struct_param_symbol, name));
+
+      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/hlsl.h

@@ -43,6 +43,9 @@ class Hlsl : public Transform {
   /// the array usage statement.
   /// See crbug.com/tint/406 for more details
   void PromoteArrayInitializerToConstVar(CloneContext& ctx) const;
+
+  /// Hoist entry point parameters out to struct members.
+  void HandleEntryPointIOTypes(CloneContext& ctx) const;
 };
 
 }  // namespace transform

src/transform/hlsl_test.cc

@@ -143,6 +143,105 @@ fn main() -> void {
   EXPECT_EQ(expect, str(got));
 }
 
+TEST_F(HlslTest, HandleEntryPointIOTypes_Parameters) {
+  auto* src = R"(
+struct FragIn {
+  [[location(2)]]
+  loc2 : f32;
+};
+
+[[stage(fragment)]]
+fn frag_main([[builtin(frag_coord)]] coord : vec4<f32>,
+             [[location(1)]] loc1 : f32,
+             frag_in : FragIn) -> void {
+  var col : f32 = (coord.x * loc1 + frag_in.loc2);
+}
+)";
+
+  auto* expect = R"(
+struct tint_symbol_3 {
+  [[builtin(frag_coord)]]
+  coord : vec4<f32>;
+  [[location(1)]]
+  loc1 : f32;
+};
+
+struct FragIn {
+  [[location(2)]]
+  loc2 : f32;
+};
+
+[[stage(fragment)]]
+fn frag_main(tint_symbol_4 : tint_symbol_3, frag_in : FragIn) -> void {
+  const coord : vec4<f32> = tint_symbol_4.coord;
+  const loc1 : f32 = tint_symbol_4.loc1;
+  var col : f32 = ((coord.x * loc1) + frag_in.loc2);
+}
+)";
+
+  auto got = Transform<Hlsl>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
+TEST_F(HlslTest, HandleEntryPointIOTypes_OnlyStructParameters) {
+  // Expect no change.
+  auto* src = R"(
+struct FragBuiltins {
+  [[builtin(frag_coord)]]
+  coord : vec4<f32>;
+};
+
+struct FragInputs {
+  [[location(1)]]
+  loc1 : f32;
+  [[location(2)]]
+  loc2 : vec4<u32>;
+};
+
+[[stage(fragment)]]
+fn frag_main(builtins : FragBuiltins, inputs : FragInputs) -> void {
+  var col : f32 = (builtins.coord.x * inputs.loc1);
+}
+)";
+
+  auto got = Transform<Hlsl>(src);
+
+  EXPECT_EQ(src, str(got));
+}
+
+TEST_F(HlslTest, 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 {
+  [[builtin(frag_coord)]]
+  coord : vec4<f32>;
+  [[location(1)]]
+  loc1 : f32;
+  [[location(2)]]
+  loc2 : vec4<u32>;
+};
+
+[[stage(fragment)]]
+fn frag_main(tint_symbol_5 : tint_symbol_4) -> void {
+  const coord : vec4<f32> = tint_symbol_5.coord;
+  const loc1 : f32 = tint_symbol_5.loc1;
+  const loc2 : vec4<u32> = tint_symbol_5.loc2;
+}
+)";
+
+  auto got = Transform<Hlsl>(src);
+
+  EXPECT_EQ(expect, str(got));
+}
+
 }  // namespace
 }  // namespace transform
 }  // namespace tint