[wgsl-writer] Emit stage decorations.

This CL updates the wgsl writer to emit the stage decorations.

Change-Id: I0501088e074aa5d027f83fd0a2ac45ba3f36332f
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/28701
Commit-Queue: dan sinclair <dsinclair@chromium.org>
Reviewed-by: David Neto <dneto@google.com>
Reviewed-by: Sarah Mashayekhi <sarahmashay@google.com>

src/writer/wgsl/generator_impl.cc

diff --git a/src/writer/wgsl/generator_impl.cc b/src/writer/wgsl/generator_impl.cc
index 08afc35..8b9e63c 100644
--- a/src/writer/wgsl/generator_impl.cc
+++ b/src/writer/wgsl/generator_impl.cc
@@ -44,6 +44,7 @@
 #include "src/ast/scalar_constructor_expression.h"
 #include "src/ast/set_decoration.h"
 #include "src/ast/sint_literal.h"
+#include "src/ast/stage_decoration.h"
 #include "src/ast/statement.h"
 #include "src/ast/struct.h"
 #include "src/ast/struct_member.h"
@@ -434,6 +435,9 @@
       out_ << "workgroup_size(" << std::to_string(x) << ", "
            << std::to_string(y) << ", " << std::to_string(z) << ")";
     }
+    if (deco->IsStage()) {
+      out_ << "stage(" << deco->AsStage()->value() << ")";
+    }
     out_ << "]]" << std::endl;
   }
 

src/writer/wgsl/generator_impl_function_test.cc

diff --git a/src/writer/wgsl/generator_impl_function_test.cc b/src/writer/wgsl/generator_impl_function_test.cc
index 25bac24..95ad844 100644
--- a/src/writer/wgsl/generator_impl_function_test.cc
+++ b/src/writer/wgsl/generator_impl_function_test.cc
@@ -15,7 +15,9 @@
 #include "gtest/gtest.h"
 #include "src/ast/discard_statement.h"
 #include "src/ast/function.h"
+#include "src/ast/pipeline_stage.h"
 #include "src/ast/return_statement.h"
+#include "src/ast/stage_decoration.h"
 #include "src/ast/type/f32_type.h"
 #include "src/ast/type/i32_type.h"
 #include "src/ast/type/void_type.h"
@@ -78,7 +80,7 @@
 )");
 }
 
-TEST_F(WgslGeneratorImplTest, Emit_Function_WithDecorations) {
+TEST_F(WgslGeneratorImplTest, Emit_Function_WithDecoration_WorkgroupSize) {
   auto body = std::make_unique<ast::BlockStatement>();
   body->append(std::make_unique<ast::DiscardStatement>());
   body->append(std::make_unique<ast::ReturnStatement>());
@@ -100,6 +102,54 @@
 )");
 }
 
+TEST_F(WgslGeneratorImplTest, Emit_Function_WithDecoration_Stage) {
+  auto body = std::make_unique<ast::BlockStatement>();
+  body->append(std::make_unique<ast::DiscardStatement>());
+  body->append(std::make_unique<ast::ReturnStatement>());
+
+  ast::type::VoidType void_type;
+  ast::Function func("my_func", {}, &void_type);
+  func.add_decoration(
+      std::make_unique<ast::StageDecoration>(ast::PipelineStage::kFragment));
+  func.set_body(std::move(body));
+
+  GeneratorImpl g;
+  g.increment_indent();
+
+  ASSERT_TRUE(g.EmitFunction(&func));
+  EXPECT_EQ(g.result(), R"(  [[stage(fragment)]]
+  fn my_func() -> void {
+    discard;
+    return;
+  }
+)");
+}
+
+TEST_F(WgslGeneratorImplTest, Emit_Function_WithDecoration_Multiple) {
+  auto body = std::make_unique<ast::BlockStatement>();
+  body->append(std::make_unique<ast::DiscardStatement>());
+  body->append(std::make_unique<ast::ReturnStatement>());
+
+  ast::type::VoidType void_type;
+  ast::Function func("my_func", {}, &void_type);
+  func.add_decoration(
+      std::make_unique<ast::StageDecoration>(ast::PipelineStage::kFragment));
+  func.add_decoration(std::make_unique<ast::WorkgroupDecoration>(2u, 4u, 6u));
+  func.set_body(std::move(body));
+
+  GeneratorImpl g;
+  g.increment_indent();
+
+  ASSERT_TRUE(g.EmitFunction(&func));
+  EXPECT_EQ(g.result(), R"(  [[stage(fragment)]]
+  [[workgroup_size(2, 4, 6)]]
+  fn my_func() -> void {
+    discard;
+    return;
+  }
+)");
+}
+
 }  // namespace
 }  // namespace wgsl
 }  // namespace writer