ast: Remove Function::set_body()

Bug: tint:390
Change-Id: I525086868c6470ba39fe8c4ede7390e74c9489ff
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/35012
Commit-Queue: Ben Clayton <bclayton@google.com>
Reviewed-by: dan sinclair <dsinclair@chromium.org>

src/ast/function.h

@@ -153,9 +153,6 @@ class Function : public Castable<Function, Node> {
   // function is empty
   const Statement* get_last_statement() const;
 
-  /// Sets the body of the function
-  /// @param body the function body
-  void set_body(BlockStatement* body) { body_ = body; }
   /// @returns the function body
   const BlockStatement* body() const { return body_; }
   /// @returns the function body

src/ast/function_test.cc

@@ -185,12 +185,11 @@ TEST_F(FunctionTest, IsValid) {
   params.push_back(
       create<Variable>(Source{}, "var", StorageClass::kNone, &i32));
 
-  auto* block = create<BlockStatement>();
+  auto* body = create<BlockStatement>();
-  block->append(create<DiscardStatement>());
+  body->append(create<DiscardStatement>());
 
-  Function f(Source{}, "func", params, &void_type, create<BlockStatement>(),
+  Function f(Source{}, "func", params, &void_type, body,
              FunctionDecorationList{});
-  f.set_body(block);
   EXPECT_TRUE(f.IsValid());
 }
 
@@ -253,13 +252,13 @@ TEST_F(FunctionTest, IsValid_NullBodyStatement) {
   params.push_back(
       create<Variable>(Source{}, "var", StorageClass::kNone, &i32));
 
-  auto* block = create<BlockStatement>();
+  auto* body = create<BlockStatement>();
-  block->append(create<DiscardStatement>());
+  body->append(create<DiscardStatement>());
-  block->append(nullptr);
+  body->append(nullptr);
 
-  Function f(Source{}, "func", params, &void_type, create<BlockStatement>(),
+  Function f(Source{}, "func", params, &void_type, body,
              FunctionDecorationList{});
-  f.set_body(block);
+
   EXPECT_FALSE(f.IsValid());
 }
 
@@ -271,13 +270,12 @@ TEST_F(FunctionTest, IsValid_InvalidBodyStatement) {
   params.push_back(
       create<Variable>(Source{}, "var", StorageClass::kNone, &i32));
 
-  auto* block = create<BlockStatement>();
+  auto* body = create<BlockStatement>();
-  block->append(create<DiscardStatement>());
+  body->append(create<DiscardStatement>());
-  block->append(nullptr);
+  body->append(nullptr);
 
-  Function f(Source{}, "func", params, &void_type, create<BlockStatement>(),
+  Function f(Source{}, "func", params, &void_type, body,
              FunctionDecorationList{});
-  f.set_body(block);
   EXPECT_FALSE(f.IsValid());
 }
 
@@ -285,12 +283,10 @@ TEST_F(FunctionTest, ToStr) {
   type::Void void_type;
   type::I32 i32;
 
-  auto* block = create<BlockStatement>();
+  auto* body = create<BlockStatement>();
-  block->append(create<DiscardStatement>());
+  body->append(create<DiscardStatement>());
 
-  Function f(Source{}, "func", {}, &void_type, create<BlockStatement>(),
+  Function f(Source{}, "func", {}, &void_type, body, FunctionDecorationList{});
-             FunctionDecorationList{});
-  f.set_body(block);
 
   std::ostringstream out;
   f.to_str(out, 2);
@@ -306,11 +302,11 @@ TEST_F(FunctionTest, ToStr_WithDecoration) {
   type::Void void_type;
   type::I32 i32;
 
-  auto* block = create<BlockStatement>();
+  auto* body = create<BlockStatement>();
-  block->append(create<DiscardStatement>());
+  body->append(create<DiscardStatement>());
 
   Function f(
-      Source{}, "func", {}, &void_type, block,
+      Source{}, "func", {}, &void_type, body,
       FunctionDecorationList{create<WorkgroupDecoration>(2, 4, 6, Source{})});
 
   std::ostringstream out;
@@ -332,12 +328,11 @@ TEST_F(FunctionTest, ToStr_WithParams) {
   params.push_back(
       create<Variable>(Source{}, "var", StorageClass::kNone, &i32));
 
-  auto* block = create<BlockStatement>();
+  auto* body = create<BlockStatement>();
-  block->append(create<DiscardStatement>());
+  body->append(create<DiscardStatement>());
 
-  Function f(Source{}, "func", params, &void_type, create<BlockStatement>(),
+  Function f(Source{}, "func", params, &void_type, body,
              FunctionDecorationList{});
-  f.set_body(block);
 
   std::ostringstream out;
   f.to_str(out, 2);
@@ -386,9 +381,8 @@ TEST_F(FunctionTest, GetLastStatement) {
   auto* body = create<BlockStatement>();
   auto* stmt = create<DiscardStatement>();
   body->append(stmt);
-  Function f(Source{}, "func", params, &void_type, create<BlockStatement>(),
+  Function f(Source{}, "func", params, &void_type, body,
              FunctionDecorationList{});
-  f.set_body(body);
 
   EXPECT_EQ(f.get_last_statement(), stmt);
 }
@@ -398,9 +392,8 @@ TEST_F(FunctionTest, GetLastStatement_nullptr) {
 
   VariableList params;
   auto* body = create<BlockStatement>();
-  Function f(Source{}, "func", params, &void_type, create<BlockStatement>(),
+  Function f(Source{}, "func", params, &void_type, body,
              FunctionDecorationList{});
-  f.set_body(body);
 
   EXPECT_EQ(f.get_last_statement(), nullptr);
 }

src/reader/spirv/function.cc

@@ -702,7 +702,8 @@ bool FunctionEmitter::Emit() {
     return true;
   }
 
-  if (!EmitFunctionDeclaration()) {
+  FunctionDeclaration decl;
+  if (!ParseFunctionDeclaration(&decl)) {
     return false;
   }
 
@@ -716,8 +717,12 @@ bool FunctionEmitter::Emit() {
                      "element but has "
                   << statements_stack_.size();
   }
+
   auto* body = statements_stack_[0].statements_;
-  parser_impl_.get_module().functions().back()->set_body(body);
+  ast_module_.AddFunction(create<ast::Function>(
+      decl.source, decl.name, std::move(decl.params), decl.return_type, body,
+      std::move(decl.decorations)));
+
   // Maintain the invariant by repopulating the one and only element.
   statements_stack_.clear();
   PushNewStatementBlock(constructs_[0].get(), 0, nullptr, nullptr, nullptr);
@@ -725,7 +730,7 @@ bool FunctionEmitter::Emit() {
   return success();
 }
 
-bool FunctionEmitter::EmitFunctionDeclaration() {
+bool FunctionEmitter::ParseFunctionDeclaration(FunctionDeclaration* decl) {
   if (failed()) {
     return false;
   }
@@ -774,11 +779,11 @@ bool FunctionEmitter::EmitFunctionDeclaration() {
   if (ep_info_ != nullptr) {
     decos.emplace_back(create<ast::StageDecoration>(ep_info_->stage, Source{}));
   }
-  auto* ast_fn =
-      create<ast::Function>(Source{}, name, std::move(ast_params), ret_ty,
-                            create<ast::BlockStatement>(), std::move(decos));
 
-  ast_module_.AddFunction(ast_fn);
+  decl->name = name;
+  decl->params = std::move(ast_params);
+  decl->return_type = ret_ty;
+  decl->decorations = std::move(decos);
 
   return success();
 }
@@ -4114,6 +4119,9 @@ ast::Expression* FunctionEmitter::ConvertTexelForStorage(
   return ast_module_.create<ast::BitcastExpression>(dest_type, texel_prefix);
 }
 
+FunctionEmitter::FunctionDeclaration::FunctionDeclaration() = default;
+FunctionEmitter::FunctionDeclaration::~FunctionDeclaration() = default;
+
 }  // namespace spirv
 }  // namespace reader
 }  // namespace tint

src/reader/spirv/function.h

@@ -328,12 +328,6 @@ class FunctionEmitter {
   /// @returns the parser implementation
   ParserImpl* parser() { return &parser_impl_; }
 
-  /// Emits the declaration, which comprises the name, parameters, and
-  /// return type. The function AST node is appended to the module
-  /// AST node.
-  /// @returns true if emission has not yet failed.
-  bool EmitFunctionDeclaration();
-
   /// Emits the function body, populating the bottom entry of the statements
   /// stack.
   /// @returns false if emission failed.
@@ -700,6 +694,31 @@ class FunctionEmitter {
       const spvtools::opt::Instruction& image_access);
 
  private:
+  /// FunctionDeclaration contains the parsed information for a function header.
+  struct FunctionDeclaration {
+    /// Constructor
+    FunctionDeclaration();
+    /// Destructor
+    ~FunctionDeclaration();
+
+    /// Parsed header source
+    Source source;
+    /// Function name
+    std::string name;
+    /// Function parameters
+    ast::VariableList params;
+    /// Function return type
+    ast::type::Type* return_type;
+    /// Function decorations
+    ast::FunctionDecorationList decorations;
+  };
+
+  /// Parse the function declaration, which comprises the name, parameters, and
+  /// return type, populating `decl`.
+  /// @param decl the FunctionDeclaration to populate
+  /// @returns true if emission has not yet failed.
+  bool ParseFunctionDeclaration(FunctionDeclaration* decl);
+
   /// @returns the store type for the OpVariable instruction, or
   /// null on failure.
   ast::type::Type* GetVariableStoreType(

src/reader/spirv/function_decl_test.cc

@@ -49,7 +49,7 @@ std::string CommonTypes() {
   )";
 }
 
-TEST_F(SpvParserTest, EmitFunctionDeclaration_VoidFunctionWithoutParams) {
+TEST_F(SpvParserTest, Emit_VoidFunctionWithoutParams) {
   auto p = parser(test::Assemble(CommonTypes() + R"(
      %100 = OpFunction %void None %voidfn
      %entry = OpLabel
@@ -58,15 +58,20 @@ TEST_F(SpvParserTest, EmitFunctionDeclaration_VoidFunctionWithoutParams) {
   )"));
   ASSERT_TRUE(p->BuildAndParseInternalModuleExceptFunctions());
   FunctionEmitter fe(p.get(), *spirv_function(p.get(), 100));
-  EXPECT_TRUE(fe.EmitFunctionDeclaration());
+  EXPECT_TRUE(fe.Emit());
-  EXPECT_THAT(p->module().to_str(), HasSubstr(R"(
+  auto got = p->module().to_str();
+  auto* expect = R"(Module{
   Function x_100 -> __void
   ()
   {
-  })"));
+    Return{}
+  }
+}
+)";
+  EXPECT_EQ(got, expect);
 }
 
-TEST_F(SpvParserTest, EmitFunctionDeclaration_NonVoidResultType) {
+TEST_F(SpvParserTest, Emit_NonVoidResultType) {
   auto p = parser(test::Assemble(CommonTypes() + R"(
      %fn_ret_float = OpTypeFunction %float
      %100 = OpFunction %float None %fn_ret_float
@@ -76,16 +81,25 @@ TEST_F(SpvParserTest, EmitFunctionDeclaration_NonVoidResultType) {
   )"));
   ASSERT_TRUE(p->BuildAndParseInternalModuleExceptFunctions());
   FunctionEmitter fe(p.get(), *spirv_function(p.get(), 100));
-  EXPECT_TRUE(fe.EmitFunctionDeclaration());
+  EXPECT_TRUE(fe.Emit());
 
-  EXPECT_THAT(p->module().to_str(), HasSubstr(R"(
+  auto got = p->module().to_str();
+  auto* expect = R"(Module{
   Function x_100 -> __f32
   ()
   {
-  })"));
+    Return{
+      {
+        ScalarConstructor[not set]{0.000000}
+      }
+    }
+  }
+}
+)";
+  EXPECT_EQ(got, expect);
 }
 
-TEST_F(SpvParserTest, EmitFunctionDeclaration_MixedParamTypes) {
+TEST_F(SpvParserTest, Emit_MixedParamTypes) {
   auto p = parser(test::Assemble(Names({"a", "b", "c"}) + CommonTypes() + R"(
      %fn_mixed_params = OpTypeFunction %float %uint %float %int
 
@@ -99,9 +113,10 @@ TEST_F(SpvParserTest, EmitFunctionDeclaration_MixedParamTypes) {
   )"));
   ASSERT_TRUE(p->BuildAndParseInternalModuleExceptFunctions());
   FunctionEmitter fe(p.get(), *spirv_function(p.get(), 100));
-  EXPECT_TRUE(fe.EmitFunctionDeclaration());
+  EXPECT_TRUE(fe.Emit());
 
-  EXPECT_THAT(p->module().to_str(), HasSubstr(R"(
+  auto got = p->module().to_str();
+  auto* expect = R"(Module{
   Function x_100 -> __void
   (
     VariableConst{
@@ -121,10 +136,14 @@ TEST_F(SpvParserTest, EmitFunctionDeclaration_MixedParamTypes) {
     }
   )
   {
-  })"));
+    Return{}
+  }
+}
+)";
+  EXPECT_EQ(got, expect);
 }
 
-TEST_F(SpvParserTest, EmitFunctionDeclaration_GenerateParamNames) {
+TEST_F(SpvParserTest, Emit_GenerateParamNames) {
   auto p = parser(test::Assemble(CommonTypes() + R"(
      %fn_mixed_params = OpTypeFunction %float %uint %float %int
 
@@ -138,9 +157,10 @@ TEST_F(SpvParserTest, EmitFunctionDeclaration_GenerateParamNames) {
   )"));
   ASSERT_TRUE(p->BuildAndParseInternalModuleExceptFunctions());
   FunctionEmitter fe(p.get(), *spirv_function(p.get(), 100));
-  EXPECT_TRUE(fe.EmitFunctionDeclaration());
+  EXPECT_TRUE(fe.Emit());
 
-  EXPECT_THAT(p->module().to_str(), HasSubstr(R"(
+  auto got = p->module().to_str();
+  auto* expect = R"(Module{
   Function x_100 -> __void
   (
     VariableConst{
@@ -160,7 +180,11 @@ TEST_F(SpvParserTest, EmitFunctionDeclaration_GenerateParamNames) {
     }
   )
   {
-  })"));
+    Return{}
+  }
+}
+)";
+  EXPECT_EQ(got, expect);
 }
 
 }  // namespace

src/writer/spirv/builder_function_decoration_test.cc

@@ -143,12 +143,6 @@ TEST_F(BuilderTest, FunctionDecoration_Stage_WithUsedInterfaceIds) {
   ast::type::F32 f32;
   ast::type::Void void_type;
 
-  ast::Function func(
-      Source{}, "main", {}, &void_type, create<ast::BlockStatement>(),
-      ast::FunctionDecorationList{
-          create<ast::StageDecoration>(ast::PipelineStage::kVertex, Source{}),
-      });
-
   auto* body = create<ast::BlockStatement>();
   body->append(create<ast::AssignmentStatement>(
       create<ast::IdentifierExpression>("my_out"),
@@ -160,7 +154,12 @@ TEST_F(BuilderTest, FunctionDecoration_Stage_WithUsedInterfaceIds) {
   body->append(create<ast::AssignmentStatement>(
       create<ast::IdentifierExpression>("my_out"),
       create<ast::IdentifierExpression>("my_in")));
-  func.set_body(body);
+
+  ast::Function func(
+      Source{}, "main", {}, &void_type, body,
+      ast::FunctionDecorationList{
+          create<ast::StageDecoration>(ast::PipelineStage::kVertex, Source{}),
+      });
 
   auto* v_in =
       create<ast::Variable>(Source{}, "my_in", ast::StorageClass::kInput, &f32);