From f4ff6af0b4fa43993bc56ff4e26155eab327c2b0 Mon Sep 17 00:00:00 2001
From: James Price <jrprice@google.com>
Date: Wed, 10 Mar 2021 23:50:49 +0000
Subject: [PATCH] [wgsl-writer] Emit decorations on function parameters

This is needed to correctly generate entry point IO parameters.

Bug: tint:576
Change-Id: I9b96886d5ea90a54a568dd36506da563227afde7
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/44082
Commit-Queue: James Price <jrprice@google.com>
Auto-Submit: James Price <jrprice@google.com>
Reviewed-by: Ben Clayton <bclayton@google.com>
---
 src/writer/wgsl/generator_impl.cc             | 11 +++++++++
 .../wgsl/generator_impl_function_test.cc      | 24 +++++++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/src/writer/wgsl/generator_impl.cc b/src/writer/wgsl/generator_impl.cc
index dbe65f3c71..eefba24d02 100644
--- a/src/writer/wgsl/generator_impl.cc
+++ b/src/writer/wgsl/generator_impl.cc
@@ -307,6 +307,17 @@ bool GeneratorImpl::EmitFunction(ast::Function* func) {
     }
     first = false;
 
+    for (auto* deco : v->decorations()) {
+      out_ << "[[";
+      if (auto* builtin = deco->As<ast::BuiltinDecoration>()) {
+        out_ << "builtin(" << builtin->value() << ")";
+      }
+      if (auto* location = deco->As<ast::LocationDecoration>()) {
+        out_ << "location(" << location->value() << ")";
+      }
+      out_ << "]] ";
+    }
+
     out_ << program_->Symbols().NameFor(v->symbol()) << " : ";
 
     if (!EmitType(v->type())) {
diff --git a/src/writer/wgsl/generator_impl_function_test.cc b/src/writer/wgsl/generator_impl_function_test.cc
index 7fd4574f31..dfc99cf62a 100644
--- a/src/writer/wgsl/generator_impl_function_test.cc
+++ b/src/writer/wgsl/generator_impl_function_test.cc
@@ -143,6 +143,30 @@ TEST_F(WgslGeneratorImplTest, Emit_Function_WithDecoration_Multiple) {
 )");
 }
 
+TEST_F(WgslGeneratorImplTest, Emit_Function_EntryPoint_Parameters) {
+  auto* vec4 = ty.vec4<f32>();
+  auto* coord = Var("coord", vec4, ast::StorageClass::kInput, nullptr,
+                    {create<ast::BuiltinDecoration>(ast::Builtin::kFragCoord)});
+  auto* loc1 = Var("loc1", ty.f32(), ast::StorageClass::kInput, nullptr,
+                   {create<ast::LocationDecoration>(1u)});
+  auto* func =
+      Func("frag_main", ast::VariableList{coord, loc1}, ty.void_(),
+           ast::StatementList{},
+           ast::FunctionDecorationList{
+               create<ast::StageDecoration>(ast::PipelineStage::kFragment),
+           });
+
+  GeneratorImpl& gen = Build();
+
+  gen.increment_indent();
+
+  ASSERT_TRUE(gen.EmitFunction(func));
+  EXPECT_EQ(gen.result(), R"(  [[stage(fragment)]]
+  fn frag_main([[builtin(frag_coord)]] coord : vec4<f32>, [[location(1)]] loc1 : f32) -> void {
+  }
+)");
+}
+
 // https://crbug.com/tint/297
 TEST_F(WgslGeneratorImplTest,
        Emit_Function_Multiple_EntryPoint_With_Same_ModuleVar) {