[spirv-writer] Handle non-struct entry point return values

Generate a global variable for the return value and replace return
statements with assignments to this variable.

Add a list of return statements to semantic::Function.

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

src/resolver/resolver.cc

@@ -321,6 +321,7 @@ bool Resolver::Statement(ast::Statement* stmt) {
     });
   }
   if (auto* r = stmt->As<ast::ReturnStatement>()) {
+    current_function_->return_statements.push_back(r);
     return Expression(r->value());
   }
   if (auto* s = stmt->As<ast::SwitchStatement>()) {
@@ -1215,7 +1216,7 @@ void Resolver::CreateSemanticNodes() const {
 
     auto* sem_func = builder_->create<semantic::Function>(
         info->declaration, remap_vars(info->referenced_module_vars),
-        remap_vars(info->local_referenced_module_vars),
+        remap_vars(info->local_referenced_module_vars), info->return_statements,
         ancestor_entry_points[func->symbol()]);
     func_info_to_sem_func.emplace(info, sem_func);
     sem.Add(func, sem_func);

src/resolver/resolver.h

@@ -38,6 +38,7 @@ class ConstructorExpression;
 class Function;
 class IdentifierExpression;
 class MemberAccessorExpression;
+class ReturnStatement;
 class UnaryOpExpression;
 class Variable;
 }  // namespace ast
@@ -92,6 +93,7 @@ class Resolver {
     ast::Function* const declaration;
     UniqueVector<VariableInfo*> referenced_module_vars;
     UniqueVector<VariableInfo*> local_referenced_module_vars;
+    std::vector<const ast::ReturnStatement*> return_statements;
 
     // List of transitive calls this function makes
     UniqueVector<FunctionInfo*> transitive_calls;

src/resolver/resolver_test.cc

@@ -858,6 +858,30 @@ TEST_F(ResolverTest, Function_NotRegisterFunctionVariable) {
   EXPECT_EQ(func_sem->ReferencedModuleVariables().size(), 0u);
 }
 
+TEST_F(ResolverTest, Function_ReturnStatements) {
+  auto* var = Var("foo", ty.f32(), ast::StorageClass::kFunction);
+
+  auto* ret_1 = create<ast::ReturnStatement>(Expr(1.f));
+  auto* ret_foo = create<ast::ReturnStatement>(Expr("foo"));
+
+  auto* func = Func("my_func", ast::VariableList{}, ty.f32(),
+                    ast::StatementList{
+                        create<ast::VariableDeclStatement>(var),
+                        If(Expr(true), Block(ret_1)),
+                        ret_foo,
+                    },
+                    ast::DecorationList{});
+
+  EXPECT_TRUE(r()->Resolve()) << r()->error();
+
+  auto* func_sem = Sem().Get(func);
+  ASSERT_NE(func_sem, nullptr);
+
+  EXPECT_EQ(func_sem->ReturnStatements().size(), 2u);
+  EXPECT_EQ(func_sem->ReturnStatements()[0], ret_1);
+  EXPECT_EQ(func_sem->ReturnStatements()[1], ret_foo);
+}
+
 TEST_F(ResolverTest, Expr_MemberAccessor_Struct) {
   auto* strct = create<ast::Struct>(
       ast::StructMemberList{Member("first_member", ty.i32()),