From cc193de025d6bd3e927115c928a5dc91e8838312 Mon Sep 17 00:00:00 2001
From: James Price <jrprice@google.com>
Date: Mon, 15 Mar 2021 17:15:23 +0000
Subject: [PATCH] [wgsl-writer] Generate decorations on function return types

Bug: tint:576
Change-Id: I60d046fb4db558daf52178600b962eaeb6d3f843
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/44602
Commit-Queue: James Price <jrprice@google.com>
Reviewed-by: Ben Clayton <bclayton@google.com>
---
 src/writer/wgsl/generator_impl.cc             |  7 ++++++
 .../wgsl/generator_impl_function_test.cc      | 25 +++++++++++++++++++
 2 files changed, 32 insertions(+)

diff --git a/src/writer/wgsl/generator_impl.cc b/src/writer/wgsl/generator_impl.cc
index 304f720b81..273f934fd2 100644
--- a/src/writer/wgsl/generator_impl.cc
+++ b/src/writer/wgsl/generator_impl.cc
@@ -329,6 +329,13 @@ bool GeneratorImpl::EmitFunction(ast::Function* func) {
 
   out_ << ") -> ";
 
+  if (!func->return_type_decorations().empty()) {
+    if (!EmitDecorations(func->return_type_decorations())) {
+      return false;
+    }
+    out_ << " ";
+  }
+
   if (!EmitType(func->return_type())) {
     return false;
   }
diff --git a/src/writer/wgsl/generator_impl_function_test.cc b/src/writer/wgsl/generator_impl_function_test.cc
index f79d0d6d99..9da8478870 100644
--- a/src/writer/wgsl/generator_impl_function_test.cc
+++ b/src/writer/wgsl/generator_impl_function_test.cc
@@ -167,6 +167,31 @@ TEST_F(WgslGeneratorImplTest, Emit_Function_EntryPoint_Parameters) {
 )");
 }
 
+TEST_F(WgslGeneratorImplTest, Emit_Function_EntryPoint_ReturnValue) {
+  auto* func =
+      Func("frag_main", ast::VariableList{}, ty.f32(),
+           ast::StatementList{
+               create<ast::ReturnStatement>(Expr(1.f)),
+           },
+           ast::DecorationList{
+               create<ast::StageDecoration>(ast::PipelineStage::kFragment),
+           },
+           ast::DecorationList{
+               create<ast::LocationDecoration>(1u),
+           });
+
+  GeneratorImpl& gen = Build();
+
+  gen.increment_indent();
+
+  ASSERT_TRUE(gen.EmitFunction(func));
+  EXPECT_EQ(gen.result(), R"(  [[stage(fragment)]]
+  fn frag_main() -> [[location(1)]] f32 {
+    return 1.0;
+  }
+)");
+}
+
 // https://crbug.com/tint/297
 TEST_F(WgslGeneratorImplTest,
        Emit_Function_Multiple_EntryPoint_With_Same_ModuleVar) {