From b29b09fba6e2a3c8455130b09e54fe154aed8e32 Mon Sep 17 00:00:00 2001
From: Ben Clayton <bclayton@google.com>
Date: Sat, 14 Nov 2020 01:13:04 +0000
Subject: [PATCH] ast: Add body parameter to ast::Function constructors

In a near-future change, AST nodes, such as ast::BlockStatement will no longer
be std::unique_ptrs, and will have to be constructed and owned by an external
class. This means AST nodes can no longer allocate default child nodes.

Bug: tint:322
Change-Id: Iddb5605b9bc0de80ad2710ced0e429f89410af2f
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/32675
Commit-Queue: dan sinclair <dsinclair@chromium.org>
Reviewed-by: dan sinclair <dsinclair@chromium.org>
---
 src/ast/function.cc                           |  10 +-
 src/ast/function.h                            |   8 +-
 src/ast/function_test.cc                      |  54 +++--
 src/ast/module_test.cc                        |   6 +-
 src/inspector/inspector_test.cc               |  55 ++---
 src/reader/spirv/function.cc                  |   3 +-
 src/reader/wgsl/parser_impl.cc                |   5 +-
 .../bound_array_accessors_transform_test.cc   |   4 +-
 .../vertex_pulling_transform_test.cc          |   6 +-
 src/type_determiner_test.cc                   | 116 +++++------
 src/validator/validator_function_test.cc      |  73 +++----
 src/validator/validator_test.cc               |  26 +--
 src/writer/hlsl/generator_impl_binary_test.cc |   3 +-
 src/writer/hlsl/generator_impl_call_test.cc   |   9 +-
 ...tor_impl_function_entry_point_data_test.cc |  82 ++++----
 .../hlsl/generator_impl_function_test.cc      | 196 ++++++++----------
 src/writer/hlsl/generator_impl_test.cc        |   3 +-
 src/writer/msl/generator_impl_call_test.cc    |   9 +-
 ...tor_impl_function_entry_point_data_test.cc |  82 ++++----
 .../msl/generator_impl_function_test.cc       | 182 ++++++++--------
 src/writer/msl/generator_impl_test.cc         |   3 +-
 src/writer/spirv/builder_call_test.cc         |  14 +-
 .../spirv/builder_function_decoration_test.cc |  18 +-
 src/writer/spirv/builder_function_test.cc     |  35 ++--
 src/writer/spirv/builder_intrinsic_test.cc    |  58 +++---
 src/writer/spirv/builder_switch_test.cc       |  12 +-
 .../wgsl/generator_impl_function_test.cc      |  35 ++--
 src/writer/wgsl/generator_impl_test.cc        |   5 +-
 28 files changed, 545 insertions(+), 567 deletions(-)

diff --git a/src/ast/function.cc b/src/ast/function.cc
index ab1b735f4e..7ff3774467 100644
--- a/src/ast/function.cc
+++ b/src/ast/function.cc
@@ -28,22 +28,24 @@ Function::Function() = default;
 
 Function::Function(const std::string& name,
                    VariableList params,
-                   type::Type* return_type)
+                   type::Type* return_type,
+                   std::unique_ptr<BlockStatement> body)
     : Node(),
       name_(name),
       params_(std::move(params)),
       return_type_(return_type),
-      body_(std::make_unique<BlockStatement>()) {}
+      body_(std::move(body)) {}
 
 Function::Function(const Source& source,
                    const std::string& name,
                    VariableList params,
-                   type::Type* return_type)
+                   type::Type* return_type,
+                   std::unique_ptr<BlockStatement> body)
     : Node(source),
       name_(name),
       params_(std::move(params)),
       return_type_(return_type),
-      body_(std::make_unique<BlockStatement>()) {}
+      body_(std::move(body)) {}
 
 Function::Function(Function&&) = default;
 
diff --git a/src/ast/function.h b/src/ast/function.h
index 2dbeab4a61..cbd4d20705 100644
--- a/src/ast/function.h
+++ b/src/ast/function.h
@@ -56,18 +56,22 @@ class Function : public Node {
   /// @param name the function name
   /// @param params the function parameters
   /// @param return_type the return type
+  /// @param body the function body
   Function(const std::string& name,
            VariableList params,
-           type::Type* return_type);
+           type::Type* return_type,
+           std::unique_ptr<BlockStatement> body);
   /// Create a function
   /// @param source the variable source
   /// @param name the function name
   /// @param params the function parameters
   /// @param return_type the return type
+  /// @param body the function body
   Function(const Source& source,
            const std::string& name,
            VariableList params,
-           type::Type* return_type);
+           type::Type* return_type,
+           std::unique_ptr<BlockStatement> body);
   /// Move constructor
   Function(Function&&);
 
