[wgsl-writer] Generate builtin and location decorations on struct members

Bug: tint:576
Change-Id: Ie8ace8dd77095abedcca97caca330e2d11a7559c
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/44321
Commit-Queue: James Price <jrprice@google.com>
Reviewed-by: Ben Clayton <bclayton@chromium.org>
diff --git a/src/writer/wgsl/generator_impl.cc b/src/writer/wgsl/generator_impl.cc
index eefba24..7ea7a8f 100644
--- a/src/writer/wgsl/generator_impl.cc
+++ b/src/writer/wgsl/generator_impl.cc
@@ -500,14 +500,14 @@
 
   increment_indent();
   for (auto* mem : impl->members()) {
-    for (auto* deco : mem->decorations()) {
+    if (!mem->decorations().empty()) {
       make_indent();
-
-      // TODO(dsinclair): Split this out when we have more then one
-      auto* offset = deco->As<ast::StructMemberOffsetDecoration>();
-      assert(offset != nullptr);
-      out_ << "[[offset(" << offset->offset() << ")]]" << std::endl;
+      if (!EmitDecorations(mem->decorations())) {
+        return false;
+      }
+      out_ << std::endl;
     }
+
     make_indent();
     out_ << program_->Symbols().NameFor(mem->symbol()) << " : ";
     if (!EmitType(mem->type())) {
@@ -527,8 +527,11 @@
 
   make_indent();
 
-  if (!var->decorations().empty() && !EmitVariableDecorations(sem)) {
-    return false;
+  if (!var->decorations().empty()) {
+    if (!EmitDecorations(var->decorations())) {
+      return false;
+    }
+    out_ << " ";
   }
 
   if (var->is_const()) {
@@ -558,12 +561,10 @@
   return true;
 }
 
-bool GeneratorImpl::EmitVariableDecorations(const semantic::Variable* var) {
-  auto* decl = var->Declaration();
-
+bool GeneratorImpl::EmitDecorations(const ast::DecorationList& decos) {
   out_ << "[[";
   bool first = true;
-  for (auto* deco : decl->decorations()) {
+  for (auto* deco : decos) {
     if (!first) {
       out_ << ", ";
     }
@@ -579,12 +580,14 @@
       out_ << "builtin(" << builtin->value() << ")";
     } else if (auto* constant = deco->As<ast::ConstantIdDecoration>()) {
       out_ << "constant_id(" << constant->value() << ")";
+    } else if (auto* offset = deco->As<ast::StructMemberOffsetDecoration>()) {
+      out_ << "offset(" << offset->offset() << ")";
     } else {
       diagnostics_.add_error("unknown variable decoration");
       return false;
     }
   }
-  out_ << "]] ";
+  out_ << "]]";
 
   return true;
 }
diff --git a/src/writer/wgsl/generator_impl.h b/src/writer/wgsl/generator_impl.h
index a584f0f..b5b50a6 100644
--- a/src/writer/wgsl/generator_impl.h
+++ b/src/writer/wgsl/generator_impl.h
@@ -193,10 +193,10 @@
   /// @param var the variable to generate
   /// @returns true if the variable was emitted
   bool EmitVariable(ast::Variable* var);
-  /// Handles generating variable decorations
-  /// @param var the decorated variable
-  /// @returns true if the variable decoration was emitted
-  bool EmitVariableDecorations(const semantic::Variable* var);
+  /// Handles generating a decoration list
+  /// @param decos the decoration list
+  /// @returns true if the decorations were emitted
+  bool EmitDecorations(const ast::DecorationList& decos);
 
  private:
   Program const* const program_;
diff --git a/src/writer/wgsl/generator_impl_test.cc b/src/writer/wgsl/generator_impl_test.cc
index 515c8b7..77c5a1a 100644
--- a/src/writer/wgsl/generator_impl_test.cc
+++ b/src/writer/wgsl/generator_impl_test.cc
@@ -53,10 +53,10 @@
 
   GeneratorImpl& gen = Build();
 
-  gen.EmitVariableDecorations(program->Sem().Get(var));
+  gen.EmitDecorations(var->decorations());
 
   EXPECT_EQ(gen.result(),
-            "[[builtin(" + std::string(params.attribute_name) + ")]] ");
+            "[[builtin(" + std::string(params.attribute_name) + ")]]");
 }
 INSTANTIATE_TEST_SUITE_P(
     WgslGeneratorImplTest,
diff --git a/src/writer/wgsl/generator_impl_type_test.cc b/src/writer/wgsl/generator_impl_type_test.cc
index 0bd2acf..b3e5553 100644
--- a/src/writer/wgsl/generator_impl_type_test.cc
+++ b/src/writer/wgsl/generator_impl_type_test.cc
@@ -210,6 +210,31 @@
 )");
 }
 
+TEST_F(WgslGeneratorImplTest, EmitType_Struct_WithEntryPointDecorations) {
+  ast::DecorationList decos;
+  decos.push_back(create<ast::StructBlockDecoration>());
+
+  auto* str = create<ast::Struct>(
+      ast::StructMemberList{
+          Member("a", ty.u32(),
+                 {create<ast::BuiltinDecoration>(ast::Builtin::kVertexIndex)}),
+          Member("b", ty.f32(), {create<ast::LocationDecoration>(2u)})},
+      decos);
+
+  auto* s = ty.struct_("S", str);
+  GeneratorImpl& gen = Build();
+
+  ASSERT_TRUE(gen.EmitStructType(s)) << gen.error();
+  EXPECT_EQ(gen.result(), R"([[block]]
+struct S {
+  [[builtin(vertex_index)]]
+  a : u32;
+  [[location(2)]]
+  b : f32;
+};
+)");
+}
+
 TEST_F(WgslGeneratorImplTest, EmitType_U32) {
   auto* u32 = ty.u32();