diff --git a/src/ast/function_test.cc b/src/ast/function_test.cc
index 91785bc860..0eb838de2d 100644
--- a/src/ast/function_test.cc
+++ b/src/ast/function_test.cc
@@ -40,7 +40,8 @@ TEST_F(FunctionTest, Creation) {
   params.push_back(create<Variable>("var", StorageClass::kNone, &i32));
   auto* var_ptr = params[0].get();
 
-  Function f("func", std::move(params), &void_type);
+  Function f("func", std::move(params), &void_type,
+             create<ast::BlockStatement>());
   EXPECT_EQ(f.name(), "func");
   ASSERT_EQ(f.params().size(), 1u);
   EXPECT_EQ(f.return_type(), &void_type);
@@ -55,7 +56,7 @@ TEST_F(FunctionTest, Creation_WithSource) {
   params.push_back(create<Variable>("var", StorageClass::kNone, &i32));
 
   Function f(Source{Source::Location{20, 2}}, "func", std::move(params),
-             &void_type);
+             &void_type, create<ast::BlockStatement>());
   auto src = f.source();
   EXPECT_EQ(src.range.begin.line, 20u);
   EXPECT_EQ(src.range.begin.column, 2u);
@@ -66,7 +67,7 @@ TEST_F(FunctionTest, AddDuplicateReferencedVariables) {
   type::I32Type i32;
 
   Variable v("var", StorageClass::kInput, &i32);
-  Function f("func", VariableList{}, &void_type);
+  Function f("func", VariableList{}, &void_type, create<ast::BlockStatement>());
 
   f.add_referenced_module_variable(&v);
   ASSERT_EQ(f.referenced_module_variables().size(), 1u);
@@ -108,7 +109,7 @@ TEST_F(FunctionTest, GetReferenceLocations) {
       create<ast::BuiltinDecoration>(ast::Builtin::kFragDepth, Source{}));
   builtin2.set_decorations(std::move(decos));
 
-  Function f("func", VariableList{}, &void_type);
+  Function f("func", VariableList{}, &void_type, create<ast::BlockStatement>());
 
   f.add_referenced_module_variable(&loc1);
   f.add_referenced_module_variable(&builtin1);
@@ -151,7 +152,7 @@ TEST_F(FunctionTest, GetReferenceBuiltins) {
       create<ast::BuiltinDecoration>(ast::Builtin::kFragDepth, Source{}));
   builtin2.set_decorations(std::move(decos));
 
-  Function f("func", VariableList{}, &void_type);
+  Function f("func", VariableList{}, &void_type, create<ast::BlockStatement>());
 
   f.add_referenced_module_variable(&loc1);
   f.add_referenced_module_variable(&builtin1);
@@ -169,7 +170,7 @@ TEST_F(FunctionTest, GetReferenceBuiltins) {
 
 TEST_F(FunctionTest, AddDuplicateEntryPoints) {
   ast::type::VoidType void_type;
-  Function f("func", VariableList{}, &void_type);
+  Function f("func", VariableList{}, &void_type, create<ast::BlockStatement>());
 
   f.add_ancestor_entry_point("main");
   ASSERT_EQ(1u, f.ancestor_entry_points().size());
@@ -190,7 +191,8 @@ TEST_F(FunctionTest, IsValid) {
   auto block = create<ast::BlockStatement>();
   block->append(create<DiscardStatement>());
 
-  Function f("func", std::move(params), &void_type);
+  Function f("func", std::move(params), &void_type,
+             create<ast::BlockStatement>());
   f.set_body(std::move(block));
   EXPECT_TRUE(f.IsValid());
 }
@@ -202,7 +204,7 @@ TEST_F(FunctionTest, IsValid_EmptyName) {
   VariableList params;
   params.push_back(create<Variable>("var", StorageClass::kNone, &i32));
 
-  Function f("", std::move(params), &void_type);
+  Function f("", std::move(params), &void_type, create<ast::BlockStatement>());
   EXPECT_FALSE(f.IsValid());
 }
 
@@ -212,7 +214,7 @@ TEST_F(FunctionTest, IsValid_MissingReturnType) {
   VariableList params;
   params.push_back(create<Variable>("var", StorageClass::kNone, &i32));
 
-  Function f("func", std::move(params), nullptr);
+  Function f("func", std::move(params), nullptr, create<ast::BlockStatement>());
   EXPECT_FALSE(f.IsValid());
 }
 
@@ -224,7 +226,8 @@ TEST_F(FunctionTest, IsValid_NullParam) {
   params.push_back(create<Variable>("var", StorageClass::kNone, &i32));
   params.push_back(nullptr);
 
-  Function f("func", std::move(params), &void_type);
+  Function f("func", std::move(params), &void_type,
+             create<ast::BlockStatement>());
   EXPECT_FALSE(f.IsValid());
 }
 
@@ -234,7 +237,8 @@ TEST_F(FunctionTest, IsValid_InvalidParam) {
   VariableList params;
   params.push_back(create<Variable>("var", StorageClass::kNone, nullptr));
 
-  Function f("func", std::move(params), &void_type);
+  Function f("func", std::move(params), &void_type,
+             create<ast::BlockStatement>());
   EXPECT_FALSE(f.IsValid());
 }
 
@@ -249,7 +253,8 @@ TEST_F(FunctionTest, IsValid_NullBodyStatement) {
   block->append(create<DiscardStatement>());
   block->append(nullptr);
 
-  Function f("func", std::move(params), &void_type);
+  Function f("func", std::move(params), &void_type,
+             create<ast::BlockStatement>());
   f.set_body(std::move(block));
   EXPECT_FALSE(f.IsValid());
 }
@@ -265,7 +270,8 @@ TEST_F(FunctionTest, IsValid_InvalidBodyStatement) {
   block->append(create<DiscardStatement>());
   block->append(nullptr);
 
-  Function f("func", std::move(params), &void_type);
+  Function f("func", std::move(params), &void_type,
+             create<ast::BlockStatement>());
   f.set_body(std::move(block));
   EXPECT_FALSE(f.IsValid());
 }
@@ -277,7 +283,7 @@ TEST_F(FunctionTest, ToStr) {
   auto block = create<ast::BlockStatement>();
   block->append(create<DiscardStatement>());
 
-  Function f("func", {}, &void_type);
+  Function f("func", {}, &void_type, create<ast::BlockStatement>());
   f.set_body(std::move(block));
 
   std::ostringstream out;
@@ -297,7 +303,7 @@ TEST_F(FunctionTest, ToStr_WithDecoration) {
   auto block = create<ast::BlockStatement>();
   block->append(create<DiscardStatement>());
 
-  Function f("func", {}, &void_type);
+  Function f("func", {}, &void_type, create<ast::BlockStatement>());
   f.set_body(std::move(block));
   f.add_decoration(create<WorkgroupDecoration>(2, 4, 6, Source{}));
 
@@ -322,7 +328,8 @@ TEST_F(FunctionTest, ToStr_WithParams) {
   auto block = create<ast::BlockStatement>();
   block->append(create<DiscardStatement>());
 
-  Function f("func", std::move(params), &void_type);
+  Function f("func", std::move(params), &void_type,
+             create<ast::BlockStatement>());
   f.set_body(std::move(block));
 
   std::ostringstream out;
@@ -344,7 +351,7 @@ TEST_F(FunctionTest, ToStr_WithParams) {
 TEST_F(FunctionTest, TypeName) {
   type::VoidType void_type;
 
-  Function f("func", {}, &void_type);
+  Function f("func", {}, &void_type, create<ast::BlockStatement>());
   EXPECT_EQ(f.type_name(), "__func__void");
 }
 
@@ -357,7 +364,8 @@ TEST_F(FunctionTest, TypeName_WithParams) {
   params.push_back(create<Variable>("var1", StorageClass::kNone, &i32));
   params.push_back(create<Variable>("var2", StorageClass::kNone, &f32));
 
-  Function f("func", std::move(params), &void_type);
+  Function f("func", std::move(params), &void_type,
+             create<ast::BlockStatement>());
   EXPECT_EQ(f.type_name(), "__func__void__i32__f32");
 }
 
@@ -369,7 +377,8 @@ TEST_F(FunctionTest, GetLastStatement) {
   auto stmt = create<DiscardStatement>();
   auto* stmt_ptr = stmt.get();
   body->append(std::move(stmt));
-  Function f("func", std::move(params), &void_type);
+  Function f("func", std::move(params), &void_type,
+             create<ast::BlockStatement>());
   f.set_body(std::move(body));
 
   EXPECT_EQ(f.get_last_statement(), stmt_ptr);
@@ -380,7 +389,8 @@ TEST_F(FunctionTest, GetLastStatement_nullptr) {
 
   VariableList params;
   auto body = create<ast::BlockStatement>();
-  Function f("func", std::move(params), &void_type);
+  Function f("func", std::move(params), &void_type,
+             create<ast::BlockStatement>());
   f.set_body(std::move(body));
 
   EXPECT_EQ(f.get_last_statement(), nullptr);
@@ -388,7 +398,7 @@ TEST_F(FunctionTest, GetLastStatement_nullptr) {
 
 TEST_F(FunctionTest, WorkgroupSize_NoneSet) {
   type::VoidType void_type;
-  Function f("f", {}, &void_type);
+  Function f("f", {}, &void_type, create<ast::BlockStatement>());
   uint32_t x = 0;
   uint32_t y = 0;
   uint32_t z = 0;
@@ -400,7 +410,7 @@ TEST_F(FunctionTest, WorkgroupSize_NoneSet) {
 
 TEST_F(FunctionTest, WorkgroupSize) {
   type::VoidType void_type;
-  Function f("f", {}, &void_type);
+  Function f("f", {}, &void_type, create<ast::BlockStatement>());
   f.add_decoration(create<WorkgroupDecoration>(2u, 4u, 6u, Source{}));
 
   uint32_t x = 0;
diff --git a/src/ast/module_test.cc b/src/ast/module_test.cc
index eaddc076e1..9b6af13473 100644
--- a/src/ast/module_test.cc
+++ b/src/ast/module_test.cc
@@ -48,7 +48,8 @@ TEST_F(ModuleTest, LookupFunction) {
   type::F32Type f32;
   Module m;
 
-  auto func = create<Function>("main", VariableList{}, &f32);
+  auto func = create<Function>("main", VariableList{}, &f32,
+                               create<ast::BlockStatement>());
   auto* func_ptr = func.get();
   m.AddFunction(std::move(func));
   EXPECT_EQ(func_ptr, m.FindFunctionByName("main"));
@@ -124,7 +125,8 @@ TEST_F(ModuleTest, IsValid_Struct_EmptyName) {
 
 TEST_F(ModuleTest, IsValid_Function) {
   type::F32Type f32;
-  auto func = create<Function>("main", VariableList(), &f32);
+  auto func = create<Function>("main", VariableList(), &f32,
+                               create<ast::BlockStatement>());
 
   Module m;
   m.AddFunction(std::move(func));
diff --git a/src/inspector/inspector_test.cc b/src/inspector/inspector_test.cc
index da3848e117..03104bd02d 100644
--- a/src/inspector/inspector_test.cc
+++ b/src/inspector/inspector_test.cc
@@ -81,10 +81,8 @@ class InspectorHelper {
   std::unique_ptr<ast::Function> MakeEmptyBodyFunction(std::string name) {
     auto body = create<ast::BlockStatement>();
     body->append(create<ast::ReturnStatement>());
-    std::unique_ptr<ast::Function> func =
-        create<ast::Function>(name, ast::VariableList(), void_type());
-    func->set_body(std::move(body));
-    return func;
+    return create<ast::Function>(name, ast::VariableList(), void_type(),
+                                 std::move(body));
   }
 
   /// Generates a function that calls another
@@ -99,10 +97,8 @@ class InspectorHelper {
                                                  ast::ExpressionList());
     body->append(create<ast::CallStatement>(std::move(call_expr)));
     body->append(create<ast::ReturnStatement>());
-    std::unique_ptr<ast::Function> func =
-        create<ast::Function>(caller, ast::VariableList(), void_type());
-    func->set_body(std::move(body));
-    return func;
+    return create<ast::Function>(caller, ast::VariableList(), void_type(),
+                                 std::move(body));
   }
 
   /// Add In/Out variables to the global variables
@@ -139,9 +135,8 @@ class InspectorHelper {
           create<ast::IdentifierExpression>(in)));
     }
     body->append(create<ast::ReturnStatement>());
-    auto func = create<ast::Function>(name, ast::VariableList(), void_type());
-    func->set_body(std::move(body));
-    return func;
+    return create<ast::Function>(name, ast::VariableList(), void_type(),
+                                 std::move(body));
   }
 
   /// Generates a function that references in/out variables and calls another
@@ -168,9 +163,8 @@ class InspectorHelper {
                                                  ast::ExpressionList());
     body->append(create<ast::CallStatement>(std::move(call_expr)));
     body->append(create<ast::ReturnStatement>());
-    auto func = create<ast::Function>(caller, ast::VariableList(), void_type());
-    func->set_body(std::move(body));
-    return func;
+    return create<ast::Function>(caller, ast::VariableList(), void_type(),
+                                 std::move(body));
   }
 
   /// Add a Constant ID to the global variables.
@@ -415,10 +409,8 @@ class InspectorHelper {
     }
 
     body->append(create<ast::ReturnStatement>());
-    auto func =
-        create<ast::Function>(func_name, ast::VariableList(), void_type());
-    func->set_body(std::move(body));
-    return func;
+    return create<ast::Function>(func_name, ast::VariableList(), void_type(),
+                                 std::move(body));
   }
 
   /// Adds a regular sampler variable to the module
@@ -539,10 +531,8 @@ class InspectorHelper {
         std::move(call_expr)));
     body->append(create<ast::ReturnStatement>());
 
-    auto func =
-        create<ast::Function>(func_name, ast::VariableList(), void_type());
-    func->set_body(std::move(body));
-    return func;
+    return create<ast::Function>(func_name, ast::VariableList(), void_type(),
+                                 std::move(body));
   }
 
   /// Generates a function that references a specific comparison sampler
@@ -582,10 +572,8 @@ class InspectorHelper {
         std::move(call_expr)));
     body->append(create<ast::ReturnStatement>());
 
-    auto func =
-        create<ast::Function>(func_name, ast::VariableList(), void_type());
-    func->set_body(std::move(body));
-    return func;
+    return create<ast::Function>(func_name, ast::VariableList(), void_type(),
+                                 std::move(body));
   }
 
   /// Gets an appropriate type for the data in a given texture type.
@@ -1345,9 +1333,8 @@ TEST_F(InspectorGetUniformBufferResourceBindingsTest, MultipleUniformBuffers) {
   AddFuncCall(body.get(), "ub_baz_func");
 
   body->append(create<ast::ReturnStatement>());
-  std::unique_ptr<ast::Function> func =
-      create<ast::Function>("ep_func", ast::VariableList(), void_type());
-  func->set_body(std::move(body));
+  std::unique_ptr<ast::Function> func = create<ast::Function>(
+      "ep_func", ast::VariableList(), void_type(), std::move(body));
 
   func->add_decoration(
       create<ast::StageDecoration>(ast::PipelineStage::kVertex, Source{}));
@@ -1486,9 +1473,8 @@ TEST_F(InspectorGetStorageBufferResourceBindingsTest, MultipleStorageBuffers) {
   AddFuncCall(body.get(), "sb_baz_func");
 
   body->append(create<ast::ReturnStatement>());
-  std::unique_ptr<ast::Function> func =
-      create<ast::Function>("ep_func", ast::VariableList(), void_type());
-  func->set_body(std::move(body));
+  std::unique_ptr<ast::Function> func = create<ast::Function>(
+      "ep_func", ast::VariableList(), void_type(), std::move(body));
 
   func->add_decoration(
       create<ast::StageDecoration>(ast::PipelineStage::kVertex, Source{}));
@@ -1652,9 +1638,8 @@ TEST_F(InspectorGetReadOnlyStorageBufferResourceBindingsTest,
   AddFuncCall(body.get(), "sb_baz_func");
 
   body->append(create<ast::ReturnStatement>());
-  std::unique_ptr<ast::Function> func =
-      create<ast::Function>("ep_func", ast::VariableList(), void_type());
-  func->set_body(std::move(body));
+  std::unique_ptr<ast::Function> func = create<ast::Function>(
+      "ep_func", ast::VariableList(), void_type(), std::move(body));
 
   func->add_decoration(
       create<ast::StageDecoration>(ast::PipelineStage::kVertex, Source{}));
diff --git a/src/reader/spirv/function.cc b/src/reader/spirv/function.cc
index b42713fbbc..2ab7cf52f1 100644
--- a/src/reader/spirv/function.cc
+++ b/src/reader/spirv/function.cc
@@ -646,7 +646,8 @@ bool FunctionEmitter::EmitFunctionDeclaration() {
   }
 
   auto ast_fn =
-      std::make_unique<ast::Function>(name, std::move(ast_params), ret_ty);
+      std::make_unique<ast::Function>(name, std::move(ast_params), ret_ty,
+                                      std::make_unique<ast::BlockStatement>());
 
   if (ep_info_ != nullptr) {
     ast_fn->add_decoration(
diff --git a/src/reader/wgsl/parser_impl.cc b/src/reader/wgsl/parser_impl.cc
index d3c357b6cd..e70d2a4d31 100644
--- a/src/reader/wgsl/parser_impl.cc
+++ b/src/reader/wgsl/parser_impl.cc
@@ -1231,8 +1231,9 @@ Maybe<std::unique_ptr<ast::Function>> ParserImpl::function_header() {
   if (errored)
     return Failure::kErrored;
 
-  return std::make_unique<ast::Function>(source, name.value,
-                                         std::move(params.value), type.value);
+  return std::make_unique<ast::Function>(
+      source, name.value, std::move(params.value), type.value,
+      std::make_unique<ast::BlockStatement>());
 }
 
 // param_list
diff --git a/src/transform/bound_array_accessors_transform_test.cc b/src/transform/bound_array_accessors_transform_test.cc
index e165ef03ef..26c897388e 100644
--- a/src/transform/bound_array_accessors_transform_test.cc
+++ b/src/transform/bound_array_accessors_transform_test.cc
@@ -52,10 +52,10 @@ class BoundArrayAccessorsTest : public testing::Test {
   BoundArrayAccessorsTest() : td_(&ctx_, &mod_), transform_(&ctx_, &mod_) {}
 
   ast::BlockStatement* SetupFunctionAndBody() {
-    auto func = create<ast::Function>("func", ast::VariableList{}, &void_type_);
     auto block = create<ast::BlockStatement>();
     body_ = block.get();
-    func->set_body(std::move(block));
+    auto func = create<ast::Function>("func", ast::VariableList{}, &void_type_,
+                                      std::move(block));
     mod_.AddFunction(std::move(func));
     return body_;
   }
diff --git a/src/transform/vertex_pulling_transform_test.cc b/src/transform/vertex_pulling_transform_test.cc
index a9b0c8c2a9..2b2d34bef8 100644
--- a/src/transform/vertex_pulling_transform_test.cc
+++ b/src/transform/vertex_pulling_transform_test.cc
@@ -43,7 +43,8 @@ class VertexPullingTransformHelper {
   void InitBasicModule() {
     auto func = create<ast::Function>(
         "main", ast::VariableList{},
-        ctx_.type_mgr().Get(std::make_unique<ast::type::VoidType>()));
+        ctx_.type_mgr().Get(std::make_unique<ast::type::VoidType>()),
+        create<ast::BlockStatement>());
     func->add_decoration(
         create<ast::StageDecoration>(ast::PipelineStage ::kVertex, Source{}));
     mod()->AddFunction(std::move(func));
@@ -118,7 +119,8 @@ TEST_F(VertexPullingTransformTest, Error_InvalidEntryPoint) {
 TEST_F(VertexPullingTransformTest, Error_EntryPointWrongStage) {
   auto func = create<ast::Function>(
       "main", ast::VariableList{},
-      ctx()->type_mgr().Get(std::make_unique<ast::type::VoidType>()));
+      ctx()->type_mgr().Get(std::make_unique<ast::type::VoidType>()),
+      create<ast::BlockStatement>());
   func->add_decoration(
       create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
   mod()->AddFunction(std::move(func));
diff --git a/src/type_determiner_test.cc b/src/type_determiner_test.cc
index 970c0bf224..089a14c99b 100644
--- a/src/type_determiner_test.cc
+++ b/src/type_determiner_test.cc
@@ -381,7 +381,8 @@ TEST_F(TypeDeterminerTest, Stmt_Call) {
   ast::type::F32Type f32;
 
   ast::VariableList params;
-  auto func = create<ast::Function>("my_func", std::move(params), &f32);
+  auto func = create<ast::Function>("my_func", std::move(params), &f32,
+                                    create<ast::BlockStatement>());
   mod()->AddFunction(std::move(func));
 
   // Register the function
@@ -408,17 +409,17 @@ TEST_F(TypeDeterminerTest, Stmt_Call_undeclared) {
                                       Source{Source::Location{12, 34}}, "func"),
                                   std::move(call_params));
   ast::VariableList params0;
-  auto func_main = create<ast::Function>("main", std::move(params0), &f32);
   auto main_body = create<ast::BlockStatement>();
   main_body->append(create<ast::CallStatement>(std::move(call_expr)));
   main_body->append(create<ast::ReturnStatement>());
-  func_main->set_body(std::move(main_body));
+  auto func_main = create<ast::Function>("main", std::move(params0), &f32,
+                                         std::move(main_body));
   mod()->AddFunction(std::move(func_main));
 
-  auto func = create<ast::Function>("func", std::move(params0), &f32);
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::ReturnStatement>());
-  func->set_body(std::move(body));
+  auto func =
+      create<ast::Function>("func", std::move(params0), &f32, std::move(body));
   mod()->AddFunction(std::move(func));
 
   EXPECT_FALSE(td()->Determine()) << td()->error();
@@ -625,7 +626,8 @@ TEST_F(TypeDeterminerTest, Expr_Call) {
   ast::type::F32Type f32;
 
   ast::VariableList params;
-  auto func = create<ast::Function>("my_func", std::move(params), &f32);
+  auto func = create<ast::Function>("my_func", std::move(params), &f32,
+                                    create<ast::BlockStatement>());
   mod()->AddFunction(std::move(func));
 
   // Register the function
@@ -643,7 +645,8 @@ TEST_F(TypeDeterminerTest, Expr_Call_WithParams) {
   ast::type::F32Type f32;
 
   ast::VariableList params;
-  auto func = create<ast::Function>("my_func", std::move(params), &f32);
+  auto func = create<ast::Function>("my_func", std::move(params), &f32,
+                                    create<ast::BlockStatement>());
   mod()->AddFunction(std::move(func));
 
   // Register the function
@@ -769,8 +772,7 @@ TEST_F(TypeDeterminerTest, Expr_Identifier_FunctionVariable_Const) {
   body->append(create<ast::AssignmentStatement>(
       std::move(my_var), create<ast::IdentifierExpression>("my_var")));
 
-  ast::Function f("my_func", {}, &f32);
-  f.set_body(std::move(body));
+  ast::Function f("my_func", {}, &f32, std::move(body));
 
   EXPECT_TRUE(td()->DetermineFunction(&f));
 
@@ -791,8 +793,7 @@ TEST_F(TypeDeterminerTest, Expr_Identifier_FunctionVariable) {
   body->append(create<ast::AssignmentStatement>(
       std::move(my_var), create<ast::IdentifierExpression>("my_var")));
 
-  ast::Function f("my_func", {}, &f32);
-  f.set_body(std::move(body));
+  ast::Function f("my_func", {}, &f32, std::move(body));
 
   EXPECT_TRUE(td()->DetermineFunction(&f));
 
@@ -815,8 +816,7 @@ TEST_F(TypeDeterminerTest, Expr_Identifier_Function_Ptr) {
   body->append(create<ast::AssignmentStatement>(
       std::move(my_var), create<ast::IdentifierExpression>("my_var")));
 
-  ast::Function f("my_func", {}, &f32);
-  f.set_body(std::move(body));
+  ast::Function f("my_func", {}, &f32, std::move(body));
 
   EXPECT_TRUE(td()->DetermineFunction(&f));
 
@@ -829,7 +829,8 @@ TEST_F(TypeDeterminerTest, Expr_Identifier_Function) {
   ast::type::F32Type f32;
 
   ast::VariableList params;
-  auto func = create<ast::Function>("my_func", std::move(params), &f32);
+  auto func = create<ast::Function>("my_func", std::move(params), &f32,
+                                    create<ast::BlockStatement>());
   mod()->AddFunction(std::move(func));
 
   // Register the function
@@ -873,9 +874,6 @@ TEST_F(TypeDeterminerTest, Function_RegisterInputOutputVariables) {
   mod()->AddGlobalVariable(std::move(priv_var));
 
   ast::VariableList params;
-  auto func = create<ast::Function>("my_func", std::move(params), &f32);
-  auto* func_ptr = func.get();
-
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::AssignmentStatement>(
       create<ast::IdentifierExpression>("out_var"),
@@ -889,7 +887,9 @@ TEST_F(TypeDeterminerTest, Function_RegisterInputOutputVariables) {
   body->append(create<ast::AssignmentStatement>(
       create<ast::IdentifierExpression>("priv_var"),
       create<ast::IdentifierExpression>("priv_var")));
-  func->set_body(std::move(body));
+  auto func = create<ast::Function>("my_func", std::move(params), &f32,
+                                    std::move(body));
+  auto* func_ptr = func.get();
 
   mod()->AddFunction(std::move(func));
 
@@ -931,9 +931,6 @@ TEST_F(TypeDeterminerTest, Function_RegisterInputOutputVariables_SubFunction) {
   mod()->AddGlobalVariable(std::move(wg_var));
   mod()->AddGlobalVariable(std::move(priv_var));
 
-  ast::VariableList params;
-  auto func = create<ast::Function>("my_func", std::move(params), &f32);
-
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::AssignmentStatement>(
       create<ast::IdentifierExpression>("out_var"),
@@ -947,19 +944,20 @@ TEST_F(TypeDeterminerTest, Function_RegisterInputOutputVariables_SubFunction) {
   body->append(create<ast::AssignmentStatement>(
       create<ast::IdentifierExpression>("priv_var"),
       create<ast::IdentifierExpression>("priv_var")));
-  func->set_body(std::move(body));
+  ast::VariableList params;
+  auto func = create<ast::Function>("my_func", std::move(params), &f32,
+                                    std::move(body));
 
   mod()->AddFunction(std::move(func));
 
-  auto func2 = create<ast::Function>("func", std::move(params), &f32);
-  auto* func2_ptr = func2.get();
-
   body = create<ast::BlockStatement>();
   body->append(create<ast::AssignmentStatement>(
       create<ast::IdentifierExpression>("out_var"),
       create<ast::CallExpression>(create<ast::IdentifierExpression>("my_func"),
                                   ast::ExpressionList{})));
-  func2->set_body(std::move(body));
+  auto func2 =
+      create<ast::Function>("func", std::move(params), &f32, std::move(body));
+  auto* func2_ptr = func2.get();
 
   mod()->AddFunction(std::move(func2));
 
@@ -981,17 +979,17 @@ TEST_F(TypeDeterminerTest, Function_NotRegisterFunctionVariable) {
   auto var =
       create<ast::Variable>("in_var", ast::StorageClass::kFunction, &f32);
 
-  ast::VariableList params;
-  auto func = create<ast::Function>("my_func", std::move(params), &f32);
-  auto* func_ptr = func.get();
-
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::VariableDeclStatement>(std::move(var)));
   body->append(create<ast::AssignmentStatement>(
       create<ast::IdentifierExpression>("var"),
       create<ast::ScalarConstructorExpression>(
           create<ast::FloatLiteral>(&f32, 1.f))));
-  func->set_body(std::move(body));
+
+  ast::VariableList params;
+  auto func = create<ast::Function>("my_func", std::move(params), &f32,
+                                    std::move(body));
+  auto* func_ptr = func.get();
 
   mod()->AddFunction(std::move(func));
 
@@ -2339,11 +2337,10 @@ TEST_F(TypeDeterminerTest, StorageClass_SetsIfMissing) {
   auto* var_ptr = var.get();
   auto stmt = create<ast::VariableDeclStatement>(std::move(var));
 
-  auto func = create<ast::Function>("func", ast::VariableList{}, &i32);
-
   auto body = create<ast::BlockStatement>();
   body->append(std::move(stmt));
-  func->set_body(std::move(body));
+  auto func =
+      create<ast::Function>("func", ast::VariableList{}, &i32, std::move(body));
 
   mod()->AddFunction(std::move(func));
 
@@ -2359,11 +2356,10 @@ TEST_F(TypeDeterminerTest, StorageClass_DoesNotSetOnConst) {
   auto* var_ptr = var.get();
   auto stmt = create<ast::VariableDeclStatement>(std::move(var));
 
-  auto func = create<ast::Function>("func", ast::VariableList{}, &i32);
-
   auto body = create<ast::BlockStatement>();
   body->append(std::move(stmt));
-  func->set_body(std::move(body));
+  auto func =
+      create<ast::Function>("func", ast::VariableList{}, &i32, std::move(body));
 
   mod()->AddFunction(std::move(func));
 
@@ -2377,11 +2373,10 @@ TEST_F(TypeDeterminerTest, StorageClass_NonFunctionClassError) {
   auto var = create<ast::Variable>("var", ast::StorageClass::kWorkgroup, &i32);
   auto stmt = create<ast::VariableDeclStatement>(std::move(var));
 
-  auto func = create<ast::Function>("func", ast::VariableList{}, &i32);
-
   auto body = create<ast::BlockStatement>();
   body->append(std::move(stmt));
-  func->set_body(std::move(body));
+  auto func =
+      create<ast::Function>("func", ast::VariableList{}, &i32, std::move(body));
 
   mod()->AddFunction(std::move(func));
 
@@ -4583,36 +4578,28 @@ TEST_F(TypeDeterminerTest, Function_EntryPoints_StageDecoration) {
   // ep_2 -> {}
 
   ast::VariableList params;
-  auto func_b = create<ast::Function>("b", std::move(params), &f32);
-  auto* func_b_ptr = func_b.get();
-
   auto body = create<ast::BlockStatement>();
-  func_b->set_body(std::move(body));
-
-  auto func_c = create<ast::Function>("c", std::move(params), &f32);
-  auto* func_c_ptr = func_c.get();
+  auto func_b =
+      create<ast::Function>("b", std::move(params), &f32, std::move(body));
+  auto* func_b_ptr = func_b.get();
 
   body = create<ast::BlockStatement>();
   body->append(create<ast::AssignmentStatement>(
       create<ast::IdentifierExpression>("second"),
       create<ast::CallExpression>(create<ast::IdentifierExpression>("b"),
                                   ast::ExpressionList{})));
-  func_c->set_body(std::move(body));
-
-  auto func_a = create<ast::Function>("a", std::move(params), &f32);
-  auto* func_a_ptr = func_a.get();
+  auto func_c =
+      create<ast::Function>("c", std::move(params), &f32, std::move(body));
+  auto* func_c_ptr = func_c.get();
 
   body = create<ast::BlockStatement>();
   body->append(create<ast::AssignmentStatement>(
       create<ast::IdentifierExpression>("first"),
       create<ast::CallExpression>(create<ast::IdentifierExpression>("c"),
                                   ast::ExpressionList{})));
-  func_a->set_body(std::move(body));
-
-  auto ep_1 = create<ast::Function>("ep_1", std::move(params), &f32);
-  ep_1->add_decoration(
-      create<ast::StageDecoration>(ast::PipelineStage::kVertex, Source{}));
-  auto* ep_1_ptr = ep_1.get();
+  auto func_a =
+      create<ast::Function>("a", std::move(params), &f32, std::move(body));
+  auto* func_a_ptr = func_a.get();
 
   body = create<ast::BlockStatement>();
   body->append(create<ast::AssignmentStatement>(
@@ -4623,19 +4610,22 @@ TEST_F(TypeDeterminerTest, Function_EntryPoints_StageDecoration) {
       create<ast::IdentifierExpression>("call_b"),
       create<ast::CallExpression>(create<ast::IdentifierExpression>("b"),
                                   ast::ExpressionList{})));
-  ep_1->set_body(std::move(body));
-
-  auto ep_2 = create<ast::Function>("ep_2", std::move(params), &f32);
-  ep_2->add_decoration(
+  auto ep_1 =
+      create<ast::Function>("ep_1", std::move(params), &f32, std::move(body));
+  ep_1->add_decoration(
       create<ast::StageDecoration>(ast::PipelineStage::kVertex, Source{}));
-  auto* ep_2_ptr = ep_2.get();
+  auto* ep_1_ptr = ep_1.get();
 
   body = create<ast::BlockStatement>();
   body->append(create<ast::AssignmentStatement>(
       create<ast::IdentifierExpression>("call_c"),
       create<ast::CallExpression>(create<ast::IdentifierExpression>("c"),
                                   ast::ExpressionList{})));
-  ep_2->set_body(std::move(body));
+  auto ep_2 =
+      create<ast::Function>("ep_2", std::move(params), &f32, std::move(body));
+  ep_2->add_decoration(
+      create<ast::StageDecoration>(ast::PipelineStage::kVertex, Source{}));
+  auto* ep_2_ptr = ep_2.get();
 
   mod()->AddFunction(std::move(func_b));
   mod()->AddFunction(std::move(func_c));
diff --git a/src/validator/validator_function_test.cc b/src/validator/validator_function_test.cc
index 6ad4acd1c3..97abf87771 100644
--- a/src/validator/validator_function_test.cc
+++ b/src/validator/validator_function_test.cc
@@ -45,11 +45,11 @@ TEST_F(ValidateFunctionTest, VoidFunctionEndWithoutReturnStatement_Pass) {
 
   ast::VariableList params;
   ast::type::VoidType void_type;
-  auto func = create<ast::Function>(Source{Source::Location{12, 34}}, "func",
-                                    std::move(params), &void_type);
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::VariableDeclStatement>(std::move(var)));
-  func->set_body(std::move(body));
+  auto func =
+      create<ast::Function>(Source{Source::Location{12, 34}}, "func",
+                            std::move(params), &void_type, std::move(body));
   func->add_decoration(
       create<ast::StageDecoration>(ast::PipelineStage::kVertex, Source{}));
   mod()->AddFunction(std::move(func));
@@ -65,7 +65,8 @@ TEST_F(ValidateFunctionTest,
   ast::type::VoidType void_type;
   ast::VariableList params;
   auto func = create<ast::Function>(Source{Source::Location{12, 34}}, "func",
-                                    std::move(params), &void_type);
+                                    std::move(params), &void_type,
+                                    create<ast::BlockStatement>());
   func->add_decoration(
       create<ast::StageDecoration>(ast::PipelineStage::kVertex, Source{}));
   mod()->AddFunction(std::move(func));
@@ -84,11 +85,10 @@ TEST_F(ValidateFunctionTest, FunctionEndWithoutReturnStatement_Fail) {
 
   ast::VariableList params;
   ast::type::VoidType void_type;
-  auto func = create<ast::Function>(Source{Source::Location{12, 34}}, "func",
-                                    std::move(params), &i32);
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::VariableDeclStatement>(std::move(var)));
-  func->set_body(std::move(body));
+  auto func = create<ast::Function>(Source{Source::Location{12, 34}}, "func",
+                                    std::move(params), &i32, std::move(body));
   mod()->AddFunction(std::move(func));
 
   EXPECT_TRUE(td()->Determine()) << td()->error();
@@ -104,7 +104,8 @@ TEST_F(ValidateFunctionTest, FunctionEndWithoutReturnStatementEmptyBody_Fail) {
   ast::type::I32Type i32;
   ast::VariableList params;
   auto func = create<ast::Function>(Source{Source::Location{12, 34}}, "func",
-                                    std::move(params), &i32);
+                                    std::move(params), &i32,
+                                    create<ast::BlockStatement>());
   mod()->AddFunction(std::move(func));
 
   EXPECT_TRUE(td()->Determine()) << td()->error();
@@ -120,10 +121,10 @@ TEST_F(ValidateFunctionTest, FunctionTypeMustMatchReturnStatementType_Pass) {
   ast::type::VoidType void_type;
   ast::VariableList params;
 
-  auto func = create<ast::Function>("func", std::move(params), &void_type);
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::ReturnStatement>());
-  func->set_body(std::move(body));
+  auto func = create<ast::Function>("func", std::move(params), &void_type,
+                                    std::move(body));
   func->add_decoration(
       create<ast::StageDecoration>(ast::PipelineStage::kVertex, Source{}));
   mod()->AddFunction(std::move(func));
@@ -137,14 +138,14 @@ TEST_F(ValidateFunctionTest, FunctionTypeMustMatchReturnStatementType_fail) {
   ast::type::VoidType void_type;
   ast::type::I32Type i32;
   ast::VariableList params;
-  auto func = create<ast::Function>("func", std::move(params), &void_type);
   auto body = create<ast::BlockStatement>();
   auto return_expr = create<ast::ScalarConstructorExpression>(
       create<ast::SintLiteral>(&i32, 2));
 
   body->append(create<ast::ReturnStatement>(Source{Source::Location{12, 34}},
                                             std::move(return_expr)));
-  func->set_body(std::move(body));
+  auto func = create<ast::Function>("func", std::move(params), &void_type,
+                                    std::move(body));
   mod()->AddFunction(std::move(func));
 
   EXPECT_TRUE(td()->Determine()) << td()->error();
@@ -160,14 +161,14 @@ TEST_F(ValidateFunctionTest, FunctionTypeMustMatchReturnStatementTypeF32_fail) {
   ast::type::I32Type i32;
   ast::type::F32Type f32;
   ast::VariableList params;
-  auto func = create<ast::Function>("func", std::move(params), &f32);
   auto body = create<ast::BlockStatement>();
   auto return_expr = create<ast::ScalarConstructorExpression>(
       create<ast::SintLiteral>(&i32, 2));
 
   body->append(create<ast::ReturnStatement>(Source{Source::Location{12, 34}},
                                             std::move(return_expr)));
-  func->set_body(std::move(body));
+  auto func =
+      create<ast::Function>("func", std::move(params), &f32, std::move(body));
   mod()->AddFunction(std::move(func));
 
   EXPECT_TRUE(td()->Determine()) << td()->error();
@@ -185,23 +186,23 @@ TEST_F(ValidateFunctionTest, FunctionNamesMustBeUnique_fail) {
   ast::type::I32Type i32;
 
   ast::VariableList params;
-  auto func = create<ast::Function>("func", std::move(params), &i32);
   auto body = create<ast::BlockStatement>();
   auto return_expr = create<ast::ScalarConstructorExpression>(
       create<ast::SintLiteral>(&i32, 2));
 
   body->append(create<ast::ReturnStatement>(std::move(return_expr)));
-  func->set_body(std::move(body));
+  auto func =
+      create<ast::Function>("func", std::move(params), &i32, std::move(body));
 
   ast::VariableList params_copy;
-  auto func_copy = create<ast::Function>(Source{Source::Location{12, 34}},
-                                         "func", std::move(params_copy), &i32);
   auto body_copy = create<ast::BlockStatement>();
   auto return_expr_copy = create<ast::ScalarConstructorExpression>(
       create<ast::SintLiteral>(&i32, 2));
 
   body_copy->append(create<ast::ReturnStatement>(std::move(return_expr_copy)));
-  func_copy->set_body(std::move(body_copy));
+  auto func_copy =
+      create<ast::Function>(Source{Source::Location{12, 34}}, "func",
+                            std::move(params_copy), &i32, std::move(body_copy));
 
   mod()->AddFunction(std::move(func));
   mod()->AddFunction(std::move(func_copy));
@@ -221,11 +222,11 @@ TEST_F(ValidateFunctionTest, RecursionIsNotAllowed_Fail) {
       Source{Source::Location{12, 34}},
       create<ast::IdentifierExpression>("func"), std::move(call_params));
   ast::VariableList params0;
-  auto func0 = create<ast::Function>("func", std::move(params0), &f32);
   auto body0 = create<ast::BlockStatement>();
   body0->append(create<ast::CallStatement>(std::move(call_expr)));
   body0->append(create<ast::ReturnStatement>());
-  func0->set_body(std::move(body0));
+  auto func0 =
+      create<ast::Function>("func", std::move(params0), &f32, std::move(body0));
   mod()->AddFunction(std::move(func0));
 
   EXPECT_TRUE(td()->Determine()) << td()->error();
@@ -243,14 +244,14 @@ TEST_F(ValidateFunctionTest, RecursionIsNotAllowedExpr_Fail) {
       create<ast::IdentifierExpression>("func"), std::move(call_params));
   var->set_constructor(std::move(call_expr));
   ast::VariableList params0;
-  auto func0 = create<ast::Function>("func", std::move(params0), &i32);
   auto body0 = create<ast::BlockStatement>();
   body0->append(create<ast::VariableDeclStatement>(std::move(var)));
   auto return_expr = create<ast::ScalarConstructorExpression>(
       create<ast::SintLiteral>(&i32, 2));
 
   body0->append(create<ast::ReturnStatement>(std::move(return_expr)));
-  func0->set_body(std::move(body0));
+  auto func0 =
+      create<ast::Function>("func", std::move(params0), &i32, std::move(body0));
   mod()->AddFunction(std::move(func0));
 
   EXPECT_TRUE(td()->Determine()) << td()->error();
@@ -263,14 +264,14 @@ TEST_F(ValidateFunctionTest, Function_WithPipelineStage_NotVoid_Fail) {
   // fn vtx_main() -> i32 { return 0; }
   ast::type::I32Type i32;
   ast::VariableList params;
-  auto func = create<ast::Function>(Source{Source::Location{12, 34}},
-                                    "vtx_main", std::move(params), &i32);
   auto return_expr = create<ast::ScalarConstructorExpression>(
       create<ast::SintLiteral>(&i32, 0));
 
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::ReturnStatement>(std::move(return_expr)));
-  func->set_body(std::move(body));
+  auto func =
+      create<ast::Function>(Source{Source::Location{12, 34}}, "vtx_main",
+                            std::move(params), &i32, std::move(body));
   func->add_decoration(
       create<ast::StageDecoration>(ast::PipelineStage::kVertex, Source{}));
 
@@ -288,11 +289,11 @@ TEST_F(ValidateFunctionTest, Function_WithPipelineStage_WithParams_Fail) {
   ast::type::VoidType void_type;
   ast::VariableList params;
   params.push_back(create<ast::Variable>("a", ast::StorageClass::kNone, &i32));
-  auto func = create<ast::Function>(Source{Source::Location{12, 34}},
-                                    "vtx_func", std::move(params), &void_type);
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::ReturnStatement>());
-  func->set_body(std::move(body));
+  auto func =
+      create<ast::Function>(Source{Source::Location{12, 34}}, "vtx_func",
+                            std::move(params), &void_type, std::move(body));
   func->add_decoration(
       create<ast::StageDecoration>(ast::PipelineStage::kVertex, Source{}));
 
@@ -310,11 +311,11 @@ TEST_F(ValidateFunctionTest, PipelineStage_MustBeUnique_Fail) {
   // fn main() -> void { return; }
   ast::type::VoidType void_type;
   ast::VariableList params;
-  auto func = create<ast::Function>(Source{Source::Location{12, 34}}, "main",
-                                    std::move(params), &void_type);
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::ReturnStatement>());
-  func->set_body(std::move(body));
+  auto func =
+      create<ast::Function>(Source{Source::Location{12, 34}}, "main",
+                            std::move(params), &void_type, std::move(body));
   func->add_decoration(
       create<ast::StageDecoration>(ast::PipelineStage::kVertex, Source{}));
   func->add_decoration(
@@ -332,10 +333,10 @@ TEST_F(ValidateFunctionTest, OnePipelineStageFunctionMustBePresent_Pass) {
   // fn vtx_func() -> void { return; }
   ast::type::VoidType void_type;
   ast::VariableList params;
-  auto func = create<ast::Function>("vtx_func", std::move(params), &void_type);
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::ReturnStatement>());
-  func->set_body(std::move(body));
+  auto func = create<ast::Function>("vtx_func", std::move(params), &void_type,
+                                    std::move(body));
   func->add_decoration(
       create<ast::StageDecoration>(ast::PipelineStage::kVertex, Source{}));
   mod()->AddFunction(std::move(func));
@@ -348,10 +349,10 @@ TEST_F(ValidateFunctionTest, OnePipelineStageFunctionMustBePresent_Fail) {
   // fn vtx_func() -> void { return; }
   ast::type::VoidType void_type;
   ast::VariableList params;
-  auto func = create<ast::Function>("vtx_func", std::move(params), &void_type);
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::ReturnStatement>());
-  func->set_body(std::move(body));
+  auto func = create<ast::Function>("vtx_func", std::move(params), &void_type,
+                                    std::move(body));
   mod()->AddFunction(std::move(func));
 
   EXPECT_TRUE(td()->Determine()) << td()->error();
diff --git a/src/validator/validator_test.cc b/src/validator/validator_test.cc
index a6bed90b29..4e955f9db0 100644
--- a/src/validator/validator_test.cc
+++ b/src/validator/validator_test.cc
@@ -297,12 +297,12 @@ TEST_F(ValidatorTest, UsingUndefinedVariableGlobalVariable_Fail) {
       create<ast::FloatLiteral>(&f32, 3.14f));
 
   ast::VariableList params;
-  auto func = create<ast::Function>("my_func", std::move(params), &f32);
-
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::AssignmentStatement>(
       Source{Source::Location{12, 34}}, std::move(lhs), std::move(rhs)));
-  func->set_body(std::move(body));
+
+  auto func = create<ast::Function>("my_func", std::move(params), &f32,
+                                    std::move(body));
   mod()->AddFunction(std::move(func));
 
   EXPECT_FALSE(v()->Validate(mod()));
@@ -329,13 +329,13 @@ TEST_F(ValidatorTest, UsingUndefinedVariableGlobalVariable_Pass) {
       create<ast::FloatLiteral>(&f32, 3.14f));
 
   ast::VariableList params;
-  auto func = create<ast::Function>("my_func", std::move(params), &void_type);
 
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::AssignmentStatement>(
       Source{Source::Location{12, 34}}, std::move(lhs), std::move(rhs)));
   body->append(create<ast::ReturnStatement>());
-  func->set_body(std::move(body));
+  auto func = create<ast::Function>("my_func", std::move(params), &void_type,
+                                    std::move(body));
   func->add_decoration(
       create<ast::StageDecoration>(ast::PipelineStage::kVertex, Source{}));
   mod()->AddFunction(std::move(func));
@@ -507,11 +507,11 @@ TEST_F(ValidatorTest, GlobalVariableFunctionVariableNotUnique_Fail) {
   var->set_constructor(create<ast::ScalarConstructorExpression>(
       create<ast::FloatLiteral>(&f32, 2.0)));
   ast::VariableList params;
-  auto func = create<ast::Function>("my_func", std::move(params), &void_type);
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::VariableDeclStatement>(
       Source{Source::Location{12, 34}}, std::move(var)));
-  func->set_body(std::move(body));
+  auto func = create<ast::Function>("my_func", std::move(params), &void_type,
+                                    std::move(body));
   auto* func_ptr = func.get();
   mod()->AddFunction(std::move(func));
 
@@ -538,12 +538,12 @@ TEST_F(ValidatorTest, RedeclaredIndentifier_Fail) {
       create<ast::FloatLiteral>(&f32, 0.1)));
 
   ast::VariableList params;
-  auto func = create<ast::Function>("my_func", std::move(params), &void_type);
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::VariableDeclStatement>(std::move(var)));
   body->append(create<ast::VariableDeclStatement>(
       Source{Source::Location{12, 34}}, std::move(var_a_float)));
-  func->set_body(std::move(body));
+  auto func = create<ast::Function>("my_func", std::move(params), &void_type,
+                                    std::move(body));
   auto* func_ptr = func.get();
   mod()->AddFunction(std::move(func));
 
@@ -631,20 +631,20 @@ TEST_F(ValidatorTest, RedeclaredIdentifierDifferentFunctions_Pass) {
       create<ast::FloatLiteral>(&f32, 1.0)));
 
   ast::VariableList params0;
-  auto func0 = create<ast::Function>("func0", std::move(params0), &void_type);
   auto body0 = create<ast::BlockStatement>();
   body0->append(create<ast::VariableDeclStatement>(
       Source{Source::Location{12, 34}}, std::move(var0)));
   body0->append(create<ast::ReturnStatement>());
-  func0->set_body(std::move(body0));
+  auto func0 = create<ast::Function>("func0", std::move(params0), &void_type,
+                                     std::move(body0));
 
   ast::VariableList params1;
-  auto func1 = create<ast::Function>("func1", std::move(params1), &void_type);
   auto body1 = create<ast::BlockStatement>();
   body1->append(create<ast::VariableDeclStatement>(
       Source{Source::Location{13, 34}}, std::move(var1)));
   body1->append(create<ast::ReturnStatement>());
-  func1->set_body(std::move(body1));
+  auto func1 = create<ast::Function>("func1", std::move(params1), &void_type,
+                                     std::move(body1));
   func1->add_decoration(
       create<ast::StageDecoration>(ast::PipelineStage::kVertex, Source{}));
 
diff --git a/src/writer/hlsl/generator_impl_binary_test.cc b/src/writer/hlsl/generator_impl_binary_test.cc
index c201f4742f..32918bdfcd 100644
--- a/src/writer/hlsl/generator_impl_binary_test.cc
+++ b/src/writer/hlsl/generator_impl_binary_test.cc
@@ -326,7 +326,8 @@ TEST_F(HlslGeneratorImplTest_Binary, Call_WithLogical) {
 
   ast::type::VoidType void_type;
 
-  auto func = create<ast::Function>("foo", ast::VariableList{}, &void_type);
+  auto func = create<ast::Function>("foo", ast::VariableList{}, &void_type,
+                                    create<ast::BlockStatement>());
   mod.AddFunction(std::move(func));
 
   ast::ExpressionList params;
diff --git a/src/writer/hlsl/generator_impl_call_test.cc b/src/writer/hlsl/generator_impl_call_test.cc
index 6a50881d8c..afac31a411 100644
--- a/src/writer/hlsl/generator_impl_call_test.cc
+++ b/src/writer/hlsl/generator_impl_call_test.cc
@@ -35,7 +35,8 @@ TEST_F(HlslGeneratorImplTest_Call, EmitExpression_Call_WithoutParams) {
   auto id = create<ast::IdentifierExpression>("my_func");
   ast::CallExpression call(std::move(id), {});
 
-  auto func = create<ast::Function>("my_func", ast::VariableList{}, &void_type);
+  auto func = create<ast::Function>("my_func", ast::VariableList{}, &void_type,
+                                    create<ast::BlockStatement>());
   mod.AddFunction(std::move(func));
 
   ASSERT_TRUE(gen.EmitExpression(pre, out, &call)) << gen.error();
@@ -51,7 +52,8 @@ TEST_F(HlslGeneratorImplTest_Call, EmitExpression_Call_WithParams) {
   params.push_back(create<ast::IdentifierExpression>("param2"));
   ast::CallExpression call(std::move(id), std::move(params));
 
-  auto func = create<ast::Function>("my_func", ast::VariableList{}, &void_type);
+  auto func = create<ast::Function>("my_func", ast::VariableList{}, &void_type,
+                                    create<ast::BlockStatement>());
   mod.AddFunction(std::move(func));
 
   ASSERT_TRUE(gen.EmitExpression(pre, out, &call)) << gen.error();
@@ -68,7 +70,8 @@ TEST_F(HlslGeneratorImplTest_Call, EmitStatement_Call) {
   ast::CallStatement call(
       create<ast::CallExpression>(std::move(id), std::move(params)));
 
-  auto func = create<ast::Function>("my_func", ast::VariableList{}, &void_type);
+  auto func = create<ast::Function>("my_func", ast::VariableList{}, &void_type,
+                                    create<ast::BlockStatement>());
   mod.AddFunction(std::move(func));
   gen.increment_indent();
   ASSERT_TRUE(gen.EmitStatement(out, &call)) << gen.error();
diff --git a/src/writer/hlsl/generator_impl_function_entry_point_data_test.cc b/src/writer/hlsl/generator_impl_function_entry_point_data_test.cc
index 1803fd0527..594442f54a 100644
--- a/src/writer/hlsl/generator_impl_function_entry_point_data_test.cc
+++ b/src/writer/hlsl/generator_impl_function_entry_point_data_test.cc
@@ -72,11 +72,6 @@ TEST_F(HlslGeneratorImplTest_EntryPoint,
   mod.AddGlobalVariable(std::move(bar_var));
 
   ast::VariableList params;
-  auto func = create<ast::Function>("vtx_main", std::move(params), &f32);
-  func->add_decoration(
-      create<ast::StageDecoration>(ast::PipelineStage::kVertex, Source{}));
-  auto* func_ptr = func.get();
-
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::AssignmentStatement>(
       create<ast::IdentifierExpression>("foo"),
@@ -84,7 +79,12 @@ TEST_F(HlslGeneratorImplTest_EntryPoint,
   body->append(create<ast::AssignmentStatement>(
       create<ast::IdentifierExpression>("bar"),
       create<ast::IdentifierExpression>("bar")));
-  func->set_body(std::move(body));
+
+  auto func = create<ast::Function>("vtx_main", std::move(params), &f32,
+                                    std::move(body));
+  func->add_decoration(
+      create<ast::StageDecoration>(ast::PipelineStage::kVertex, Source{}));
+  auto* func_ptr = func.get();
 
   mod.AddFunction(std::move(func));
 
@@ -132,10 +132,6 @@ TEST_F(HlslGeneratorImplTest_EntryPoint,
   mod.AddGlobalVariable(std::move(bar_var));
 
   ast::VariableList params;
-  auto func = create<ast::Function>("vtx_main", std::move(params), &f32);
-  func->add_decoration(
-      create<ast::StageDecoration>(ast::PipelineStage::kVertex, Source{}));
-  auto* func_ptr = func.get();
 
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::AssignmentStatement>(
@@ -144,7 +140,12 @@ TEST_F(HlslGeneratorImplTest_EntryPoint,
   body->append(create<ast::AssignmentStatement>(
       create<ast::IdentifierExpression>("bar"),
       create<ast::IdentifierExpression>("bar")));
-  func->set_body(std::move(body));
+
+  auto func = create<ast::Function>("vtx_main", std::move(params), &f32,
+                                    std::move(body));
+  func->add_decoration(
+      create<ast::StageDecoration>(ast::PipelineStage::kVertex, Source{}));
+  auto* func_ptr = func.get();
 
   mod.AddFunction(std::move(func));
 
@@ -192,10 +193,6 @@ TEST_F(HlslGeneratorImplTest_EntryPoint,
   mod.AddGlobalVariable(std::move(bar_var));
 
   ast::VariableList params;
-  auto func = create<ast::Function>("main", std::move(params), &f32);
-  func->add_decoration(
-      create<ast::StageDecoration>(ast::PipelineStage::kVertex, Source{}));
-  auto* func_ptr = func.get();
 
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::AssignmentStatement>(
@@ -204,7 +201,12 @@ TEST_F(HlslGeneratorImplTest_EntryPoint,
   body->append(create<ast::AssignmentStatement>(
       create<ast::IdentifierExpression>("bar"),
       create<ast::IdentifierExpression>("bar")));
-  func->set_body(std::move(body));
+
+  auto func =
+      create<ast::Function>("main", std::move(params), &f32, std::move(body));
+  func->add_decoration(
+      create<ast::StageDecoration>(ast::PipelineStage::kVertex, Source{}));
+  auto* func_ptr = func.get();
 
   mod.AddFunction(std::move(func));
 
@@ -252,11 +254,6 @@ TEST_F(HlslGeneratorImplTest_EntryPoint,
   mod.AddGlobalVariable(std::move(bar_var));
 
   ast::VariableList params;
-  auto func = create<ast::Function>("main", std::move(params), &f32);
-  func->add_decoration(
-      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
-  auto* func_ptr = func.get();
-
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::AssignmentStatement>(
       create<ast::IdentifierExpression>("foo"),
@@ -264,7 +261,12 @@ TEST_F(HlslGeneratorImplTest_EntryPoint,
   body->append(create<ast::AssignmentStatement>(
       create<ast::IdentifierExpression>("bar"),
       create<ast::IdentifierExpression>("bar")));
-  func->set_body(std::move(body));
+
+  auto func =
+      create<ast::Function>("main", std::move(params), &f32, std::move(body));
+  func->add_decoration(
+      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
+  auto* func_ptr = func.get();
 
   mod.AddFunction(std::move(func));
 
@@ -309,11 +311,6 @@ TEST_F(HlslGeneratorImplTest_EntryPoint,
   mod.AddGlobalVariable(std::move(bar_var));
 
   ast::VariableList params;
-  auto func = create<ast::Function>("main", std::move(params), &f32);
-  func->add_decoration(
-      create<ast::StageDecoration>(ast::PipelineStage::kCompute, Source{}));
-  auto* func_ptr = func.get();
-
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::AssignmentStatement>(
       create<ast::IdentifierExpression>("foo"),
@@ -321,7 +318,12 @@ TEST_F(HlslGeneratorImplTest_EntryPoint,
   body->append(create<ast::AssignmentStatement>(
       create<ast::IdentifierExpression>("bar"),
       create<ast::IdentifierExpression>("bar")));
-  func->set_body(std::move(body));
+
+  auto func =
+      create<ast::Function>("main", std::move(params), &f32, std::move(body));
+  func->add_decoration(
+      create<ast::StageDecoration>(ast::PipelineStage::kCompute, Source{}));
+  auto* func_ptr = func.get();
 
   mod.AddFunction(std::move(func));
 
@@ -361,11 +363,6 @@ TEST_F(HlslGeneratorImplTest_EntryPoint,
   mod.AddGlobalVariable(std::move(bar_var));
 
   ast::VariableList params;
-  auto func = create<ast::Function>("main", std::move(params), &f32);
-  func->add_decoration(
-      create<ast::StageDecoration>(ast::PipelineStage::kCompute, Source{}));
-  auto* func_ptr = func.get();
-
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::AssignmentStatement>(
       create<ast::IdentifierExpression>("foo"),
@@ -373,7 +370,12 @@ TEST_F(HlslGeneratorImplTest_EntryPoint,
   body->append(create<ast::AssignmentStatement>(
       create<ast::IdentifierExpression>("bar"),
       create<ast::IdentifierExpression>("bar")));
-  func->set_body(std::move(body));
+
+  auto func =
+      create<ast::Function>("main", std::move(params), &f32, std::move(body));
+  func->add_decoration(
+      create<ast::StageDecoration>(ast::PipelineStage::kCompute, Source{}));
+  auto* func_ptr = func.get();
 
   mod.AddFunction(std::move(func));
 
@@ -422,18 +424,18 @@ TEST_F(HlslGeneratorImplTest_EntryPoint,
   mod.AddGlobalVariable(std::move(depth_var));
 
   ast::VariableList params;
-  auto func = create<ast::Function>("main", std::move(params), &void_type);
-  func->add_decoration(
-      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
-  auto* func_ptr = func.get();
-
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::AssignmentStatement>(
       create<ast::IdentifierExpression>("depth"),
       create<ast::MemberAccessorExpression>(
           create<ast::IdentifierExpression>("coord"),
           create<ast::IdentifierExpression>("x"))));
-  func->set_body(std::move(body));
+
+  auto func = create<ast::Function>("main", std::move(params), &void_type,
+                                    std::move(body));
+  func->add_decoration(
+      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
+  auto* func_ptr = func.get();
 
   mod.AddFunction(std::move(func));
 
diff --git a/src/writer/hlsl/generator_impl_function_test.cc b/src/writer/hlsl/generator_impl_function_test.cc
index a4e9df3fce..dd0c106b93 100644
--- a/src/writer/hlsl/generator_impl_function_test.cc
+++ b/src/writer/hlsl/generator_impl_function_test.cc
@@ -57,11 +57,10 @@ using HlslGeneratorImplTest_Function = TestHelper;
 TEST_F(HlslGeneratorImplTest_Function, Emit_Function) {
   ast::type::VoidType void_type;
 
-  auto func = create<ast::Function>("my_func", ast::VariableList{}, &void_type);
-
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::ReturnStatement>());
-  func->set_body(std::move(body));
+  auto func = create<ast::Function>("my_func", ast::VariableList{}, &void_type,
+                                    std::move(body));
 
   mod.AddFunction(std::move(func));
   gen.increment_indent();
@@ -77,12 +76,10 @@ TEST_F(HlslGeneratorImplTest_Function, Emit_Function) {
 TEST_F(HlslGeneratorImplTest_Function, Emit_Function_Name_Collision) {
   ast::type::VoidType void_type;
 
-  auto func =
-      create<ast::Function>("GeometryShader", ast::VariableList{}, &void_type);
-
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::ReturnStatement>());
-  func->set_body(std::move(body));
+  auto func = create<ast::Function>("GeometryShader", ast::VariableList{},
+                                    &void_type, std::move(body));
 
   mod.AddFunction(std::move(func));
   gen.increment_indent();
@@ -104,11 +101,11 @@ TEST_F(HlslGeneratorImplTest_Function, Emit_Function_WithParams) {
   params.push_back(create<ast::Variable>("b", ast::StorageClass::kNone, &i32));
 
   ast::type::VoidType void_type;
-  auto func = create<ast::Function>("my_func", std::move(params), &void_type);
 
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::ReturnStatement>());
-  func->set_body(std::move(body));
+  auto func = create<ast::Function>("my_func", std::move(params), &void_type,
+                                    std::move(body));
 
   mod.AddFunction(std::move(func));
   gen.increment_indent();
@@ -145,16 +142,15 @@ TEST_F(HlslGeneratorImplTest_Function,
   mod.AddGlobalVariable(std::move(bar_var));
 
   ast::VariableList params;
-  auto func = create<ast::Function>("frag_main", std::move(params), &void_type);
-  func->add_decoration(
-      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
-
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::AssignmentStatement>(
       create<ast::IdentifierExpression>("bar"),
       create<ast::IdentifierExpression>("foo")));
   body->append(create<ast::ReturnStatement>());
-  func->set_body(std::move(body));
+  auto func = create<ast::Function>("frag_main", std::move(params), &void_type,
+                                    std::move(body));
+  func->add_decoration(
+      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
 
   mod.AddFunction(std::move(func));
 
@@ -204,10 +200,6 @@ TEST_F(HlslGeneratorImplTest_Function,
   mod.AddGlobalVariable(std::move(depth_var));
 
   ast::VariableList params;
-  auto func = create<ast::Function>("frag_main", std::move(params), &void_type);
-  func->add_decoration(
-      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
-
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::AssignmentStatement>(
       create<ast::IdentifierExpression>("depth"),
@@ -215,7 +207,10 @@ TEST_F(HlslGeneratorImplTest_Function,
           create<ast::IdentifierExpression>("coord"),
           create<ast::IdentifierExpression>("x"))));
   body->append(create<ast::ReturnStatement>());
-  func->set_body(std::move(body));
+  auto func = create<ast::Function>("frag_main", std::move(params), &void_type,
+                                    std::move(body));
+  func->add_decoration(
+      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
 
   mod.AddFunction(std::move(func));
 
@@ -256,10 +251,6 @@ TEST_F(HlslGeneratorImplTest_Function,
   mod.AddGlobalVariable(std::move(coord_var));
 
   ast::VariableList params;
-  auto func = create<ast::Function>("frag_main", std::move(params), &void_type);
-  func->add_decoration(
-      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
-
   auto var = create<ast::Variable>("v", ast::StorageClass::kFunction, &f32);
   var->set_constructor(create<ast::MemberAccessorExpression>(
       create<ast::IdentifierExpression>("coord"),
@@ -268,7 +259,10 @@ TEST_F(HlslGeneratorImplTest_Function,
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::VariableDeclStatement>(std::move(var)));
   body->append(create<ast::ReturnStatement>());
-  func->set_body(std::move(body));
+  auto func = create<ast::Function>("frag_main", std::move(params), &void_type,
+                                    std::move(body));
+  func->add_decoration(
+      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
 
   mod.AddFunction(std::move(func));
 
@@ -315,10 +309,6 @@ TEST_F(HlslGeneratorImplTest_Function,
   mod.AddGlobalVariable(std::move(coord_var));
 
   ast::VariableList params;
-  auto func = create<ast::Function>("frag_main", std::move(params), &void_type);
-  func->add_decoration(
-      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
-
   auto var = create<ast::Variable>("v", ast::StorageClass::kFunction, &f32);
   var->set_constructor(create<ast::MemberAccessorExpression>(
       create<ast::MemberAccessorExpression>(
@@ -329,7 +319,10 @@ TEST_F(HlslGeneratorImplTest_Function,
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::VariableDeclStatement>(std::move(var)));
   body->append(create<ast::ReturnStatement>());
-  func->set_body(std::move(body));
+  auto func = create<ast::Function>("frag_main", std::move(params), &void_type,
+                                    std::move(body));
+  func->add_decoration(
+      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
 
   mod.AddFunction(std::move(func));
 
@@ -382,10 +375,6 @@ TEST_F(HlslGeneratorImplTest_Function,
   mod.AddGlobalVariable(std::move(coord_var));
 
   ast::VariableList params;
-  auto func = create<ast::Function>("frag_main", std::move(params), &void_type);
-  func->add_decoration(
-      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
-
   auto var = create<ast::Variable>("v", ast::StorageClass::kFunction, &f32);
   var->set_constructor(create<ast::MemberAccessorExpression>(
       create<ast::IdentifierExpression>("coord"),
@@ -394,7 +383,10 @@ TEST_F(HlslGeneratorImplTest_Function,
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::VariableDeclStatement>(std::move(var)));
   body->append(create<ast::ReturnStatement>());
-  func->set_body(std::move(body));
+  auto func = create<ast::Function>("frag_main", std::move(params), &void_type,
+                                    std::move(body));
+  func->add_decoration(
+      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
 
   mod.AddFunction(std::move(func));
 
@@ -443,10 +435,6 @@ TEST_F(HlslGeneratorImplTest_Function,
   mod.AddGlobalVariable(std::move(coord_var));
 
   ast::VariableList params;
-  auto func = create<ast::Function>("frag_main", std::move(params), &void_type);
-  func->add_decoration(
-      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
-
   auto var = create<ast::Variable>("v", ast::StorageClass::kFunction, &f32);
   var->set_constructor(create<ast::MemberAccessorExpression>(
       create<ast::IdentifierExpression>("coord"),
@@ -455,7 +443,10 @@ TEST_F(HlslGeneratorImplTest_Function,
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::VariableDeclStatement>(std::move(var)));
   body->append(create<ast::ReturnStatement>());
-  func->set_body(std::move(body));
+  auto func = create<ast::Function>("frag_main", std::move(params), &void_type,
+                                    std::move(body));
+  func->add_decoration(
+      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
 
   mod.AddFunction(std::move(func));
 
@@ -505,10 +496,6 @@ TEST_F(HlslGeneratorImplTest_Function,
   mod.AddGlobalVariable(std::move(coord_var));
 
   ast::VariableList params;
-  auto func = create<ast::Function>("frag_main", std::move(params), &void_type);
-  func->add_decoration(
-      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
-
   auto assign = create<ast::AssignmentStatement>(
       create<ast::MemberAccessorExpression>(
           create<ast::IdentifierExpression>("coord"),
@@ -519,7 +506,10 @@ TEST_F(HlslGeneratorImplTest_Function,
   auto body = create<ast::BlockStatement>();
   body->append(std::move(assign));
   body->append(create<ast::ReturnStatement>());
-  func->set_body(std::move(body));
+  auto func = create<ast::Function>("frag_main", std::move(params), &void_type,
+                                    std::move(body));
+  func->add_decoration(
+      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
 
   mod.AddFunction(std::move(func));
 
@@ -569,7 +559,6 @@ TEST_F(
   ast::VariableList params;
   params.push_back(
       create<ast::Variable>("param", ast::StorageClass::kFunction, &f32));
-  auto sub_func = create<ast::Function>("sub_func", std::move(params), &f32);
 
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::AssignmentStatement>(
@@ -580,14 +569,11 @@ TEST_F(
       create<ast::IdentifierExpression>("param")));
   body->append(
       create<ast::ReturnStatement>(create<ast::IdentifierExpression>("foo")));
-  sub_func->set_body(std::move(body));
+  auto sub_func = create<ast::Function>("sub_func", std::move(params), &f32,
+                                        std::move(body));
 
   mod.AddFunction(std::move(sub_func));
 
-  auto func_1 = create<ast::Function>("ep_1", std::move(params), &void_type);
-  func_1->add_decoration(
-      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
-
   ast::ExpressionList expr;
   expr.push_back(create<ast::ScalarConstructorExpression>(
       create<ast::FloatLiteral>(&f32, 1.0f)));
@@ -598,7 +584,10 @@ TEST_F(
       create<ast::CallExpression>(create<ast::IdentifierExpression>("sub_func"),
                                   std::move(expr))));
   body->append(create<ast::ReturnStatement>());
-  func_1->set_body(std::move(body));
+  auto func_1 = create<ast::Function>("ep_1", std::move(params), &void_type,
+                                      std::move(body));
+  func_1->add_decoration(
+      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
 
   mod.AddFunction(std::move(func_1));
 
@@ -649,19 +638,15 @@ TEST_F(HlslGeneratorImplTest_Function,
   ast::VariableList params;
   params.push_back(
       create<ast::Variable>("param", ast::StorageClass::kFunction, &f32));
-  auto sub_func = create<ast::Function>("sub_func", std::move(params), &f32);
 
   auto body = create<ast::BlockStatement>();
   body->append(
       create<ast::ReturnStatement>(create<ast::IdentifierExpression>("param")));
-  sub_func->set_body(std::move(body));
+  auto sub_func = create<ast::Function>("sub_func", std::move(params), &f32,
+                                        std::move(body));
 
   mod.AddFunction(std::move(sub_func));
 
-  auto func_1 = create<ast::Function>("ep_1", std::move(params), &void_type);
-  func_1->add_decoration(
-      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
-
   ast::ExpressionList expr;
   expr.push_back(create<ast::ScalarConstructorExpression>(
       create<ast::FloatLiteral>(&f32, 1.0f)));
@@ -672,7 +657,10 @@ TEST_F(HlslGeneratorImplTest_Function,
       create<ast::CallExpression>(create<ast::IdentifierExpression>("sub_func"),
                                   std::move(expr))));
   body->append(create<ast::ReturnStatement>());
-  func_1->set_body(std::move(body));
+  auto func_1 = create<ast::Function>("ep_1", std::move(params), &void_type,
+                                      std::move(body));
+  func_1->add_decoration(
+      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
 
   mod.AddFunction(std::move(func_1));
 
@@ -725,7 +713,6 @@ TEST_F(
   ast::VariableList params;
   params.push_back(
       create<ast::Variable>("param", ast::StorageClass::kFunction, &f32));
-  auto sub_func = create<ast::Function>("sub_func", std::move(params), &f32);
 
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::AssignmentStatement>(
@@ -735,14 +722,11 @@ TEST_F(
           create<ast::IdentifierExpression>("x"))));
   body->append(
       create<ast::ReturnStatement>(create<ast::IdentifierExpression>("param")));
-  sub_func->set_body(std::move(body));
+  auto sub_func = create<ast::Function>("sub_func", std::move(params), &f32,
+                                        std::move(body));
 
   mod.AddFunction(std::move(sub_func));
 
-  auto func_1 = create<ast::Function>("ep_1", std::move(params), &void_type);
-  func_1->add_decoration(
-      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
-
   ast::ExpressionList expr;
   expr.push_back(create<ast::ScalarConstructorExpression>(
       create<ast::FloatLiteral>(&f32, 1.0f)));
@@ -753,7 +737,10 @@ TEST_F(
       create<ast::CallExpression>(create<ast::IdentifierExpression>("sub_func"),
                                   std::move(expr))));
   body->append(create<ast::ReturnStatement>());
-  func_1->set_body(std::move(body));
+  auto func_1 = create<ast::Function>("ep_1", std::move(params), &void_type,
+                                      std::move(body));
+  func_1->add_decoration(
+      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
 
   mod.AddFunction(std::move(func_1));
 
@@ -802,21 +789,17 @@ TEST_F(HlslGeneratorImplTest_Function,
   ast::VariableList params;
   params.push_back(
       create<ast::Variable>("param", ast::StorageClass::kFunction, &f32));
-  auto sub_func = create<ast::Function>("sub_func", std::move(params), &f32);
 
   auto body = create<ast::BlockStatement>();
   body->append(
       create<ast::ReturnStatement>(create<ast::MemberAccessorExpression>(
           create<ast::IdentifierExpression>("coord"),
           create<ast::IdentifierExpression>("x"))));
-  sub_func->set_body(std::move(body));
+  auto sub_func = create<ast::Function>("sub_func", std::move(params), &f32,
+                                        std::move(body));
 
   mod.AddFunction(std::move(sub_func));
 
-  auto func = create<ast::Function>("frag_main", std::move(params), &void_type);
-  func->add_decoration(
-      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
-
   ast::ExpressionList expr;
   expr.push_back(create<ast::ScalarConstructorExpression>(
       create<ast::FloatLiteral>(&f32, 1.0f)));
@@ -828,7 +811,10 @@ TEST_F(HlslGeneratorImplTest_Function,
   body = create<ast::BlockStatement>();
   body->append(create<ast::VariableDeclStatement>(std::move(var)));
   body->append(create<ast::ReturnStatement>());
-  func->set_body(std::move(body));
+  auto func = create<ast::Function>("frag_main", std::move(params), &void_type,
+                                    std::move(body));
+  func->add_decoration(
+      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
 
   mod.AddFunction(std::move(func));
 
@@ -871,21 +857,17 @@ TEST_F(HlslGeneratorImplTest_Function,
   ast::VariableList params;
   params.push_back(
       create<ast::Variable>("param", ast::StorageClass::kFunction, &f32));
-  auto sub_func = create<ast::Function>("sub_func", std::move(params), &f32);
 
   auto body = create<ast::BlockStatement>();
   body->append(
       create<ast::ReturnStatement>(create<ast::MemberAccessorExpression>(
           create<ast::IdentifierExpression>("coord"),
           create<ast::IdentifierExpression>("x"))));
-  sub_func->set_body(std::move(body));
+  auto sub_func = create<ast::Function>("sub_func", std::move(params), &f32,
+                                        std::move(body));
 
   mod.AddFunction(std::move(sub_func));
 
-  auto func = create<ast::Function>("frag_main", std::move(params), &void_type);
-  func->add_decoration(
-      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
-
   ast::ExpressionList expr;
   expr.push_back(create<ast::ScalarConstructorExpression>(
       create<ast::FloatLiteral>(&f32, 1.0f)));
@@ -897,7 +879,10 @@ TEST_F(HlslGeneratorImplTest_Function,
   body = create<ast::BlockStatement>();
   body->append(create<ast::VariableDeclStatement>(std::move(var)));
   body->append(create<ast::ReturnStatement>());
-  func->set_body(std::move(body));
+  auto func = create<ast::Function>("frag_main", std::move(params), &void_type,
+                                    std::move(body));
+  func->add_decoration(
+      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
 
   mod.AddFunction(std::move(func));
 
@@ -933,10 +918,6 @@ TEST_F(HlslGeneratorImplTest_Function,
   mod.AddGlobalVariable(std::move(bar_var));
 
   ast::VariableList params;
-  auto func_1 = create<ast::Function>("ep_1", std::move(params), &void_type);
-  func_1->add_decoration(
-      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
-
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::AssignmentStatement>(
       create<ast::IdentifierExpression>("bar"),
@@ -955,7 +936,10 @@ TEST_F(HlslGeneratorImplTest_Function,
       std::move(list)));
 
   body->append(create<ast::ReturnStatement>());
-  func_1->set_body(std::move(body));
+  auto func_1 = create<ast::Function>("ep_1", std::move(params), &void_type,
+                                      std::move(body));
+  func_1->add_decoration(
+      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
 
   mod.AddFunction(std::move(func_1));
 
@@ -981,8 +965,8 @@ TEST_F(HlslGeneratorImplTest_Function,
        Emit_FunctionDecoration_EntryPoint_WithNameCollision) {
   ast::type::VoidType void_type;
 
-  auto func =
-      create<ast::Function>("GeometryShader", ast::VariableList{}, &void_type);
+  auto func = create<ast::Function>("GeometryShader", ast::VariableList{},
+                                    &void_type, create<ast::BlockStatement>());
   func->add_decoration(
       create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
 
@@ -1000,13 +984,12 @@ TEST_F(HlslGeneratorImplTest_Function,
   ast::type::VoidType void_type;
 
   ast::VariableList params;
-  auto func = create<ast::Function>("main", std::move(params), &void_type);
-  func->add_decoration(
-      create<ast::StageDecoration>(ast::PipelineStage::kCompute, Source{}));
-
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::ReturnStatement>());
-  func->set_body(std::move(body));
+  auto func = create<ast::Function>("main", std::move(params), &void_type,
+                                    std::move(body));
+  func->add_decoration(
+      create<ast::StageDecoration>(ast::PipelineStage::kCompute, Source{}));
 
   mod.AddFunction(std::move(func));
 
@@ -1025,15 +1008,14 @@ TEST_F(HlslGeneratorImplTest_Function,
   ast::type::VoidType void_type;
 
   ast::VariableList params;
-  auto func = create<ast::Function>("main", std::move(params), &void_type);
+  auto body = create<ast::BlockStatement>();
+  body->append(create<ast::ReturnStatement>());
+  auto func = create<ast::Function>("main", std::move(params), &void_type,
+                                    std::move(body));
   func->add_decoration(
       create<ast::StageDecoration>(ast::PipelineStage::kCompute, Source{}));
   func->add_decoration(create<ast::WorkgroupDecoration>(2u, 4u, 6u, Source{}));
 
-  auto body = create<ast::BlockStatement>();
-  body->append(create<ast::ReturnStatement>());
-  func->set_body(std::move(body));
-
   mod.AddFunction(std::move(func));
 
   ASSERT_TRUE(td.Determine()) << td.error();
@@ -1054,11 +1036,11 @@ TEST_F(HlslGeneratorImplTest_Function, Emit_Function_WithArrayParams) {
   params.push_back(create<ast::Variable>("a", ast::StorageClass::kNone, &ary));
 
   ast::type::VoidType void_type;
-  auto func = create<ast::Function>("my_func", std::move(params), &void_type);
 
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::ReturnStatement>());
-  func->set_body(std::move(body));
+  auto func = create<ast::Function>("my_func", std::move(params), &void_type,
+                                    std::move(body));
 
   mod.AddFunction(std::move(func));
   gen.increment_indent();
@@ -1119,10 +1101,6 @@ TEST_F(HlslGeneratorImplTest_Function,
 
   {
     ast::VariableList params;
-    auto func = create<ast::Function>("a", std::move(params), &void_type);
-    func->add_decoration(
-        create<ast::StageDecoration>(ast::PipelineStage::kCompute, Source{}));
-
     auto var = create<ast::Variable>("v", ast::StorageClass::kFunction, &f32);
     var->set_constructor(create<ast::MemberAccessorExpression>(
         create<ast::IdentifierExpression>("data"),
@@ -1131,17 +1109,16 @@ TEST_F(HlslGeneratorImplTest_Function,
     auto body = create<ast::BlockStatement>();
     body->append(create<ast::VariableDeclStatement>(std::move(var)));
     body->append(create<ast::ReturnStatement>());
-    func->set_body(std::move(body));
+    auto func = create<ast::Function>("a", std::move(params), &void_type,
+                                      std::move(body));
+    func->add_decoration(
+        create<ast::StageDecoration>(ast::PipelineStage::kCompute, Source{}));
 
     mod.AddFunction(std::move(func));
   }
 
   {
     ast::VariableList params;
-    auto func = create<ast::Function>("b", std::move(params), &void_type);
-    func->add_decoration(
-        create<ast::StageDecoration>(ast::PipelineStage::kCompute, Source{}));
-
     auto var = create<ast::Variable>("v", ast::StorageClass::kFunction, &f32);
     var->set_constructor(create<ast::MemberAccessorExpression>(
         create<ast::IdentifierExpression>("data"),
@@ -1150,7 +1127,10 @@ TEST_F(HlslGeneratorImplTest_Function,
     auto body = create<ast::BlockStatement>();
     body->append(create<ast::VariableDeclStatement>(std::move(var)));
     body->append(create<ast::ReturnStatement>());
-    func->set_body(std::move(body));
+    auto func = create<ast::Function>("b", std::move(params), &void_type,
+                                      std::move(body));
+    func->add_decoration(
+        create<ast::StageDecoration>(ast::PipelineStage::kCompute, Source{}));
 
     mod.AddFunction(std::move(func));
   }
diff --git a/src/writer/hlsl/generator_impl_test.cc b/src/writer/hlsl/generator_impl_test.cc
index a92a5ed031..9e097728df 100644
--- a/src/writer/hlsl/generator_impl_test.cc
+++ b/src/writer/hlsl/generator_impl_test.cc
@@ -29,7 +29,8 @@ using HlslGeneratorImplTest = TestHelper;
 
 TEST_F(HlslGeneratorImplTest, Generate) {
   ast::type::VoidType void_type;
-  auto func = create<ast::Function>("my_func", ast::VariableList{}, &void_type);
+  auto func = create<ast::Function>("my_func", ast::VariableList{}, &void_type,
+                                    create<ast::BlockStatement>());
   mod.AddFunction(std::move(func));
 
   ASSERT_TRUE(gen.Generate(out)) << gen.error();
diff --git a/src/writer/msl/generator_impl_call_test.cc b/src/writer/msl/generator_impl_call_test.cc
index e01f87bf20..6a6b07cc47 100644
--- a/src/writer/msl/generator_impl_call_test.cc
+++ b/src/writer/msl/generator_impl_call_test.cc
@@ -37,7 +37,8 @@ TEST_F(MslGeneratorImplTest, EmitExpression_Call_WithoutParams) {
   auto id = create<ast::IdentifierExpression>("my_func");
   ast::CallExpression call(std::move(id), {});
 
-  auto func = create<ast::Function>("my_func", ast::VariableList{}, &void_type);
+  auto func = create<ast::Function>("my_func", ast::VariableList{}, &void_type,
+                                    create<ast::BlockStatement>());
   mod.AddFunction(std::move(func));
 
   ASSERT_TRUE(gen.EmitExpression(&call)) << gen.error();
@@ -53,7 +54,8 @@ TEST_F(MslGeneratorImplTest, EmitExpression_Call_WithParams) {
   params.push_back(create<ast::IdentifierExpression>("param2"));
   ast::CallExpression call(std::move(id), std::move(params));
 
-  auto func = create<ast::Function>("my_func", ast::VariableList{}, &void_type);
+  auto func = create<ast::Function>("my_func", ast::VariableList{}, &void_type,
+                                    create<ast::BlockStatement>());
   mod.AddFunction(std::move(func));
 
   ASSERT_TRUE(gen.EmitExpression(&call)) << gen.error();
@@ -70,7 +72,8 @@ TEST_F(MslGeneratorImplTest, EmitStatement_Call) {
   ast::CallStatement call(
       create<ast::CallExpression>(std::move(id), std::move(params)));
 
-  auto func = create<ast::Function>("my_func", ast::VariableList{}, &void_type);
+  auto func = create<ast::Function>("my_func", ast::VariableList{}, &void_type,
+                                    create<ast::BlockStatement>());
   mod.AddFunction(std::move(func));
 
   gen.increment_indent();
diff --git a/src/writer/msl/generator_impl_function_entry_point_data_test.cc b/src/writer/msl/generator_impl_function_entry_point_data_test.cc
index 4a0a400c46..3bf391fe56 100644
--- a/src/writer/msl/generator_impl_function_entry_point_data_test.cc
+++ b/src/writer/msl/generator_impl_function_entry_point_data_test.cc
@@ -71,11 +71,6 @@ TEST_F(MslGeneratorImplTest, Emit_Function_EntryPointData_Vertex_Input) {
   mod.AddGlobalVariable(std::move(bar_var));
 
   ast::VariableList params;
-  auto func = create<ast::Function>("vtx_main", std::move(params), &f32);
-  func->add_decoration(
-      create<ast::StageDecoration>(ast::PipelineStage::kVertex, Source{}));
-  auto* func_ptr = func.get();
-
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::AssignmentStatement>(
       create<ast::IdentifierExpression>("foo"),
@@ -83,7 +78,12 @@ TEST_F(MslGeneratorImplTest, Emit_Function_EntryPointData_Vertex_Input) {
   body->append(create<ast::AssignmentStatement>(
       create<ast::IdentifierExpression>("bar"),
       create<ast::IdentifierExpression>("bar")));
-  func->set_body(std::move(body));
+
+  auto func = create<ast::Function>("vtx_main", std::move(params), &f32,
+                                    std::move(body));
+  func->add_decoration(
+      create<ast::StageDecoration>(ast::PipelineStage::kVertex, Source{}));
+  auto* func_ptr = func.get();
 
   mod.AddFunction(std::move(func));
 
@@ -129,11 +129,6 @@ TEST_F(MslGeneratorImplTest, Emit_Function_EntryPointData_Vertex_Output) {
   mod.AddGlobalVariable(std::move(bar_var));
 
   ast::VariableList params;
-  auto func = create<ast::Function>("vtx_main", std::move(params), &f32);
-  func->add_decoration(
-      create<ast::StageDecoration>(ast::PipelineStage::kVertex, Source{}));
-  auto* func_ptr = func.get();
-
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::AssignmentStatement>(
       create<ast::IdentifierExpression>("foo"),
@@ -141,7 +136,12 @@ TEST_F(MslGeneratorImplTest, Emit_Function_EntryPointData_Vertex_Output) {
   body->append(create<ast::AssignmentStatement>(
       create<ast::IdentifierExpression>("bar"),
       create<ast::IdentifierExpression>("bar")));
-  func->set_body(std::move(body));
+
+  auto func = create<ast::Function>("vtx_main", std::move(params), &f32,
+                                    std::move(body));
+  func->add_decoration(
+      create<ast::StageDecoration>(ast::PipelineStage::kVertex, Source{}));
+  auto* func_ptr = func.get();
 
   mod.AddFunction(std::move(func));
 
@@ -187,11 +187,6 @@ TEST_F(MslGeneratorImplTest, Emit_Function_EntryPointData_Fragment_Input) {
   mod.AddGlobalVariable(std::move(bar_var));
 
   ast::VariableList params;
-  auto func = create<ast::Function>("main", std::move(params), &f32);
-  func->add_decoration(
-      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
-  auto* func_ptr = func.get();
-
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::AssignmentStatement>(
       create<ast::IdentifierExpression>("foo"),
@@ -199,7 +194,11 @@ TEST_F(MslGeneratorImplTest, Emit_Function_EntryPointData_Fragment_Input) {
   body->append(create<ast::AssignmentStatement>(
       create<ast::IdentifierExpression>("bar"),
       create<ast::IdentifierExpression>("bar")));
-  func->set_body(std::move(body));
+  auto func =
+      create<ast::Function>("main", std::move(params), &f32, std::move(body));
+  func->add_decoration(
+      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
+  auto* func_ptr = func.get();
 
   mod.AddFunction(std::move(func));
 
@@ -245,11 +244,6 @@ TEST_F(MslGeneratorImplTest, Emit_Function_EntryPointData_Fragment_Output) {
   mod.AddGlobalVariable(std::move(bar_var));
 
   ast::VariableList params;
-  auto func = create<ast::Function>("main", std::move(params), &f32);
-  func->add_decoration(
-      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
-  auto* func_ptr = func.get();
-
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::AssignmentStatement>(
       create<ast::IdentifierExpression>("foo"),
@@ -257,7 +251,12 @@ TEST_F(MslGeneratorImplTest, Emit_Function_EntryPointData_Fragment_Output) {
   body->append(create<ast::AssignmentStatement>(
       create<ast::IdentifierExpression>("bar"),
       create<ast::IdentifierExpression>("bar")));
-  func->set_body(std::move(body));
+
+  auto func =
+      create<ast::Function>("main", std::move(params), &f32, std::move(body));
+  func->add_decoration(
+      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
+  auto* func_ptr = func.get();
 
   mod.AddFunction(std::move(func));
 
@@ -300,11 +299,6 @@ TEST_F(MslGeneratorImplTest, Emit_Function_EntryPointData_Compute_Input) {
   mod.AddGlobalVariable(std::move(bar_var));
 
   ast::VariableList params;
-  auto func = create<ast::Function>("main", std::move(params), &f32);
-  func->add_decoration(
-      create<ast::StageDecoration>(ast::PipelineStage::kCompute, Source{}));
-  auto* func_ptr = func.get();
-
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::AssignmentStatement>(
       create<ast::IdentifierExpression>("foo"),
@@ -312,7 +306,12 @@ TEST_F(MslGeneratorImplTest, Emit_Function_EntryPointData_Compute_Input) {
   body->append(create<ast::AssignmentStatement>(
       create<ast::IdentifierExpression>("bar"),
       create<ast::IdentifierExpression>("bar")));
-  func->set_body(std::move(body));
+
+  auto func =
+      create<ast::Function>("main", std::move(params), &f32, std::move(body));
+  func->add_decoration(
+      create<ast::StageDecoration>(ast::PipelineStage::kCompute, Source{}));
+  auto* func_ptr = func.get();
 
   mod.AddFunction(std::move(func));
 
@@ -350,11 +349,6 @@ TEST_F(MslGeneratorImplTest, Emit_Function_EntryPointData_Compute_Output) {
   mod.AddGlobalVariable(std::move(bar_var));
 
   ast::VariableList params;
-  auto func = create<ast::Function>("main", std::move(params), &f32);
-  func->add_decoration(
-      create<ast::StageDecoration>(ast::PipelineStage::kCompute, Source{}));
-  auto* func_ptr = func.get();
-
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::AssignmentStatement>(
       create<ast::IdentifierExpression>("foo"),
@@ -362,7 +356,12 @@ TEST_F(MslGeneratorImplTest, Emit_Function_EntryPointData_Compute_Output) {
   body->append(create<ast::AssignmentStatement>(
       create<ast::IdentifierExpression>("bar"),
       create<ast::IdentifierExpression>("bar")));
-  func->set_body(std::move(body));
+
+  auto func =
+      create<ast::Function>("main", std::move(params), &f32, std::move(body));
+  func->add_decoration(
+      create<ast::StageDecoration>(ast::PipelineStage::kCompute, Source{}));
+  auto* func_ptr = func.get();
 
   mod.AddFunction(std::move(func));
 
@@ -408,10 +407,6 @@ TEST_F(MslGeneratorImplTest, Emit_Function_EntryPointData_Builtins) {
   mod.AddGlobalVariable(std::move(depth_var));
 
   ast::VariableList params;
-  auto func = create<ast::Function>("main", std::move(params), &void_type);
-  func->add_decoration(
-      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
-  auto* func_ptr = func.get();
 
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::AssignmentStatement>(
@@ -419,7 +414,12 @@ TEST_F(MslGeneratorImplTest, Emit_Function_EntryPointData_Builtins) {
       create<ast::MemberAccessorExpression>(
           create<ast::IdentifierExpression>("coord"),
           create<ast::IdentifierExpression>("x"))));
-  func->set_body(std::move(body));
+
+  auto func = create<ast::Function>("main", std::move(params), &void_type,
+                                    std::move(body));
+  func->add_decoration(
+      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
+  auto* func_ptr = func.get();
 
   mod.AddFunction(std::move(func));
 
diff --git a/src/writer/msl/generator_impl_function_test.cc b/src/writer/msl/generator_impl_function_test.cc
index 093e9529c6..6f13cf9d45 100644
--- a/src/writer/msl/generator_impl_function_test.cc
+++ b/src/writer/msl/generator_impl_function_test.cc
@@ -60,11 +60,10 @@ using MslGeneratorImplTest = TestHelper;
 TEST_F(MslGeneratorImplTest, Emit_Function) {
   ast::type::VoidType void_type;
 
-  auto func = create<ast::Function>("my_func", ast::VariableList{}, &void_type);
-
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::ReturnStatement>());
-  func->set_body(std::move(body));
+  auto func = create<ast::Function>("my_func", ast::VariableList{}, &void_type,
+                                    std::move(body));
 
   mod.AddFunction(std::move(func));
   gen.increment_indent();
@@ -82,11 +81,10 @@ TEST_F(MslGeneratorImplTest, Emit_Function) {
 TEST_F(MslGeneratorImplTest, Emit_Function_Name_Collision) {
   ast::type::VoidType void_type;
 
-  auto func = create<ast::Function>("main", ast::VariableList{}, &void_type);
-
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::ReturnStatement>());
-  func->set_body(std::move(body));
+  auto func = create<ast::Function>("main", ast::VariableList{}, &void_type,
+                                    std::move(body));
 
   mod.AddFunction(std::move(func));
   gen.increment_indent();
@@ -110,11 +108,11 @@ TEST_F(MslGeneratorImplTest, Emit_Function_WithParams) {
   params.push_back(create<ast::Variable>("b", ast::StorageClass::kNone, &i32));
 
   ast::type::VoidType void_type;
-  auto func = create<ast::Function>("my_func", std::move(params), &void_type);
 
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::ReturnStatement>());
-  func->set_body(std::move(body));
+  auto func = create<ast::Function>("my_func", std::move(params), &void_type,
+                                    std::move(body));
 
   mod.AddFunction(std::move(func));
   gen.increment_indent();
@@ -152,16 +150,16 @@ TEST_F(MslGeneratorImplTest, Emit_FunctionDecoration_EntryPoint_WithInOutVars) {
   mod.AddGlobalVariable(std::move(bar_var));
 
   ast::VariableList params;
-  auto func = create<ast::Function>("frag_main", std::move(params), &void_type);
-  func->add_decoration(
-      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
-
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::AssignmentStatement>(
       create<ast::IdentifierExpression>("bar"),
       create<ast::IdentifierExpression>("foo")));
   body->append(create<ast::ReturnStatement>());
-  func->set_body(std::move(body));
+
+  auto func = create<ast::Function>("frag_main", std::move(params), &void_type,
+                                    std::move(body));
+  func->add_decoration(
+      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
 
   mod.AddFunction(std::move(func));
 
@@ -214,10 +212,6 @@ TEST_F(MslGeneratorImplTest,
   mod.AddGlobalVariable(std::move(depth_var));
 
   ast::VariableList params;
-  auto func = create<ast::Function>("frag_main", std::move(params), &void_type);
-  func->add_decoration(
-      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
-
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::AssignmentStatement>(
       create<ast::IdentifierExpression>("depth"),
@@ -225,7 +219,11 @@ TEST_F(MslGeneratorImplTest,
           create<ast::IdentifierExpression>("coord"),
           create<ast::IdentifierExpression>("x"))));
   body->append(create<ast::ReturnStatement>());
-  func->set_body(std::move(body));
+
+  auto func = create<ast::Function>("frag_main", std::move(params), &void_type,
+                                    std::move(body));
+  func->add_decoration(
+      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
 
   mod.AddFunction(std::move(func));
 
@@ -265,10 +263,6 @@ TEST_F(MslGeneratorImplTest, Emit_FunctionDecoration_EntryPoint_With_Uniform) {
   mod.AddGlobalVariable(std::move(coord_var));
 
   ast::VariableList params;
-  auto func = create<ast::Function>("frag_main", std::move(params), &void_type);
-  func->add_decoration(
-      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
-
   auto var = create<ast::Variable>("v", ast::StorageClass::kFunction, &f32);
   var->set_constructor(create<ast::MemberAccessorExpression>(
       create<ast::IdentifierExpression>("coord"),
@@ -277,7 +271,11 @@ TEST_F(MslGeneratorImplTest, Emit_FunctionDecoration_EntryPoint_With_Uniform) {
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::VariableDeclStatement>(std::move(var)));
   body->append(create<ast::ReturnStatement>());
-  func->set_body(std::move(body));
+
+  auto func = create<ast::Function>("frag_main", std::move(params), &void_type,
+                                    std::move(body));
+  func->add_decoration(
+      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
 
   mod.AddFunction(std::move(func));
 
@@ -330,10 +328,6 @@ TEST_F(MslGeneratorImplTest,
   mod.AddGlobalVariable(std::move(coord_var));
 
   ast::VariableList params;
-  auto func = create<ast::Function>("frag_main", std::move(params), &void_type);
-  func->add_decoration(
-      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
-
   auto var = create<ast::Variable>("v", ast::StorageClass::kFunction, &f32);
   var->set_constructor(create<ast::MemberAccessorExpression>(
       create<ast::IdentifierExpression>("coord"),
@@ -342,7 +336,11 @@ TEST_F(MslGeneratorImplTest,
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::VariableDeclStatement>(std::move(var)));
   body->append(create<ast::ReturnStatement>());
-  func->set_body(std::move(body));
+
+  auto func = create<ast::Function>("frag_main", std::move(params), &void_type,
+                                    std::move(body));
+  func->add_decoration(
+      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
 
   mod.AddFunction(std::move(func));
 
@@ -399,9 +397,6 @@ TEST_F(MslGeneratorImplTest,
   mod.AddGlobalVariable(std::move(coord_var));
 
   ast::VariableList params;
-  auto func = create<ast::Function>("frag_main", std::move(params), &void_type);
-  func->add_decoration(
-      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
 
   auto var = create<ast::Variable>("v", ast::StorageClass::kFunction, &f32);
   var->set_constructor(create<ast::MemberAccessorExpression>(
@@ -411,7 +406,11 @@ TEST_F(MslGeneratorImplTest,
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::VariableDeclStatement>(std::move(var)));
   body->append(create<ast::ReturnStatement>());
-  func->set_body(std::move(body));
+
+  auto func = create<ast::Function>("frag_main", std::move(params), &void_type,
+                                    std::move(body));
+  func->add_decoration(
+      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
 
   mod.AddFunction(std::move(func));
 
@@ -467,7 +466,6 @@ TEST_F(
   ast::VariableList params;
   params.push_back(
       create<ast::Variable>("param", ast::StorageClass::kFunction, &f32));
-  auto sub_func = create<ast::Function>("sub_func", std::move(params), &f32);
 
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::AssignmentStatement>(
@@ -478,14 +476,11 @@ TEST_F(
       create<ast::IdentifierExpression>("param")));
   body->append(
       create<ast::ReturnStatement>(create<ast::IdentifierExpression>("foo")));
-  sub_func->set_body(std::move(body));
+  auto sub_func = create<ast::Function>("sub_func", std::move(params), &f32,
+                                        std::move(body));
 
   mod.AddFunction(std::move(sub_func));
 
-  auto func_1 = create<ast::Function>("ep_1", std::move(params), &void_type);
-  func_1->add_decoration(
-      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
-
   ast::ExpressionList expr;
   expr.push_back(create<ast::ScalarConstructorExpression>(
       create<ast::FloatLiteral>(&f32, 1.0f)));
@@ -496,7 +491,10 @@ TEST_F(
       create<ast::CallExpression>(create<ast::IdentifierExpression>("sub_func"),
                                   std::move(expr))));
   body->append(create<ast::ReturnStatement>());
-  func_1->set_body(std::move(body));
+  auto func_1 = create<ast::Function>("ep_1", std::move(params), &void_type,
+                                      std::move(body));
+  func_1->add_decoration(
+      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
 
   mod.AddFunction(std::move(func_1));
 
@@ -550,19 +548,15 @@ TEST_F(MslGeneratorImplTest,
   ast::VariableList params;
   params.push_back(
       create<ast::Variable>("param", ast::StorageClass::kFunction, &f32));
-  auto sub_func = create<ast::Function>("sub_func", std::move(params), &f32);
 
   auto body = create<ast::BlockStatement>();
   body->append(
       create<ast::ReturnStatement>(create<ast::IdentifierExpression>("param")));
-  sub_func->set_body(std::move(body));
+  auto sub_func = create<ast::Function>("sub_func", std::move(params), &f32,
+                                        std::move(body));
 
   mod.AddFunction(std::move(sub_func));
 
-  auto func_1 = create<ast::Function>("ep_1", std::move(params), &void_type);
-  func_1->add_decoration(
-      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
-
   ast::ExpressionList expr;
   expr.push_back(create<ast::ScalarConstructorExpression>(
       create<ast::FloatLiteral>(&f32, 1.0f)));
@@ -573,7 +567,11 @@ TEST_F(MslGeneratorImplTest,
       create<ast::CallExpression>(create<ast::IdentifierExpression>("sub_func"),
                                   std::move(expr))));
   body->append(create<ast::ReturnStatement>());
-  func_1->set_body(std::move(body));
+
+  auto func_1 = create<ast::Function>("ep_1", std::move(params), &void_type,
+                                      std::move(body));
+  func_1->add_decoration(
+      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
 
   mod.AddFunction(std::move(func_1));
 
@@ -629,7 +627,6 @@ TEST_F(
   ast::VariableList params;
   params.push_back(
       create<ast::Variable>("param", ast::StorageClass::kFunction, &f32));
-  auto sub_func = create<ast::Function>("sub_func", std::move(params), &f32);
 
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::AssignmentStatement>(
@@ -639,14 +636,11 @@ TEST_F(
           create<ast::IdentifierExpression>("x"))));
   body->append(
       create<ast::ReturnStatement>(create<ast::IdentifierExpression>("param")));
-  sub_func->set_body(std::move(body));
+  auto sub_func = create<ast::Function>("sub_func", std::move(params), &f32,
+                                        std::move(body));
 
   mod.AddFunction(std::move(sub_func));
 
-  auto func_1 = create<ast::Function>("ep_1", std::move(params), &void_type);
-  func_1->add_decoration(
-      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
-
   ast::ExpressionList expr;
   expr.push_back(create<ast::ScalarConstructorExpression>(
       create<ast::FloatLiteral>(&f32, 1.0f)));
@@ -657,7 +651,10 @@ TEST_F(
       create<ast::CallExpression>(create<ast::IdentifierExpression>("sub_func"),
                                   std::move(expr))));
   body->append(create<ast::ReturnStatement>());
-  func_1->set_body(std::move(body));
+  auto func_1 = create<ast::Function>("ep_1", std::move(params), &void_type,
+                                      std::move(body));
+  func_1->add_decoration(
+      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
 
   mod.AddFunction(std::move(func_1));
 
@@ -704,21 +701,17 @@ TEST_F(MslGeneratorImplTest,
   ast::VariableList params;
   params.push_back(
       create<ast::Variable>("param", ast::StorageClass::kFunction, &f32));
-  auto sub_func = create<ast::Function>("sub_func", std::move(params), &f32);
 
   auto body = create<ast::BlockStatement>();
   body->append(
       create<ast::ReturnStatement>(create<ast::MemberAccessorExpression>(
           create<ast::IdentifierExpression>("coord"),
           create<ast::IdentifierExpression>("x"))));
-  sub_func->set_body(std::move(body));
+  auto sub_func = create<ast::Function>("sub_func", std::move(params), &f32,
+                                        std::move(body));
 
   mod.AddFunction(std::move(sub_func));
 
-  auto func = create<ast::Function>("frag_main", std::move(params), &void_type);
-  func->add_decoration(
-      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
-
   ast::ExpressionList expr;
   expr.push_back(create<ast::ScalarConstructorExpression>(
       create<ast::FloatLiteral>(&f32, 1.0f)));
@@ -730,7 +723,11 @@ TEST_F(MslGeneratorImplTest,
   body = create<ast::BlockStatement>();
   body->append(create<ast::VariableDeclStatement>(std::move(var)));
   body->append(create<ast::ReturnStatement>());
-  func->set_body(std::move(body));
+
+  auto func = create<ast::Function>("frag_main", std::move(params), &void_type,
+                                    std::move(body));
+  func->add_decoration(
+      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
 
   mod.AddFunction(std::move(func));
 
@@ -787,21 +784,17 @@ TEST_F(MslGeneratorImplTest,
   ast::VariableList params;
   params.push_back(
       create<ast::Variable>("param", ast::StorageClass::kFunction, &f32));
-  auto sub_func = create<ast::Function>("sub_func", std::move(params), &f32);
 
   auto body = create<ast::BlockStatement>();
   body->append(
       create<ast::ReturnStatement>(create<ast::MemberAccessorExpression>(
           create<ast::IdentifierExpression>("coord"),
           create<ast::IdentifierExpression>("b"))));
-  sub_func->set_body(std::move(body));
+  auto sub_func = create<ast::Function>("sub_func", std::move(params), &f32,
+                                        std::move(body));
 
   mod.AddFunction(std::move(sub_func));
 
-  auto func = create<ast::Function>("frag_main", std::move(params), &void_type);
-  func->add_decoration(
-      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
-
   ast::ExpressionList expr;
   expr.push_back(create<ast::ScalarConstructorExpression>(
       create<ast::FloatLiteral>(&f32, 1.0f)));
@@ -813,7 +806,11 @@ TEST_F(MslGeneratorImplTest,
   body = create<ast::BlockStatement>();
   body->append(create<ast::VariableDeclStatement>(std::move(var)));
   body->append(create<ast::ReturnStatement>());
-  func->set_body(std::move(body));
+
+  auto func = create<ast::Function>("frag_main", std::move(params), &void_type,
+                                    std::move(body));
+  func->add_decoration(
+      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
 
   mod.AddFunction(std::move(func));
 
@@ -876,21 +873,17 @@ TEST_F(MslGeneratorImplTest,
   ast::VariableList params;
   params.push_back(
       create<ast::Variable>("param", ast::StorageClass::kFunction, &f32));
-  auto sub_func = create<ast::Function>("sub_func", std::move(params), &f32);
 
   auto body = create<ast::BlockStatement>();
   body->append(
       create<ast::ReturnStatement>(create<ast::MemberAccessorExpression>(
           create<ast::IdentifierExpression>("coord"),
           create<ast::IdentifierExpression>("b"))));
-  sub_func->set_body(std::move(body));
+  auto sub_func = create<ast::Function>("sub_func", std::move(params), &f32,
+                                        std::move(body));
 
   mod.AddFunction(std::move(sub_func));
 
-  auto func = create<ast::Function>("frag_main", std::move(params), &void_type);
-  func->add_decoration(
-      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
-
   ast::ExpressionList expr;
   expr.push_back(create<ast::ScalarConstructorExpression>(
       create<ast::FloatLiteral>(&f32, 1.0f)));
@@ -902,7 +895,11 @@ TEST_F(MslGeneratorImplTest,
   body = create<ast::BlockStatement>();
   body->append(create<ast::VariableDeclStatement>(std::move(var)));
   body->append(create<ast::ReturnStatement>());
-  func->set_body(std::move(body));
+
+  auto func = create<ast::Function>("frag_main", std::move(params), &void_type,
+                                    std::move(body));
+  func->add_decoration(
+      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
 
   mod.AddFunction(std::move(func));
 
@@ -944,10 +941,6 @@ TEST_F(MslGeneratorImplTest,
   mod.AddGlobalVariable(std::move(bar_var));
 
   ast::VariableList params;
-  auto func_1 = create<ast::Function>("ep_1", std::move(params), &void_type);
-  func_1->add_decoration(
-      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
-
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::AssignmentStatement>(
       create<ast::IdentifierExpression>("bar"),
@@ -966,7 +959,11 @@ TEST_F(MslGeneratorImplTest,
       std::move(list)));
 
   body->append(create<ast::ReturnStatement>());
-  func_1->set_body(std::move(body));
+
+  auto func_1 = create<ast::Function>("ep_1", std::move(params), &void_type,
+                                      std::move(body));
+  func_1->add_decoration(
+      create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
 
   mod.AddFunction(std::move(func_1));
 
@@ -994,7 +991,8 @@ TEST_F(MslGeneratorImplTest,
        Emit_FunctionDecoration_EntryPoint_WithNameCollision) {
   ast::type::VoidType void_type;
 
-  auto func = create<ast::Function>("main", ast::VariableList{}, &void_type);
+  auto func = create<ast::Function>("main", ast::VariableList{}, &void_type,
+                                    create<ast::BlockStatement>());
   func->add_decoration(
       create<ast::StageDecoration>(ast::PipelineStage::kCompute, Source{}));
 
@@ -1017,11 +1015,11 @@ TEST_F(MslGeneratorImplTest, Emit_Function_WithArrayParams) {
   params.push_back(create<ast::Variable>("a", ast::StorageClass::kNone, &ary));
 
   ast::type::VoidType void_type;
-  auto func = create<ast::Function>("my_func", std::move(params), &void_type);
 
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::ReturnStatement>());
-  func->set_body(std::move(body));
+  auto func = create<ast::Function>("my_func", std::move(params), &void_type,
+                                    std::move(body));
 
   mod.AddFunction(std::move(func));
 
@@ -1086,10 +1084,6 @@ TEST_F(MslGeneratorImplTest,
 
   {
     ast::VariableList params;
-    auto func = create<ast::Function>("a", std::move(params), &void_type);
-    func->add_decoration(
-        create<ast::StageDecoration>(ast::PipelineStage::kCompute, Source{}));
-
     auto var = create<ast::Variable>("v", ast::StorageClass::kFunction, &f32);
     var->set_constructor(create<ast::MemberAccessorExpression>(
         create<ast::IdentifierExpression>("data"),
@@ -1098,17 +1092,17 @@ TEST_F(MslGeneratorImplTest,
     auto body = create<ast::BlockStatement>();
     body->append(create<ast::VariableDeclStatement>(std::move(var)));
     body->append(create<ast::ReturnStatement>());
-    func->set_body(std::move(body));
+
+    auto func = create<ast::Function>("a", std::move(params), &void_type,
+                                      std::move(body));
+    func->add_decoration(
+        create<ast::StageDecoration>(ast::PipelineStage::kCompute, Source{}));
 
     mod.AddFunction(std::move(func));
   }
 
   {
     ast::VariableList params;
-    auto func = create<ast::Function>("b", std::move(params), &void_type);
-    func->add_decoration(
-        create<ast::StageDecoration>(ast::PipelineStage::kCompute, Source{}));
-
     auto var = create<ast::Variable>("v", ast::StorageClass::kFunction, &f32);
     var->set_constructor(create<ast::MemberAccessorExpression>(
         create<ast::IdentifierExpression>("data"),
@@ -1117,7 +1111,11 @@ TEST_F(MslGeneratorImplTest,
     auto body = create<ast::BlockStatement>();
     body->append(create<ast::VariableDeclStatement>(std::move(var)));
     body->append(create<ast::ReturnStatement>());
-    func->set_body(std::move(body));
+
+    auto func = create<ast::Function>("b", std::move(params), &void_type,
+                                      std::move(body));
+    func->add_decoration(
+        create<ast::StageDecoration>(ast::PipelineStage::kCompute, Source{}));
 
     mod.AddFunction(std::move(func));
   }
diff --git a/src/writer/msl/generator_impl_test.cc b/src/writer/msl/generator_impl_test.cc
index 7c70ede211..1da11441bb 100644
--- a/src/writer/msl/generator_impl_test.cc
+++ b/src/writer/msl/generator_impl_test.cc
@@ -50,7 +50,8 @@ using MslGeneratorImplTest = TestHelper;
 TEST_F(MslGeneratorImplTest, Generate) {
   ast::type::VoidType void_type;
 
-  auto func = create<ast::Function>("my_func", ast::VariableList{}, &void_type);
+  auto func = create<ast::Function>("my_func", ast::VariableList{}, &void_type,
+                                    create<ast::BlockStatement>());
   func->add_decoration(
       create<ast::StageDecoration>(ast::PipelineStage::kCompute, Source{}));
   mod.AddFunction(std::move(func));
diff --git a/src/writer/spirv/builder_call_test.cc b/src/writer/spirv/builder_call_test.cc
index 5bcadc5778..a7afb760bb 100644
--- a/src/writer/spirv/builder_call_test.cc
+++ b/src/writer/spirv/builder_call_test.cc
@@ -49,15 +49,13 @@ TEST_F(BuilderTest, Expression_Call) {
   func_params.push_back(
       create<ast::Variable>("b", ast::StorageClass::kFunction, &f32));
 
-  ast::Function a_func("a_func", std::move(func_params), &f32);
-
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::ReturnStatement>(create<ast::BinaryExpression>(
       ast::BinaryOp::kAdd, create<ast::IdentifierExpression>("a"),
       create<ast::IdentifierExpression>("b"))));
-  a_func.set_body(std::move(body));
+  ast::Function a_func("a_func", std::move(func_params), &f32, std::move(body));
 
-  ast::Function func("main", {}, &void_type);
+  ast::Function func("main", {}, &void_type, create<ast::BlockStatement>());
 
   ast::ExpressionList call_params;
   call_params.push_back(create<ast::ScalarConstructorExpression>(
@@ -111,15 +109,15 @@ TEST_F(BuilderTest, Statement_Call) {
   func_params.push_back(
       create<ast::Variable>("b", ast::StorageClass::kFunction, &f32));
 
-  ast::Function a_func("a_func", std::move(func_params), &void_type);
-
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::ReturnStatement>(create<ast::BinaryExpression>(
       ast::BinaryOp::kAdd, create<ast::IdentifierExpression>("a"),
       create<ast::IdentifierExpression>("b"))));
-  a_func.set_body(std::move(body));
 
-  ast::Function func("main", {}, &void_type);
+  ast::Function a_func("a_func", std::move(func_params), &void_type,
+                       std::move(body));
+
+  ast::Function func("main", {}, &void_type, create<ast::BlockStatement>());
 
   ast::ExpressionList call_params;
   call_params.push_back(create<ast::ScalarConstructorExpression>(
diff --git a/src/writer/spirv/builder_function_decoration_test.cc b/src/writer/spirv/builder_function_decoration_test.cc
index 4a22a79a78..740c5d7f41 100644
--- a/src/writer/spirv/builder_function_decoration_test.cc
+++ b/src/writer/spirv/builder_function_decoration_test.cc
@@ -42,7 +42,7 @@ using BuilderTest = TestHelper;
 TEST_F(BuilderTest, FunctionDecoration_Stage) {
   ast::type::VoidType void_type;
 
-  ast::Function func("main", {}, &void_type);
+  ast::Function func("main", {}, &void_type, create<ast::BlockStatement>());
   func.add_decoration(
       create<ast::StageDecoration>(ast::PipelineStage::kVertex, Source{}));
 
@@ -66,7 +66,7 @@ TEST_P(FunctionDecoration_StageTest, Emit) {
 
   ast::type::VoidType void_type;
 
-  ast::Function func("main", {}, &void_type);
+  ast::Function func("main", {}, &void_type, create<ast::BlockStatement>());
   func.add_decoration(create<ast::StageDecoration>(params.stage, Source{}));
 
   ASSERT_TRUE(b.GenerateFunction(&func)) << b.error();
@@ -92,7 +92,7 @@ TEST_F(BuilderTest, FunctionDecoration_Stage_WithUnusedInterfaceIds) {
   ast::type::F32Type f32;
   ast::type::VoidType void_type;
 
-  ast::Function func("main", {}, &void_type);
+  ast::Function func("main", {}, &void_type, create<ast::BlockStatement>());
   func.add_decoration(
       create<ast::StageDecoration>(ast::PipelineStage::kVertex, Source{}));
   auto v_in = create<ast::Variable>("my_in", ast::StorageClass::kInput, &f32);
@@ -135,7 +135,7 @@ TEST_F(BuilderTest, FunctionDecoration_Stage_WithUsedInterfaceIds) {
   ast::type::F32Type f32;
   ast::type::VoidType void_type;
 
-  ast::Function func("main", {}, &void_type);
+  ast::Function func("main", {}, &void_type, create<ast::BlockStatement>());
   func.add_decoration(
       create<ast::StageDecoration>(ast::PipelineStage::kVertex, Source{}));
 
@@ -197,7 +197,7 @@ OpName %11 "tint_6d61696e"
 TEST_F(BuilderTest, FunctionDecoration_ExecutionMode_Fragment_OriginUpperLeft) {
   ast::type::VoidType void_type;
 
-  ast::Function func("main", {}, &void_type);
+  ast::Function func("main", {}, &void_type, create<ast::BlockStatement>());
   func.add_decoration(
       create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
 
@@ -210,7 +210,7 @@ TEST_F(BuilderTest, FunctionDecoration_ExecutionMode_Fragment_OriginUpperLeft) {
 TEST_F(BuilderTest, FunctionDecoration_WorkgroupSize_Default) {
   ast::type::VoidType void_type;
 
-  ast::Function func("main", {}, &void_type);
+  ast::Function func("main", {}, &void_type, create<ast::BlockStatement>());
   func.add_decoration(
       create<ast::StageDecoration>(ast::PipelineStage::kCompute, Source{}));
 
@@ -223,7 +223,7 @@ TEST_F(BuilderTest, FunctionDecoration_WorkgroupSize_Default) {
 TEST_F(BuilderTest, FunctionDecoration_WorkgroupSize) {
   ast::type::VoidType void_type;
 
-  ast::Function func("main", {}, &void_type);
+  ast::Function func("main", {}, &void_type, create<ast::BlockStatement>());
   func.add_decoration(create<ast::WorkgroupDecoration>(2u, 4u, 6u, Source{}));
   func.add_decoration(
       create<ast::StageDecoration>(ast::PipelineStage::kCompute, Source{}));
@@ -237,11 +237,11 @@ TEST_F(BuilderTest, FunctionDecoration_WorkgroupSize) {
 TEST_F(BuilderTest, FunctionDecoration_ExecutionMode_MultipleFragment) {
   ast::type::VoidType void_type;
 
-  ast::Function func1("main1", {}, &void_type);
+  ast::Function func1("main1", {}, &void_type, create<ast::BlockStatement>());
   func1.add_decoration(
       create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
 
-  ast::Function func2("main2", {}, &void_type);
+  ast::Function func2("main2", {}, &void_type, create<ast::BlockStatement>());
   func2.add_decoration(
       create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
 
diff --git a/src/writer/spirv/builder_function_test.cc b/src/writer/spirv/builder_function_test.cc
index e7549f7167..4855fefe96 100644
--- a/src/writer/spirv/builder_function_test.cc
+++ b/src/writer/spirv/builder_function_test.cc
@@ -48,7 +48,7 @@ using BuilderTest = TestHelper;
 
 TEST_F(BuilderTest, Function_Empty) {
   ast::type::VoidType void_type;
-  ast::Function func("a_func", {}, &void_type);
+  ast::Function func("a_func", {}, &void_type, create<ast::BlockStatement>());
 
   ASSERT_TRUE(b.GenerateFunction(&func));
   EXPECT_EQ(DumpInstructions(b.debug()), R"(OpName %3 "tint_615f66756e63"
@@ -76,12 +76,10 @@ TEST_F(BuilderTest, Function_WithParams) {
   var_b->set_is_const(true);
   params.push_back(std::move(var_b));
 
-  ast::Function func("a_func", std::move(params), &f32);
-
   auto body = create<ast::BlockStatement>();
   body->append(
       create<ast::ReturnStatement>(create<ast::IdentifierExpression>("a")));
-  func.set_body(std::move(body));
+  ast::Function func("a_func", std::move(params), &f32, std::move(body));
 
   td.RegisterVariableForTesting(func.params()[0].get());
   td.RegisterVariableForTesting(func.params()[1].get());
@@ -109,8 +107,7 @@ TEST_F(BuilderTest, Function_WithBody) {
   auto body = create<ast::BlockStatement>();
   body->append(create<ast::ReturnStatement>());
 
-  ast::Function func("a_func", {}, &void_type);
-  func.set_body(std::move(body));
+  ast::Function func("a_func", {}, &void_type, std::move(body));
 
   ASSERT_TRUE(b.GenerateFunction(&func));
   EXPECT_EQ(DumpBuilder(b), R"(OpName %3 "tint_615f66756e63"
@@ -125,7 +122,7 @@ OpFunctionEnd
 
 TEST_F(BuilderTest, FunctionType) {
   ast::type::VoidType void_type;
-  ast::Function func("a_func", {}, &void_type);
+  ast::Function func("a_func", {}, &void_type, create<ast::BlockStatement>());
 
   ASSERT_TRUE(b.GenerateFunction(&func));
   EXPECT_EQ(DumpInstructions(b.types()), R"(%2 = OpTypeVoid
@@ -135,8 +132,8 @@ TEST_F(BuilderTest, FunctionType) {
 
 TEST_F(BuilderTest, FunctionType_DeDuplicate) {
   ast::type::VoidType void_type;
-  ast::Function func1("a_func", {}, &void_type);
-  ast::Function func2("b_func", {}, &void_type);
+  ast::Function func1("a_func", {}, &void_type, create<ast::BlockStatement>());
+  ast::Function func2("b_func", {}, &void_type, create<ast::BlockStatement>());
 
   ASSERT_TRUE(b.GenerateFunction(&func1));
   ASSERT_TRUE(b.GenerateFunction(&func2));
@@ -193,10 +190,6 @@ TEST_F(BuilderTest, Emit_Multiple_EntryPoint_With_Same_ModuleVar) {
 
   {
     ast::VariableList params;
-    auto func = create<ast::Function>("a", std::move(params), &void_type);
-    func->add_decoration(
-        create<ast::StageDecoration>(ast::PipelineStage::kCompute, Source{}));
-
     auto var = create<ast::Variable>("v", ast::StorageClass::kFunction, &f32);
     var->set_constructor(create<ast::MemberAccessorExpression>(
         create<ast::IdentifierExpression>("data"),
@@ -205,17 +198,17 @@ TEST_F(BuilderTest, Emit_Multiple_EntryPoint_With_Same_ModuleVar) {
     auto body = create<ast::BlockStatement>();
     body->append(create<ast::VariableDeclStatement>(std::move(var)));
     body->append(create<ast::ReturnStatement>());
-    func->set_body(std::move(body));
+
+    auto func = create<ast::Function>("a", std::move(params), &void_type,
+                                      std::move(body));
+    func->add_decoration(
+        create<ast::StageDecoration>(ast::PipelineStage::kCompute, Source{}));
 
     mod.AddFunction(std::move(func));
   }
 
   {
     ast::VariableList params;
-    auto func = create<ast::Function>("b", std::move(params), &void_type);
-    func->add_decoration(
-        create<ast::StageDecoration>(ast::PipelineStage::kCompute, Source{}));
-
     auto var = create<ast::Variable>("v", ast::StorageClass::kFunction, &f32);
     var->set_constructor(create<ast::MemberAccessorExpression>(
         create<ast::IdentifierExpression>("data"),
@@ -224,7 +217,11 @@ TEST_F(BuilderTest, Emit_Multiple_EntryPoint_With_Same_ModuleVar) {
     auto body = create<ast::BlockStatement>();
     body->append(create<ast::VariableDeclStatement>(std::move(var)));
     body->append(create<ast::ReturnStatement>());
-    func->set_body(std::move(body));
+
+    auto func = create<ast::Function>("b", std::move(params), &void_type,
+                                      std::move(body));
+    func->add_decoration(
+        create<ast::StageDecoration>(ast::PipelineStage::kCompute, Source{}));
 
     mod.AddFunction(std::move(func));
   }
diff --git a/src/writer/spirv/builder_intrinsic_test.cc b/src/writer/spirv/builder_intrinsic_test.cc
index 033d83a421..88b9b23315 100644
--- a/src/writer/spirv/builder_intrinsic_test.cc
+++ b/src/writer/spirv/builder_intrinsic_test.cc
@@ -940,7 +940,7 @@ TEST_F(IntrinsicBuilderTest, Call_GLSLMethod_WithLoad) {
 
   ASSERT_TRUE(td.DetermineResultType(&expr)) << td.error();
 
-  ast::Function func("a_func", {}, void_type());
+  ast::Function func("a_func", {}, void_type(), create<ast::BlockStatement>());
 
   ASSERT_TRUE(b.GenerateGlobalVariable(var.get())) << b.error();
   ASSERT_TRUE(b.GenerateFunction(&func)) << b.error();
@@ -971,7 +971,7 @@ TEST_P(Intrinsic_Builtin_SingleParam_Float_Test, Call_Scalar) {
   auto expr = call_expr(param.name, 1.0f);
   ASSERT_TRUE(td.DetermineResultType(&expr)) << td.error();
 
-  ast::Function func("a_func", {}, void_type());
+  ast::Function func("a_func", {}, void_type(), create<ast::BlockStatement>());
 
   ASSERT_TRUE(b.GenerateFunction(&func)) << b.error();
 
@@ -996,7 +996,7 @@ TEST_P(Intrinsic_Builtin_SingleParam_Float_Test, Call_Vector) {
   auto expr = call_expr(param.name, construct(vec(f32(), 2), 1.0f, 1.0f));
   ASSERT_TRUE(td.DetermineResultType(&expr)) << td.error();
 
-  ast::Function func("a_func", {}, void_type());
+  ast::Function func("a_func", {}, void_type(), create<ast::BlockStatement>());
 
   ASSERT_TRUE(b.GenerateFunction(&func)) << b.error();
 
@@ -1047,7 +1047,7 @@ TEST_F(IntrinsicBuilderTest, Call_Length_Scalar) {
 
   ASSERT_TRUE(td.DetermineResultType(&expr)) << td.error();
 
-  ast::Function func("a_func", {}, void_type());
+  ast::Function func("a_func", {}, void_type(), create<ast::BlockStatement>());
 
   ASSERT_TRUE(b.GenerateFunction(&func)) << b.error();
 
@@ -1069,7 +1069,7 @@ TEST_F(IntrinsicBuilderTest, Call_Length_Vector) {
   auto expr = call_expr("length", construct(vec(f32(), 2), 1.0f, 1.0f));
   ASSERT_TRUE(td.DetermineResultType(&expr)) << td.error();
 
-  ast::Function func("a_func", {}, void_type());
+  ast::Function func("a_func", {}, void_type(), create<ast::BlockStatement>());
 
   ASSERT_TRUE(b.GenerateFunction(&func)) << b.error();
 
@@ -1093,7 +1093,7 @@ TEST_F(IntrinsicBuilderTest, Call_Normalize) {
   auto expr = call_expr("normalize", construct(vec(f32(), 2), 1.0f, 1.0f));
   ASSERT_TRUE(td.DetermineResultType(&expr)) << td.error();
 
-  ast::Function func("a_func", {}, void_type());
+  ast::Function func("a_func", {}, void_type(), create<ast::BlockStatement>());
 
   ASSERT_TRUE(b.GenerateFunction(&func)) << b.error();
 
@@ -1121,7 +1121,7 @@ TEST_P(Intrinsic_Builtin_DualParam_Float_Test, Call_Scalar) {
   auto expr = call_expr(param.name, 1.0f, 1.0f);
   ASSERT_TRUE(td.DetermineResultType(&expr)) << td.error();
 
-  ast::Function func("a_func", {}, void_type());
+  ast::Function func("a_func", {}, void_type(), create<ast::BlockStatement>());
 
   ASSERT_TRUE(b.GenerateFunction(&func)) << b.error();
 
@@ -1149,7 +1149,7 @@ TEST_P(Intrinsic_Builtin_DualParam_Float_Test, Call_Vector) {
 
   ASSERT_TRUE(td.DetermineResultType(&expr)) << td.error();
 
-  ast::Function func("a_func", {}, void_type());
+  ast::Function func("a_func", {}, void_type(), create<ast::BlockStatement>());
 
   ASSERT_TRUE(b.GenerateFunction(&func)) << b.error();
 
@@ -1182,7 +1182,7 @@ TEST_F(IntrinsicBuilderTest, Call_Distance_Scalar) {
   auto expr = call_expr("distance", 1.0f, 1.0f);
   ASSERT_TRUE(td.DetermineResultType(&expr)) << td.error();
 
-  ast::Function func("a_func", {}, void_type());
+  ast::Function func("a_func", {}, void_type(), create<ast::BlockStatement>());
 
   ASSERT_TRUE(b.GenerateFunction(&func)) << b.error();
 
@@ -1207,7 +1207,7 @@ TEST_F(IntrinsicBuilderTest, Call_Distance_Vector) {
 
   ASSERT_TRUE(td.DetermineResultType(&expr)) << td.error();
 
-  ast::Function func("a_func", {}, void_type());
+  ast::Function func("a_func", {}, void_type(), create<ast::BlockStatement>());
 
   ASSERT_TRUE(b.GenerateFunction(&func)) << b.error();
 
@@ -1234,7 +1234,7 @@ TEST_F(IntrinsicBuilderTest, Call_Cross) {
 
   ASSERT_TRUE(td.DetermineResultType(&expr)) << td.error();
 
-  ast::Function func("a_func", {}, void_type());
+  ast::Function func("a_func", {}, void_type(), create<ast::BlockStatement>());
 
   ASSERT_TRUE(b.GenerateFunction(&func)) << b.error();
 
@@ -1262,7 +1262,7 @@ TEST_P(Intrinsic_Builtin_ThreeParam_Float_Test, Call_Scalar) {
   auto expr = call_expr(param.name, 1.0f, 1.0f, 1.0f);
   ASSERT_TRUE(td.DetermineResultType(&expr)) << td.error();
 
-  ast::Function func("a_func", {}, void_type());
+  ast::Function func("a_func", {}, void_type(), create<ast::BlockStatement>());
 
   ASSERT_TRUE(b.GenerateFunction(&func)) << b.error();
 
@@ -1291,7 +1291,7 @@ TEST_P(Intrinsic_Builtin_ThreeParam_Float_Test, Call_Vector) {
 
   ASSERT_TRUE(td.DetermineResultType(&expr)) << td.error();
 
-  ast::Function func("a_func", {}, void_type());
+  ast::Function func("a_func", {}, void_type(), create<ast::BlockStatement>());
 
   ASSERT_TRUE(b.GenerateFunction(&func)) << b.error();
 
@@ -1329,7 +1329,7 @@ TEST_P(Intrinsic_Builtin_SingleParam_Sint_Test, Call_Scalar) {
   auto expr = call_expr(param.name, 1);
   ASSERT_TRUE(td.DetermineResultType(&expr)) << td.error();
 
-  ast::Function func("a_func", {}, void_type());
+  ast::Function func("a_func", {}, void_type(), create<ast::BlockStatement>());
 
   ASSERT_TRUE(b.GenerateFunction(&func)) << b.error();
 
@@ -1354,7 +1354,7 @@ TEST_P(Intrinsic_Builtin_SingleParam_Sint_Test, Call_Vector) {
   auto expr = call_expr(param.name, construct(vec(i32(), 2), 1, 1));
   ASSERT_TRUE(td.DetermineResultType(&expr)) << td.error();
 
-  ast::Function func("a_func", {}, void_type());
+  ast::Function func("a_func", {}, void_type(), create<ast::BlockStatement>());
 
   ASSERT_TRUE(b.GenerateFunction(&func)) << b.error();
 
@@ -1386,7 +1386,7 @@ TEST_P(Intrinsic_Builtin_SingleParam_Uint_Test, Call_Scalar) {
   auto expr = call_expr(param.name, 1u);
   ASSERT_TRUE(td.DetermineResultType(&expr)) << td.error();
 
-  ast::Function func("a_func", {}, void_type());
+  ast::Function func("a_func", {}, void_type(), create<ast::BlockStatement>());
 
   ASSERT_TRUE(b.GenerateFunction(&func)) << b.error();
 
@@ -1411,7 +1411,7 @@ TEST_P(Intrinsic_Builtin_SingleParam_Uint_Test, Call_Vector) {
   auto expr = call_expr(param.name, construct(vec(u32(), 2), 1u, 1u));
   ASSERT_TRUE(td.DetermineResultType(&expr)) << td.error();
 
-  ast::Function func("a_func", {}, void_type());
+  ast::Function func("a_func", {}, void_type(), create<ast::BlockStatement>());
 
   ASSERT_TRUE(b.GenerateFunction(&func)) << b.error();
 
@@ -1443,7 +1443,7 @@ TEST_P(Intrinsic_Builtin_DualParam_SInt_Test, Call_Scalar) {
   auto expr = call_expr(param.name, 1, 1);
   ASSERT_TRUE(td.DetermineResultType(&expr)) << td.error();
 
-  ast::Function func("a_func", {}, void_type());
+  ast::Function func("a_func", {}, void_type(), create<ast::BlockStatement>());
 
   ASSERT_TRUE(b.GenerateFunction(&func)) << b.error();
 
@@ -1470,7 +1470,7 @@ TEST_P(Intrinsic_Builtin_DualParam_SInt_Test, Call_Vector) {
       call_expr(param.name, construct(vec2, 1, 1), construct(vec2, 1, 1));
   ASSERT_TRUE(td.DetermineResultType(&expr)) << td.error();
 
-  ast::Function func("a_func", {}, void_type());
+  ast::Function func("a_func", {}, void_type(), create<ast::BlockStatement>());
 
   ASSERT_TRUE(b.GenerateFunction(&func)) << b.error();
 
@@ -1503,7 +1503,7 @@ TEST_P(Intrinsic_Builtin_DualParam_UInt_Test, Call_Scalar) {
   auto expr = call_expr(param.name, 1u, 1u);
   ASSERT_TRUE(td.DetermineResultType(&expr)) << td.error();
 
-  ast::Function func("a_func", {}, void_type());
+  ast::Function func("a_func", {}, void_type(), create<ast::BlockStatement>());
 
   ASSERT_TRUE(b.GenerateFunction(&func)) << b.error();
 
@@ -1531,7 +1531,7 @@ TEST_P(Intrinsic_Builtin_DualParam_UInt_Test, Call_Vector) {
 
   ASSERT_TRUE(td.DetermineResultType(&expr)) << td.error();
 
-  ast::Function func("a_func", {}, void_type());
+  ast::Function func("a_func", {}, void_type(), create<ast::BlockStatement>());
 
   ASSERT_TRUE(b.GenerateFunction(&func)) << b.error();
 
@@ -1564,7 +1564,7 @@ TEST_P(Intrinsic_Builtin_ThreeParam_Sint_Test, Call_Scalar) {
   auto expr = call_expr(param.name, 1, 1, 1);
   ASSERT_TRUE(td.DetermineResultType(&expr)) << td.error();
 
-  ast::Function func("a_func", {}, void_type());
+  ast::Function func("a_func", {}, void_type(), create<ast::BlockStatement>());
 
   ASSERT_TRUE(b.GenerateFunction(&func)) << b.error();
 
@@ -1592,7 +1592,7 @@ TEST_P(Intrinsic_Builtin_ThreeParam_Sint_Test, Call_Vector) {
 
   ASSERT_TRUE(td.DetermineResultType(&expr)) << td.error();
 
-  ast::Function func("a_func", {}, void_type());
+  ast::Function func("a_func", {}, void_type(), create<ast::BlockStatement>());
 
   ASSERT_TRUE(b.GenerateFunction(&func)) << b.error();
 
@@ -1624,7 +1624,7 @@ TEST_P(Intrinsic_Builtin_ThreeParam_Uint_Test, Call_Scalar) {
   auto expr = call_expr(param.name, 1u, 1u, 1u);
   ASSERT_TRUE(td.DetermineResultType(&expr)) << td.error();
 
-  ast::Function func("a_func", {}, void_type());
+  ast::Function func("a_func", {}, void_type(), create<ast::BlockStatement>());
 
   ASSERT_TRUE(b.GenerateFunction(&func)) << b.error();
 
@@ -1652,7 +1652,7 @@ TEST_P(Intrinsic_Builtin_ThreeParam_Uint_Test, Call_Vector) {
 
   ASSERT_TRUE(td.DetermineResultType(&expr)) << td.error();
 
-  ast::Function func("a_func", {}, void_type());
+  ast::Function func("a_func", {}, void_type(), create<ast::BlockStatement>());
 
   ASSERT_TRUE(b.GenerateFunction(&func)) << b.error();
 
@@ -1682,7 +1682,7 @@ TEST_F(IntrinsicBuilderTest, Call_Determinant) {
 
   ASSERT_TRUE(td.DetermineResultType(&expr)) << td.error();
 
-  ast::Function func("a_func", {}, void_type());
+  ast::Function func("a_func", {}, void_type(), create<ast::BlockStatement>());
 
   ASSERT_TRUE(b.GenerateFunction(&func)) << b.error();
 
@@ -1725,7 +1725,7 @@ TEST_F(IntrinsicBuilderTest, Call_ArrayLength) {
 
   ASSERT_TRUE(td.DetermineResultType(&expr)) << td.error();
 
-  ast::Function func("a_func", {}, void_type());
+  ast::Function func("a_func", {}, void_type(), create<ast::BlockStatement>());
 
   ASSERT_TRUE(b.GenerateFunction(&func)) << b.error();
   ASSERT_TRUE(b.GenerateGlobalVariable(var.get())) << b.error();
@@ -1765,7 +1765,7 @@ TEST_F(IntrinsicBuilderTest, Call_ArrayLength_OtherMembersInStruct) {
 
   ASSERT_TRUE(td.DetermineResultType(&expr)) << td.error();
 
-  ast::Function func("a_func", {}, void_type());
+  ast::Function func("a_func", {}, void_type(), create<ast::BlockStatement>());
 
   ASSERT_TRUE(b.GenerateFunction(&func)) << b.error();
   ASSERT_TRUE(b.GenerateGlobalVariable(var.get())) << b.error();
@@ -1810,7 +1810,7 @@ TEST_F(IntrinsicBuilderTest, DISABLED_Call_ArrayLength_Ptr) {
   auto expr = call_expr("arrayLength", "ptr_var");
   ASSERT_TRUE(td.DetermineResultType(&expr)) << td.error();
 
-  ast::Function func("a_func", {}, void_type());
+  ast::Function func("a_func", {}, void_type(), create<ast::BlockStatement>());
 
   ASSERT_TRUE(b.GenerateFunction(&func)) << b.error();
   ASSERT_TRUE(b.GenerateGlobalVariable(var.get())) << b.error();
diff --git a/src/writer/spirv/builder_switch_test.cc b/src/writer/spirv/builder_switch_test.cc
index ed4b35ecd3..639af98a69 100644
--- a/src/writer/spirv/builder_switch_test.cc
+++ b/src/writer/spirv/builder_switch_test.cc
@@ -111,7 +111,7 @@ TEST_F(BuilderTest, Switch_WithCase) {
   td.RegisterVariableForTesting(a.get());
   ASSERT_TRUE(td.DetermineResultType(&expr)) << td.error();
 
-  ast::Function func("a_func", {}, &i32);
+  ast::Function func("a_func", {}, &i32, create<ast::BlockStatement>());
 
   ASSERT_TRUE(b.GenerateGlobalVariable(v.get())) << b.error();
   ASSERT_TRUE(b.GenerateGlobalVariable(a.get())) << b.error();
@@ -175,7 +175,7 @@ TEST_F(BuilderTest, Switch_WithDefault) {
   td.RegisterVariableForTesting(a.get());
   ASSERT_TRUE(td.DetermineResultType(&expr)) << td.error();
 
-  ast::Function func("a_func", {}, &i32);
+  ast::Function func("a_func", {}, &i32, create<ast::BlockStatement>());
 
   ASSERT_TRUE(b.GenerateGlobalVariable(v.get())) << b.error();
   ASSERT_TRUE(b.GenerateGlobalVariable(a.get())) << b.error();
@@ -260,7 +260,7 @@ TEST_F(BuilderTest, Switch_WithCaseAndDefault) {
   td.RegisterVariableForTesting(a.get());
   ASSERT_TRUE(td.DetermineResultType(&expr)) << td.error();
 
-  ast::Function func("a_func", {}, &i32);
+  ast::Function func("a_func", {}, &i32, create<ast::BlockStatement>());
 
   ASSERT_TRUE(b.GenerateGlobalVariable(v.get())) << b.error();
   ASSERT_TRUE(b.GenerateGlobalVariable(a.get())) << b.error();
@@ -354,7 +354,7 @@ TEST_F(BuilderTest, Switch_CaseWithFallthrough) {
   td.RegisterVariableForTesting(a.get());
   ASSERT_TRUE(td.DetermineResultType(&expr)) << td.error();
 
-  ast::Function func("a_func", {}, &i32);
+  ast::Function func("a_func", {}, &i32, create<ast::BlockStatement>());
 
   ASSERT_TRUE(b.GenerateGlobalVariable(v.get())) << b.error();
   ASSERT_TRUE(b.GenerateGlobalVariable(a.get())) << b.error();
@@ -426,7 +426,7 @@ TEST_F(BuilderTest, Switch_CaseFallthroughLastStatement) {
   td.RegisterVariableForTesting(a.get());
   ASSERT_TRUE(td.DetermineResultType(&expr)) << td.error();
 
-  ast::Function func("a_func", {}, &i32);
+  ast::Function func("a_func", {}, &i32, create<ast::BlockStatement>());
 
   ASSERT_TRUE(b.GenerateGlobalVariable(v.get())) << b.error();
   ASSERT_TRUE(b.GenerateGlobalVariable(a.get())) << b.error();
@@ -479,7 +479,7 @@ TEST_F(BuilderTest, Switch_WithNestedBreak) {
   td.RegisterVariableForTesting(a.get());
   ASSERT_TRUE(td.DetermineResultType(&expr)) << td.error();
 
-  ast::Function func("a_func", {}, &i32);
+  ast::Function func("a_func", {}, &i32, create<ast::BlockStatement>());
 
   ASSERT_TRUE(b.GenerateGlobalVariable(v.get())) << b.error();
   ASSERT_TRUE(b.GenerateGlobalVariable(a.get())) << b.error();
diff --git a/src/writer/wgsl/generator_impl_function_test.cc b/src/writer/wgsl/generator_impl_function_test.cc
index 85dac8e34f..c7f0b306b9 100644
--- a/src/writer/wgsl/generator_impl_function_test.cc
+++ b/src/writer/wgsl/generator_impl_function_test.cc
@@ -48,8 +48,7 @@ TEST_F(WgslGeneratorImplTest, Emit_Function) {
   body->append(create<ast::ReturnStatement>());
 
   ast::type::VoidType void_type;
-  ast::Function func("my_func", {}, &void_type);
-  func.set_body(std::move(body));
+  ast::Function func("my_func", {}, &void_type, std::move(body));
 
   gen.increment_indent();
 
@@ -73,8 +72,7 @@ TEST_F(WgslGeneratorImplTest, Emit_Function_WithParams) {
   params.push_back(create<ast::Variable>("b", ast::StorageClass::kNone, &i32));
 
   ast::type::VoidType void_type;
-  ast::Function func("my_func", std::move(params), &void_type);
-  func.set_body(std::move(body));
+  ast::Function func("my_func", std::move(params), &void_type, std::move(body));
 
   gen.increment_indent();
 
@@ -92,9 +90,8 @@ TEST_F(WgslGeneratorImplTest, Emit_Function_WithDecoration_WorkgroupSize) {
   body->append(create<ast::ReturnStatement>());
 
   ast::type::VoidType void_type;
-  ast::Function func("my_func", {}, &void_type);
+  ast::Function func("my_func", {}, &void_type, std::move(body));
   func.add_decoration(create<ast::WorkgroupDecoration>(2u, 4u, 6u, Source{}));
-  func.set_body(std::move(body));
 
   gen.increment_indent();
 
@@ -113,10 +110,9 @@ TEST_F(WgslGeneratorImplTest, Emit_Function_WithDecoration_Stage) {
   body->append(create<ast::ReturnStatement>());
 
   ast::type::VoidType void_type;
-  ast::Function func("my_func", {}, &void_type);
+  ast::Function func("my_func", {}, &void_type, std::move(body));
   func.add_decoration(
       create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
-  func.set_body(std::move(body));
 
   gen.increment_indent();
 
@@ -135,11 +131,10 @@ TEST_F(WgslGeneratorImplTest, Emit_Function_WithDecoration_Multiple) {
   body->append(create<ast::ReturnStatement>());
 
   ast::type::VoidType void_type;
-  ast::Function func("my_func", {}, &void_type);
+  ast::Function func("my_func", {}, &void_type, std::move(body));
   func.add_decoration(
       create<ast::StageDecoration>(ast::PipelineStage::kFragment, Source{}));
   func.add_decoration(create<ast::WorkgroupDecoration>(2u, 4u, 6u, Source{}));
-  func.set_body(std::move(body));
 
   gen.increment_indent();
 
@@ -202,10 +197,6 @@ TEST_F(WgslGeneratorImplTest,
 
   {
     ast::VariableList params;
-    auto func = create<ast::Function>("a", std::move(params), &void_type);
-    func->add_decoration(
-        create<ast::StageDecoration>(ast::PipelineStage::kCompute, Source{}));
-
     auto var = create<ast::Variable>("v", ast::StorageClass::kFunction, &f32);
     var->set_constructor(create<ast::MemberAccessorExpression>(
         create<ast::IdentifierExpression>("data"),
@@ -214,17 +205,17 @@ TEST_F(WgslGeneratorImplTest,
     auto body = create<ast::BlockStatement>();
     body->append(create<ast::VariableDeclStatement>(std::move(var)));
     body->append(create<ast::ReturnStatement>());
-    func->set_body(std::move(body));
+
+    auto func = create<ast::Function>("a", std::move(params), &void_type,
+                                      std::move(body));
+    func->add_decoration(
+        create<ast::StageDecoration>(ast::PipelineStage::kCompute, Source{}));
 
     mod.AddFunction(std::move(func));
   }
 
   {
     ast::VariableList params;
-    auto func = create<ast::Function>("b", std::move(params), &void_type);
-    func->add_decoration(
-        create<ast::StageDecoration>(ast::PipelineStage::kCompute, Source{}));
-
     auto var = create<ast::Variable>("v", ast::StorageClass::kFunction, &f32);
     var->set_constructor(create<ast::MemberAccessorExpression>(
         create<ast::IdentifierExpression>("data"),
@@ -233,7 +224,11 @@ TEST_F(WgslGeneratorImplTest,
     auto body = create<ast::BlockStatement>();
     body->append(create<ast::VariableDeclStatement>(std::move(var)));
     body->append(create<ast::ReturnStatement>());
-    func->set_body(std::move(body));
+
+    auto func = create<ast::Function>("b", std::move(params), &void_type,
+                                      std::move(body));
+    func->add_decoration(
+        create<ast::StageDecoration>(ast::PipelineStage::kCompute, Source{}));
 
     mod.AddFunction(std::move(func));
   }
diff --git a/src/writer/wgsl/generator_impl_test.cc b/src/writer/wgsl/generator_impl_test.cc
index 20f72f4656..d9ea548c9e 100644
--- a/src/writer/wgsl/generator_impl_test.cc
+++ b/src/writer/wgsl/generator_impl_test.cc
@@ -32,8 +32,9 @@ using WgslGeneratorImplTest = TestHelper;
 TEST_F(WgslGeneratorImplTest, Generate) {
   ast::type::VoidType void_type;
 
-  mod.AddFunction(
-      create<ast::Function>("my_func", ast::VariableList{}, &void_type));
+  mod.AddFunction(create<ast::Function>("my_func", ast::VariableList{},
+                                        &void_type,
+                                        create<ast::BlockStatement>()));
 
   ASSERT_TRUE(gen.Generate(mod)) << gen.error();
   EXPECT_EQ(gen.result(), R"(fn my_func() -> void {