diff --git a/src/writer/hlsl/generator_impl_array_accessor_test.cc b/src/writer/hlsl/generator_impl_array_accessor_test.cc
index 9b04593..53c42ad 100644
--- a/src/writer/hlsl/generator_impl_array_accessor_test.cc
+++ b/src/writer/hlsl/generator_impl_array_accessor_test.cc
@@ -21,8 +21,10 @@
 
 using HlslGeneratorImplTest_Expression = TestHelper;
 
-TEST_F(HlslGeneratorImplTest_Expression, EmitExpression_ArrayAccessor) {
+TEST_F(HlslGeneratorImplTest_Expression, ArrayAccessor) {
+  Global("ary", ty.array<i32, 10>(), ast::StorageClass::kPrivate);
   auto* expr = IndexAccessor("ary", 5);
+  WrapInFunction(expr);
 
   GeneratorImpl& gen = Build();
 
@@ -30,15 +32,6 @@
   EXPECT_EQ(result(), "ary[5]");
 }
 
-TEST_F(HlslGeneratorImplTest_Expression, EmitArrayAccessor) {
-  auto* expr = IndexAccessor("ary", "idx");
-
-  GeneratorImpl& gen = Build();
-
-  ASSERT_TRUE(gen.EmitExpression(pre, out, expr)) << gen.error();
-  EXPECT_EQ(result(), "ary[idx]");
-}
-
 }  // namespace
 }  // namespace hlsl
 }  // namespace writer
diff --git a/src/writer/hlsl/generator_impl_assign_test.cc b/src/writer/hlsl/generator_impl_assign_test.cc
index b53613d..cc90781 100644
--- a/src/writer/hlsl/generator_impl_assign_test.cc
+++ b/src/writer/hlsl/generator_impl_assign_test.cc
@@ -22,9 +22,10 @@
 using HlslGeneratorImplTest_Assign = TestHelper;
 
 TEST_F(HlslGeneratorImplTest_Assign, Emit_Assign) {
-  auto* lhs = Expr("lhs");
-  auto* rhs = Expr("rhs");
-  auto* assign = create<ast::AssignmentStatement>(lhs, rhs);
+  Global("lhs", ty.i32(), ast::StorageClass::kPrivate);
+  Global("rhs", ty.i32(), ast::StorageClass::kPrivate);
+  auto* assign = create<ast::AssignmentStatement>(Expr("lhs"), Expr("rhs"));
+  WrapInFunction(assign);
 
   GeneratorImpl& gen = Build();
 
diff --git a/src/writer/hlsl/generator_impl_binary_test.cc b/src/writer/hlsl/generator_impl_binary_test.cc
index 41d1289..fa06356 100644
--- a/src/writer/hlsl/generator_impl_binary_test.cc
+++ b/src/writer/hlsl/generator_impl_binary_test.cc
@@ -157,7 +157,7 @@
 }
 
 TEST_F(HlslGeneratorImplTest_Binary, Multiply_MatrixScalar) {
-  Global("mat", ty.mat3x3<f32>(), ast::StorageClass::kFunction);
+  Global("mat", ty.mat3x3<f32>(), ast::StorageClass::kPrivate);
   auto* lhs = Expr("mat");
   auto* rhs = Expr(1.f);
 
@@ -172,7 +172,7 @@
 }
 
 TEST_F(HlslGeneratorImplTest_Binary, Multiply_ScalarMatrix) {
-  Global("mat", ty.mat3x3<f32>(), ast::StorageClass::kFunction);
+  Global("mat", ty.mat3x3<f32>(), ast::StorageClass::kPrivate);
   auto* lhs = Expr(1.f);
   auto* rhs = Expr("mat");
 
@@ -187,7 +187,7 @@
 }
 
 TEST_F(HlslGeneratorImplTest_Binary, Multiply_MatrixVector) {
-  Global("mat", ty.mat3x3<f32>(), ast::StorageClass::kFunction);
+  Global("mat", ty.mat3x3<f32>(), ast::StorageClass::kPrivate);
   auto* lhs = Expr("mat");
   auto* rhs = vec3<f32>(1.f, 1.f, 1.f);
 
@@ -202,7 +202,7 @@
 }
 
 TEST_F(HlslGeneratorImplTest_Binary, Multiply_VectorMatrix) {
-  Global("mat", ty.mat3x3<f32>(), ast::StorageClass::kFunction);
+  Global("mat", ty.mat3x3<f32>(), ast::StorageClass::kPrivate);
   auto* lhs = vec3<f32>(1.f, 1.f, 1.f);
   auto* rhs = Expr("mat");
 
@@ -217,49 +217,52 @@
 }
 
 TEST_F(HlslGeneratorImplTest_Binary, Multiply_MatrixMatrix) {
-  Global("mat", ty.mat3x3<f32>(), ast::StorageClass::kFunction);
-  auto* lhs = Expr("mat");
-  auto* rhs = Expr("mat");
+  Global("lhs", ty.mat3x3<f32>(), ast::StorageClass::kPrivate);
+  Global("rhs", ty.mat3x3<f32>(), ast::StorageClass::kPrivate);
 
-  auto* expr =
-      create<ast::BinaryExpression>(ast::BinaryOp::kMultiply, lhs, rhs);
+  auto* expr = create<ast::BinaryExpression>(ast::BinaryOp::kMultiply,
+                                             Expr("lhs"), Expr("rhs"));
   WrapInFunction(expr);
 
   GeneratorImpl& gen = Build();
 
   EXPECT_TRUE(gen.EmitExpression(pre, out, expr)) << gen.error();
-  EXPECT_EQ(result(), "mul(mat, mat)");
+  EXPECT_EQ(result(), "mul(rhs, lhs)");
 }
 
 TEST_F(HlslGeneratorImplTest_Binary, Logical_And) {
-  auto* left = Expr("left");
-  auto* right = Expr("right");
+  Global("a", ty.bool_(), ast::StorageClass::kPrivate);
+  Global("b", ty.bool_(), ast::StorageClass::kPrivate);
 
-  auto* expr =
-      create<ast::BinaryExpression>(ast::BinaryOp::kLogicalAnd, left, right);
+  auto* expr = create<ast::BinaryExpression>(ast::BinaryOp::kLogicalAnd,
+                                             Expr("a"), Expr("b"));
+  WrapInFunction(expr);
 
   GeneratorImpl& gen = Build();
 
   ASSERT_TRUE(gen.EmitExpression(pre, out, expr)) << gen.error();
   EXPECT_EQ(result(), "(tint_tmp)");
-  EXPECT_EQ(pre_result(), R"(bool tint_tmp = left;
+  EXPECT_EQ(pre_result(), R"(bool tint_tmp = a;
 if (tint_tmp) {
-  tint_tmp = right;
+  tint_tmp = b;
 }
 )");
 }
 
 TEST_F(HlslGeneratorImplTest_Binary, Logical_Multi) {
   // (a && b) || (c || d)
-  auto* a = Expr("a");
-  auto* b = Expr("b");
-  auto* c = Expr("c");
-  auto* d = Expr("d");
+  Global("a", ty.bool_(), ast::StorageClass::kPrivate);
+  Global("b", ty.bool_(), ast::StorageClass::kPrivate);
+  Global("c", ty.bool_(), ast::StorageClass::kPrivate);
+  Global("d", ty.bool_(), ast::StorageClass::kPrivate);
 
   auto* expr = create<ast::BinaryExpression>(
       ast::BinaryOp::kLogicalOr,
-      create<ast::BinaryExpression>(ast::BinaryOp::kLogicalAnd, a, b),
-      create<ast::BinaryExpression>(ast::BinaryOp::kLogicalOr, c, d));
+      create<ast::BinaryExpression>(ast::BinaryOp::kLogicalAnd, Expr("a"),
+                                    Expr("b")),
+      create<ast::BinaryExpression>(ast::BinaryOp::kLogicalOr, Expr("c"),
+                                    Expr("d")));
+  WrapInFunction(expr);
 
   GeneratorImpl& gen = Build();
 
@@ -281,19 +284,20 @@
 }
 
 TEST_F(HlslGeneratorImplTest_Binary, Logical_Or) {
-  auto* left = Expr("left");
-  auto* right = Expr("right");
+  Global("a", ty.bool_(), ast::StorageClass::kPrivate);
+  Global("b", ty.bool_(), ast::StorageClass::kPrivate);
 
-  auto* expr =
-      create<ast::BinaryExpression>(ast::BinaryOp::kLogicalOr, left, right);
+  auto* expr = create<ast::BinaryExpression>(ast::BinaryOp::kLogicalOr,
+                                             Expr("a"), Expr("b"));
+  WrapInFunction(expr);
 
   GeneratorImpl& gen = Build();
 
   ASSERT_TRUE(gen.EmitExpression(pre, out, expr)) << gen.error();
   EXPECT_EQ(result(), "(tint_tmp)");
-  EXPECT_EQ(pre_result(), R"(bool tint_tmp = left;
+  EXPECT_EQ(pre_result(), R"(bool tint_tmp = a;
 if (!tint_tmp) {
-  tint_tmp = right;
+  tint_tmp = b;
 }
 )");
 }
@@ -307,13 +311,17 @@
   //   return 3;
   // }
 
+  Global("a", ty.bool_(), ast::StorageClass::kPrivate);
+  Global("b", ty.bool_(), ast::StorageClass::kPrivate);
+  Global("c", ty.bool_(), ast::StorageClass::kPrivate);
+
   auto* body = create<ast::BlockStatement>(ast::StatementList{
-      create<ast::ReturnStatement>(Expr(3)),
+      Return(Expr(3)),
   });
   auto* else_stmt = create<ast::ElseStatement>(nullptr, body);
 
   body = create<ast::BlockStatement>(ast::StatementList{
-      create<ast::ReturnStatement>(Expr(2)),
+      Return(Expr(2)),
   });
   auto* else_if_stmt = create<ast::ElseStatement>(
       create<ast::BinaryExpression>(ast::BinaryOp::kLogicalOr, Expr("b"),
@@ -321,7 +329,7 @@
       body);
 
   body = create<ast::BlockStatement>(ast::StatementList{
-      create<ast::ReturnStatement>(Expr(1)),
+      Return(Expr(1)),
   });
 
   auto* expr = create<ast::IfStatement>(
@@ -332,6 +340,7 @@
           else_if_stmt,
           else_stmt,
       });
+  Func("func", {}, ty.i32(), {WrapInStatement(expr), Return(0)});
 
   GeneratorImpl& gen = Build();
 
@@ -358,13 +367,17 @@
 
 TEST_F(HlslGeneratorImplTest_Binary, Return_WithLogical) {
   // return (a && b) || c;
-  auto* a = Expr("a");
-  auto* b = Expr("b");
-  auto* c = Expr("c");
 
-  auto* expr = create<ast::ReturnStatement>(create<ast::BinaryExpression>(
+  Global("a", ty.bool_(), ast::StorageClass::kPrivate);
+  Global("b", ty.bool_(), ast::StorageClass::kPrivate);
+  Global("c", ty.bool_(), ast::StorageClass::kPrivate);
+
+  auto* expr = Return(create<ast::BinaryExpression>(
       ast::BinaryOp::kLogicalOr,
-      create<ast::BinaryExpression>(ast::BinaryOp::kLogicalAnd, a, b), c));
+      create<ast::BinaryExpression>(ast::BinaryOp::kLogicalAnd, Expr("a"),
+                                    Expr("b")),
+      Expr("c")));
+  Func("func", {}, ty.bool_(), {WrapInStatement(expr)});
 
   GeneratorImpl& gen = Build();
 
@@ -383,16 +396,19 @@
 
 TEST_F(HlslGeneratorImplTest_Binary, Assign_WithLogical) {
   // a = (b || c) && d;
-  auto* a = Expr("a");
-  auto* b = Expr("b");
-  auto* c = Expr("c");
-  auto* d = Expr("d");
+
+  Global("a", ty.bool_(), ast::StorageClass::kPrivate);
+  Global("b", ty.bool_(), ast::StorageClass::kPrivate);
+  Global("c", ty.bool_(), ast::StorageClass::kPrivate);
+  Global("d", ty.bool_(), ast::StorageClass::kPrivate);
 
   auto* expr = create<ast::AssignmentStatement>(
-      a,
-      create<ast::BinaryExpression>(
-          ast::BinaryOp::kLogicalAnd,
-          create<ast::BinaryExpression>(ast::BinaryOp::kLogicalOr, b, c), d));
+      Expr("a"), create<ast::BinaryExpression>(
+                     ast::BinaryOp::kLogicalAnd,
+                     create<ast::BinaryExpression>(ast::BinaryOp::kLogicalOr,
+                                                   Expr("b"), Expr("c")),
+                     Expr("d")));
+  WrapInFunction(expr);
 
   GeneratorImpl& gen = Build();
 
@@ -412,22 +428,19 @@
 TEST_F(HlslGeneratorImplTest_Binary, Decl_WithLogical) {
   // var a : bool = (b && c) || d;
 
-  auto* b_decl = Decl(Var("b", ty.bool_(), ast::StorageClass::kFunction));
-  auto* c_decl = Decl(Var("c", ty.bool_(), ast::StorageClass::kFunction));
-  auto* d_decl = Decl(Var("d", ty.bool_(), ast::StorageClass::kFunction));
+  Global("b", ty.bool_(), ast::StorageClass::kPrivate);
+  Global("c", ty.bool_(), ast::StorageClass::kPrivate);
+  Global("d", ty.bool_(), ast::StorageClass::kPrivate);
 
-  auto* b = Expr("b");
-  auto* c = Expr("c");
-  auto* d = Expr("d");
-
-  auto* var = Var(
-      "a", ty.bool_(), ast::StorageClass::kFunction,
-      create<ast::BinaryExpression>(
-          ast::BinaryOp::kLogicalOr,
-          create<ast::BinaryExpression>(ast::BinaryOp::kLogicalAnd, b, c), d));
+  auto* var = Var("a", ty.bool_(), ast::StorageClass::kFunction,
+                  create<ast::BinaryExpression>(
+                      ast::BinaryOp::kLogicalOr,
+                      create<ast::BinaryExpression>(ast::BinaryOp::kLogicalAnd,
+                                                    Expr("b"), Expr("c")),
+                      Expr("d")));
 
   auto* decl = Decl(var);
-  WrapInFunction(b_decl, c_decl, d_decl, Decl(var));
+  WrapInFunction(decl);
 
   GeneratorImpl& gen = Build();
 
@@ -447,15 +460,16 @@
 TEST_F(HlslGeneratorImplTest_Binary, Bitcast_WithLogical) {
   // as<i32>(a && (b || c))
 
-  auto* a = Expr("a");
-  auto* b = Expr("b");
-  auto* c = Expr("c");
+  Global("a", ty.bool_(), ast::StorageClass::kPrivate);
+  Global("b", ty.bool_(), ast::StorageClass::kPrivate);
+  Global("c", ty.bool_(), ast::StorageClass::kPrivate);
 
   auto* expr = create<ast::BitcastExpression>(
-      ty.i32(),
-      create<ast::BinaryExpression>(
-          ast::BinaryOp::kLogicalAnd, a,
-          create<ast::BinaryExpression>(ast::BinaryOp::kLogicalOr, b, c)));
+      ty.i32(), create<ast::BinaryExpression>(
+                    ast::BinaryOp::kLogicalAnd, Expr("a"),
+                    create<ast::BinaryExpression>(ast::BinaryOp::kLogicalOr,
+                                                  Expr("b"), Expr("c"))));
+  WrapInFunction(expr);
 
   GeneratorImpl& gen = Build();
 
@@ -477,10 +491,10 @@
 
   Func("foo", ast::VariableList{}, ty.void_(), ast::StatementList{},
        ast::DecorationList{});
-  Global("a", ty.bool_(), ast::StorageClass::kInput);
-  Global("b", ty.bool_(), ast::StorageClass::kInput);
-  Global("c", ty.bool_(), ast::StorageClass::kInput);
-  Global("d", ty.bool_(), ast::StorageClass::kInput);
+  Global("a", ty.bool_(), ast::StorageClass::kPrivate);
+  Global("b", ty.bool_(), ast::StorageClass::kPrivate);
+  Global("c", ty.bool_(), ast::StorageClass::kPrivate);
+  Global("d", ty.bool_(), ast::StorageClass::kPrivate);
 
   ast::ExpressionList params;
   params.push_back(create<ast::BinaryExpression>(ast::BinaryOp::kLogicalAnd,
diff --git a/src/writer/hlsl/generator_impl_bitcast_test.cc b/src/writer/hlsl/generator_impl_bitcast_test.cc
index 51c367f..45de152 100644
--- a/src/writer/hlsl/generator_impl_bitcast_test.cc
+++ b/src/writer/hlsl/generator_impl_bitcast_test.cc
@@ -22,33 +22,33 @@
 using HlslGeneratorImplTest_Bitcast = TestHelper;
 
 TEST_F(HlslGeneratorImplTest_Bitcast, EmitExpression_Bitcast_Float) {
-  auto* id = Expr("id");
-  auto* bitcast = create<ast::BitcastExpression>(ty.f32(), id);
+  auto* bitcast = create<ast::BitcastExpression>(ty.f32(), Expr(1));
+  WrapInFunction(bitcast);
 
   GeneratorImpl& gen = Build();
 
   ASSERT_TRUE(gen.EmitExpression(pre, out, bitcast)) << gen.error();
-  EXPECT_EQ(result(), "asfloat(id)");
+  EXPECT_EQ(result(), "asfloat(1)");
 }
 
 TEST_F(HlslGeneratorImplTest_Bitcast, EmitExpression_Bitcast_Int) {
-  auto* id = Expr("id");
-  auto* bitcast = create<ast::BitcastExpression>(ty.i32(), id);
+  auto* bitcast = create<ast::BitcastExpression>(ty.i32(), Expr(1u));
+  WrapInFunction(bitcast);
 
   GeneratorImpl& gen = Build();
 
   ASSERT_TRUE(gen.EmitExpression(pre, out, bitcast)) << gen.error();
-  EXPECT_EQ(result(), "asint(id)");
+  EXPECT_EQ(result(), "asint(1u)");
 }
 
 TEST_F(HlslGeneratorImplTest_Bitcast, EmitExpression_Bitcast_Uint) {
-  auto* id = Expr("id");
-  auto* bitcast = create<ast::BitcastExpression>(ty.u32(), id);
+  auto* bitcast = create<ast::BitcastExpression>(ty.u32(), Expr(1));
+  WrapInFunction(bitcast);
 
   GeneratorImpl& gen = Build();
 
   ASSERT_TRUE(gen.EmitExpression(pre, out, bitcast)) << gen.error();
-  EXPECT_EQ(result(), "asuint(id)");
+  EXPECT_EQ(result(), "asuint(1)");
 }
 
 }  // namespace
diff --git a/src/writer/hlsl/generator_impl_block_test.cc b/src/writer/hlsl/generator_impl_block_test.cc
index 851c37c..535304b 100644
--- a/src/writer/hlsl/generator_impl_block_test.cc
+++ b/src/writer/hlsl/generator_impl_block_test.cc
@@ -22,9 +22,8 @@
 using HlslGeneratorImplTest_Block = TestHelper;
 
 TEST_F(HlslGeneratorImplTest_Block, Emit_Block) {
-  auto* b = create<ast::BlockStatement>(ast::StatementList{
-      create<ast::DiscardStatement>(),
-  });
+  auto* b = Block(create<ast::DiscardStatement>());
+  WrapInFunction(b);
 
   GeneratorImpl& gen = Build();
 
@@ -38,9 +37,8 @@
 }
 
 TEST_F(HlslGeneratorImplTest_Block, Emit_Block_WithoutNewline) {
-  auto* b = create<ast::BlockStatement>(ast::StatementList{
-      create<ast::DiscardStatement>(),
-  });
+  auto* b = Block(create<ast::DiscardStatement>());
+  WrapInFunction(b);
 
   GeneratorImpl& gen = Build();
 
diff --git a/src/writer/hlsl/generator_impl_break_test.cc b/src/writer/hlsl/generator_impl_break_test.cc
index 09fef93..dbf3241 100644
--- a/src/writer/hlsl/generator_impl_break_test.cc
+++ b/src/writer/hlsl/generator_impl_break_test.cc
@@ -23,6 +23,7 @@
 
 TEST_F(HlslGeneratorImplTest_Break, Emit_Break) {
   auto* b = create<ast::BreakStatement>();
+  WrapInFunction(Loop(Block(b)));
 
   GeneratorImpl& gen = Build();
 
diff --git a/src/writer/hlsl/generator_impl_call_test.cc b/src/writer/hlsl/generator_impl_call_test.cc
index 189e44e..b5bb089 100644
--- a/src/writer/hlsl/generator_impl_call_test.cc
+++ b/src/writer/hlsl/generator_impl_call_test.cc
@@ -38,8 +38,8 @@
 TEST_F(HlslGeneratorImplTest_Call, EmitExpression_Call_WithParams) {
   Func("my_func", ast::VariableList{}, ty.void_(), ast::StatementList{},
        ast::DecorationList{});
-  Global("param1", ty.f32(), ast::StorageClass::kInput);
-  Global("param2", ty.f32(), ast::StorageClass::kInput);
+  Global("param1", ty.f32(), ast::StorageClass::kPrivate);
+  Global("param2", ty.f32(), ast::StorageClass::kPrivate);
 
   auto* call = Call("my_func", "param1", "param2");
   WrapInFunction(call);
@@ -53,8 +53,8 @@
 TEST_F(HlslGeneratorImplTest_Call, EmitStatement_Call) {
   Func("my_func", ast::VariableList{}, ty.void_(), ast::StatementList{},
        ast::DecorationList{});
-  Global("param1", ty.f32(), ast::StorageClass::kInput);
-  Global("param2", ty.f32(), ast::StorageClass::kInput);
+  Global("param1", ty.f32(), ast::StorageClass::kPrivate);
+  Global("param2", ty.f32(), ast::StorageClass::kPrivate);
 
   auto* call = create<ast::CallStatement>(Call("my_func", "param1", "param2"));
   WrapInFunction(call);
diff --git a/src/writer/hlsl/generator_impl_case_test.cc b/src/writer/hlsl/generator_impl_case_test.cc
index 7e3c23d..b892acb 100644
--- a/src/writer/hlsl/generator_impl_case_test.cc
+++ b/src/writer/hlsl/generator_impl_case_test.cc
@@ -23,12 +23,11 @@
 using HlslGeneratorImplTest_Case = TestHelper;
 
 TEST_F(HlslGeneratorImplTest_Case, Emit_Case) {
-  auto* body = create<ast::BlockStatement>(ast::StatementList{
-      create<ast::BreakStatement>(),
-  });
+  auto* body = Block(create<ast::BreakStatement>());
   ast::CaseSelectorList lit;
   lit.push_back(Literal(5));
   auto* c = create<ast::CaseStatement>(lit, body);
+  WrapInFunction(c);
 
   GeneratorImpl& gen = Build();
 
@@ -44,8 +43,8 @@
 TEST_F(HlslGeneratorImplTest_Case, Emit_Case_BreaksByDefault) {
   ast::CaseSelectorList lit;
   lit.push_back(Literal(5));
-  auto* c = create<ast::CaseStatement>(
-      lit, create<ast::BlockStatement>(ast::StatementList{}));
+  auto* c = create<ast::CaseStatement>(lit, Block());
+  WrapInFunction(c);
 
   GeneratorImpl& gen = Build();
 
@@ -59,12 +58,11 @@
 }
 
 TEST_F(HlslGeneratorImplTest_Case, Emit_Case_WithFallthrough) {
-  auto* body = create<ast::BlockStatement>(ast::StatementList{
-      create<ast::FallthroughStatement>(),
-  });
+  auto* body = Block(create<ast::FallthroughStatement>());
   ast::CaseSelectorList lit;
   lit.push_back(Literal(5));
   auto* c = create<ast::CaseStatement>(lit, body);
+  WrapInFunction(c);
 
   GeneratorImpl& gen = Build();
 
@@ -78,13 +76,12 @@
 }
 
 TEST_F(HlslGeneratorImplTest_Case, Emit_Case_MultipleSelectors) {
-  auto* body = create<ast::BlockStatement>(ast::StatementList{
-      create<ast::BreakStatement>(),
-  });
+  auto* body = Block(create<ast::BreakStatement>());
   ast::CaseSelectorList lit;
   lit.push_back(Literal(5));
   lit.push_back(Literal(6));
   auto* c = create<ast::CaseStatement>(lit, body);
+  WrapInFunction(c);
 
   GeneratorImpl& gen = Build();
 
@@ -99,10 +96,9 @@
 }
 
 TEST_F(HlslGeneratorImplTest_Case, Emit_Case_Default) {
-  auto* body = create<ast::BlockStatement>(ast::StatementList{
-      create<ast::BreakStatement>(),
-  });
+  auto* body = Block(create<ast::BreakStatement>());
   auto* c = create<ast::CaseStatement>(ast::CaseSelectorList{}, body);
+  WrapInFunction(c);
 
   GeneratorImpl& gen = Build();
 
diff --git a/src/writer/hlsl/generator_impl_cast_test.cc b/src/writer/hlsl/generator_impl_cast_test.cc
index ac9cf54..f123536 100644
--- a/src/writer/hlsl/generator_impl_cast_test.cc
+++ b/src/writer/hlsl/generator_impl_cast_test.cc
@@ -22,21 +22,23 @@
 using HlslGeneratorImplTest_Cast = TestHelper;
 
 TEST_F(HlslGeneratorImplTest_Cast, EmitExpression_Cast_Scalar) {
-  auto* cast = Construct<f32>("id");
+  auto* cast = Construct<f32>(1);
+  WrapInFunction(cast);
 
   GeneratorImpl& gen = Build();
 
   ASSERT_TRUE(gen.EmitExpression(pre, out, cast)) << gen.error();
-  EXPECT_EQ(result(), "float(id)");
+  EXPECT_EQ(result(), "float(1)");
 }
 
 TEST_F(HlslGeneratorImplTest_Cast, EmitExpression_Cast_Vector) {
-  auto* cast = vec3<f32>("id");
+  auto* cast = vec3<f32>(vec3<i32>(1, 2, 3));
+  WrapInFunction(cast);
 
   GeneratorImpl& gen = Build();
 
   ASSERT_TRUE(gen.EmitExpression(pre, out, cast)) << gen.error();
-  EXPECT_EQ(result(), "float3(id)");
+  EXPECT_EQ(result(), "float3(int3(1, 2, 3))");
 }
 
 }  // namespace
diff --git a/src/writer/hlsl/generator_impl_continue_test.cc b/src/writer/hlsl/generator_impl_continue_test.cc
index fd1cc0b..07d6b57 100644
--- a/src/writer/hlsl/generator_impl_continue_test.cc
+++ b/src/writer/hlsl/generator_impl_continue_test.cc
@@ -23,6 +23,7 @@
 
 TEST_F(HlslGeneratorImplTest_Continue, Emit_Continue) {
   auto* c = create<ast::ContinueStatement>();
+  WrapInFunction(Loop(Block(c)));
 
   GeneratorImpl& gen = Build();
 
diff --git a/src/writer/hlsl/generator_impl_discard_test.cc b/src/writer/hlsl/generator_impl_discard_test.cc
index 50e4453..e738ef7 100644
--- a/src/writer/hlsl/generator_impl_discard_test.cc
+++ b/src/writer/hlsl/generator_impl_discard_test.cc
@@ -23,6 +23,7 @@
 
 TEST_F(HlslGeneratorImplTest_Discard, Emit_Discard) {
   auto* stmt = create<ast::DiscardStatement>();
+  WrapInFunction(stmt);
 
   GeneratorImpl& gen = Build();
 
diff --git a/src/writer/hlsl/generator_impl_identifier_test.cc b/src/writer/hlsl/generator_impl_identifier_test.cc
index 1faa7ae..161cab0 100644
--- a/src/writer/hlsl/generator_impl_identifier_test.cc
+++ b/src/writer/hlsl/generator_impl_identifier_test.cc
@@ -22,7 +22,10 @@
 using HlslGeneratorImplTest_Identifier = TestHelper;
 
 TEST_F(HlslGeneratorImplTest_Identifier, EmitIdentifierExpression) {
+  Global("foo", ty.i32(), ast::StorageClass::kPrivate);
+
   auto* i = Expr("foo");
+  WrapInFunction(i);
 
   GeneratorImpl& gen = Build();
 
diff --git a/src/writer/hlsl/generator_impl_if_test.cc b/src/writer/hlsl/generator_impl_if_test.cc
index 60d39b9..f2aece3 100644
--- a/src/writer/hlsl/generator_impl_if_test.cc
+++ b/src/writer/hlsl/generator_impl_if_test.cc
@@ -22,11 +22,12 @@
 using HlslGeneratorImplTest_If = TestHelper;
 
 TEST_F(HlslGeneratorImplTest_If, Emit_If) {
+  Global("cond", ty.bool_(), ast::StorageClass::kPrivate);
+
   auto* cond = Expr("cond");
-  auto* body = create<ast::BlockStatement>(ast::StatementList{
-      create<ast::ReturnStatement>(),
-  });
-  auto* i = create<ast::IfStatement>(cond, body, ast::ElseStatementList{});
+  auto* body = Block(Return());
+  auto* i = If(cond, body);
+  WrapInFunction(i);
 
   GeneratorImpl& gen = Build();
 
@@ -39,18 +40,18 @@
 }
 
 TEST_F(HlslGeneratorImplTest_If, Emit_IfWithElseIf) {
+  Global("cond", ty.bool_(), ast::StorageClass::kPrivate);
+  Global("else_cond", ty.bool_(), ast::StorageClass::kPrivate);
+
   auto* else_cond = Expr("else_cond");
-  auto* else_body = create<ast::BlockStatement>(ast::StatementList{
-      create<ast::ReturnStatement>(),
-  });
+  auto* else_body = Block(Return());
 
   auto* cond = Expr("cond");
-  auto* body = create<ast::BlockStatement>(ast::StatementList{
-      create<ast::ReturnStatement>(),
-  });
-  auto* i = create<ast::IfStatement>(
+  auto* body = Block(Return());
+  auto* i = If(
       cond, body,
       ast::ElseStatementList{create<ast::ElseStatement>(else_cond, else_body)});
+  WrapInFunction(i);
 
   GeneratorImpl& gen = Build();
 
@@ -68,17 +69,16 @@
 }
 
 TEST_F(HlslGeneratorImplTest_If, Emit_IfWithElse) {
-  auto* else_body = create<ast::BlockStatement>(ast::StatementList{
-      create<ast::ReturnStatement>(),
-  });
+  Global("cond", ty.bool_(), ast::StorageClass::kPrivate);
+
+  auto* else_body = Block(Return());
 
   auto* cond = Expr("cond");
-  auto* body = create<ast::BlockStatement>(ast::StatementList{
-      create<ast::ReturnStatement>(),
-  });
-  auto* i = create<ast::IfStatement>(
+  auto* body = Block(Return());
+  auto* i = If(
       cond, body,
       ast::ElseStatementList{create<ast::ElseStatement>(nullptr, else_body)});
+  WrapInFunction(i);
 
   GeneratorImpl& gen = Build();
 
@@ -94,26 +94,23 @@
 }
 
 TEST_F(HlslGeneratorImplTest_If, Emit_IfWithMultiple) {
+  Global("cond", ty.bool_(), ast::StorageClass::kPrivate);
+  Global("else_cond", ty.bool_(), ast::StorageClass::kPrivate);
+
   auto* else_cond = Expr("else_cond");
 
-  auto* else_body = create<ast::BlockStatement>(ast::StatementList{
-      create<ast::ReturnStatement>(),
-  });
+  auto* else_body = Block(Return());
 
-  auto* else_body_2 = create<ast::BlockStatement>(ast::StatementList{
-      create<ast::ReturnStatement>(),
-  });
+  auto* else_body_2 = Block(Return());
 
   auto* cond = Expr("cond");
-  auto* body = create<ast::BlockStatement>(ast::StatementList{
-      create<ast::ReturnStatement>(),
-  });
-  auto* i = create<ast::IfStatement>(
-      cond, body,
-      ast::ElseStatementList{
-          create<ast::ElseStatement>(else_cond, else_body),
-          create<ast::ElseStatement>(nullptr, else_body_2),
-      });
+  auto* body = Block(Return());
+  auto* i = If(cond, body,
+               ast::ElseStatementList{
+                   create<ast::ElseStatement>(else_cond, else_body),
+                   create<ast::ElseStatement>(nullptr, else_body_2),
+               });
+  WrapInFunction(i);
 
   GeneratorImpl& gen = Build();
 
diff --git a/src/writer/hlsl/generator_impl_loop_test.cc b/src/writer/hlsl/generator_impl_loop_test.cc
index 62b5278..045dd66 100644
--- a/src/writer/hlsl/generator_impl_loop_test.cc
+++ b/src/writer/hlsl/generator_impl_loop_test.cc
@@ -23,10 +23,8 @@
 using HlslGeneratorImplTest_Loop = TestHelper;
 
 TEST_F(HlslGeneratorImplTest_Loop, Emit_Loop) {
-  auto* body = create<ast::BlockStatement>(ast::StatementList{
-      create<ast::DiscardStatement>(),
-  });
-  auto* continuing = create<ast::BlockStatement>(ast::StatementList{});
+  auto* body = Block(create<ast::DiscardStatement>());
+  auto* continuing = Block();
   auto* l = create<ast::LoopStatement>(body, continuing);
 
   WrapInFunction(l);
@@ -43,12 +41,8 @@
 }
 
 TEST_F(HlslGeneratorImplTest_Loop, Emit_LoopWithContinuing) {
-  auto* body = create<ast::BlockStatement>(ast::StatementList{
-      create<ast::DiscardStatement>(),
-  });
-  auto* continuing = create<ast::BlockStatement>(ast::StatementList{
-      create<ast::ReturnStatement>(),
-  });
+  auto* body = Block(create<ast::DiscardStatement>());
+  auto* continuing = Block(Return());
   auto* l = create<ast::LoopStatement>(body, continuing);
 
   WrapInFunction(l);
@@ -73,27 +67,19 @@
 }
 
 TEST_F(HlslGeneratorImplTest_Loop, Emit_LoopNestedWithContinuing) {
-  Global("lhs", ty.f32(), ast::StorageClass::kInput);
-  Global("rhs", ty.f32(), ast::StorageClass::kInput);
+  Global("lhs", ty.f32(), ast::StorageClass::kPrivate);
+  Global("rhs", ty.f32(), ast::StorageClass::kPrivate);
 
-  auto* body = create<ast::BlockStatement>(ast::StatementList{
-      create<ast::DiscardStatement>(),
-  });
-  auto* continuing = create<ast::BlockStatement>(ast::StatementList{
-      create<ast::ReturnStatement>(),
-  });
+  auto* body = Block(create<ast::DiscardStatement>());
+  auto* continuing = Block(Return());
   auto* inner = create<ast::LoopStatement>(body, continuing);
 
-  body = create<ast::BlockStatement>(ast::StatementList{
-      inner,
-  });
+  body = Block(inner);
 
   auto* lhs = Expr("lhs");
   auto* rhs = Expr("rhs");
 
-  continuing = create<ast::BlockStatement>(ast::StatementList{
-      create<ast::AssignmentStatement>(lhs, rhs),
-  });
+  continuing = Block(create<ast::AssignmentStatement>(lhs, rhs));
 
   auto* outer = create<ast::LoopStatement>(body, continuing);
   WrapInFunction(outer);
@@ -149,22 +135,18 @@
   //   }
   // }
 
-  Global("rhs", ty.f32(), ast::StorageClass::kInput);
+  Global("rhs", ty.f32(), ast::StorageClass::kPrivate);
 
   auto* var = Var("lhs", ty.f32(), ast::StorageClass::kFunction, Expr(2.4f));
 
-  auto* body = create<ast::BlockStatement>(ast::StatementList{
-      create<ast::VariableDeclStatement>(var),
-      create<ast::VariableDeclStatement>(
-          Var("other", ty.f32(), ast::StorageClass::kFunction)),
-  });
+  auto* body = Block(create<ast::VariableDeclStatement>(var),
+                     create<ast::VariableDeclStatement>(
+                         Var("other", ty.f32(), ast::StorageClass::kFunction)));
 
   auto* lhs = Expr("lhs");
   auto* rhs = Expr("rhs");
 
-  auto* continuing = create<ast::BlockStatement>(ast::StatementList{
-      create<ast::AssignmentStatement>(lhs, rhs),
-  });
+  auto* continuing = Block(create<ast::AssignmentStatement>(lhs, rhs));
   auto* outer = create<ast::LoopStatement>(body, continuing);
   WrapInFunction(outer);
 
diff --git a/src/writer/hlsl/generator_impl_return_test.cc b/src/writer/hlsl/generator_impl_return_test.cc
index 48c31c6..5a4bf74 100644
--- a/src/writer/hlsl/generator_impl_return_test.cc
+++ b/src/writer/hlsl/generator_impl_return_test.cc
@@ -22,7 +22,8 @@
 using HlslGeneratorImplTest_Return = TestHelper;
 
 TEST_F(HlslGeneratorImplTest_Return, Emit_Return) {
-  auto* r = create<ast::ReturnStatement>();
+  auto* r = Return();
+  WrapInFunction(r);
 
   GeneratorImpl& gen = Build();
 
@@ -33,14 +34,15 @@
 }
 
 TEST_F(HlslGeneratorImplTest_Return, Emit_ReturnWithValue) {
-  auto* r = create<ast::ReturnStatement>(Expr("expr"));
+  auto* r = Return(123);
+  Func("f", {}, ty.i32(), {r});
 
   GeneratorImpl& gen = Build();
 
   gen.increment_indent();
 
   ASSERT_TRUE(gen.EmitStatement(out, r)) << gen.error();
-  EXPECT_EQ(result(), "  return expr;\n");
+  EXPECT_EQ(result(), "  return 123;\n");
 }
 
 }  // namespace
diff --git a/src/writer/hlsl/generator_impl_switch_test.cc b/src/writer/hlsl/generator_impl_switch_test.cc
index ba86e74..6412524 100644
--- a/src/writer/hlsl/generator_impl_switch_test.cc
+++ b/src/writer/hlsl/generator_impl_switch_test.cc
@@ -22,17 +22,15 @@
 using HlslGeneratorImplTest_Switch = TestHelper;
 
 TEST_F(HlslGeneratorImplTest_Switch, Emit_Switch) {
-  auto* def_body = create<ast::BlockStatement>(ast::StatementList{
-      create<ast::BreakStatement>(),
-  });
+  Global("cond", ty.i32(), ast::StorageClass::kPrivate);
+
+  auto* def_body = Block(create<ast::BreakStatement>());
   auto* def = create<ast::CaseStatement>(ast::CaseSelectorList{}, def_body);
 
   ast::CaseSelectorList case_val;
   case_val.push_back(Literal(5));
 
-  auto* case_body = create<ast::BlockStatement>(ast::StatementList{
-      create<ast::BreakStatement>(),
-  });
+  auto* case_body = Block(create<ast::BreakStatement>());
 
   auto* case_stmt = create<ast::CaseStatement>(case_val, case_body);
 
@@ -42,6 +40,7 @@
 
   auto* cond = Expr("cond");
   auto* s = create<ast::SwitchStatement>(cond, body);
+  WrapInFunction(s);
 
   GeneratorImpl& gen = Build();
 
diff --git a/src/writer/hlsl/generator_impl_unary_op_test.cc b/src/writer/hlsl/generator_impl_unary_op_test.cc
index 2a393d3..c1d973b 100644
--- a/src/writer/hlsl/generator_impl_unary_op_test.cc
+++ b/src/writer/hlsl/generator_impl_unary_op_test.cc
@@ -31,8 +31,10 @@
 TEST_P(HlslUnaryOpTest, Emit) {
   auto params = GetParam();
 
-  auto* expr = Expr("expr");
-  auto* op = create<ast::UnaryOpExpression>(params.op, expr);
+  Global("expr", ty.i32(), ast::StorageClass::kPrivate);
+
+  auto* op = create<ast::UnaryOpExpression>(params.op, Expr("expr"));
+  WrapInFunction(op);
 
   GeneratorImpl& gen = Build();
 
diff --git a/src/writer/hlsl/generator_impl_variable_decl_statement_test.cc b/src/writer/hlsl/generator_impl_variable_decl_statement_test.cc
index e69f8df..441312c 100644
--- a/src/writer/hlsl/generator_impl_variable_decl_statement_test.cc
+++ b/src/writer/hlsl/generator_impl_variable_decl_statement_test.cc
@@ -79,7 +79,7 @@
 
 TEST_F(HlslGeneratorImplTest_VariableDecl,
        Emit_VariableDeclStatement_Initializer_Private) {
-  Global("initializer", ty.f32(), ast::StorageClass::kInput);
+  Global("initializer", ty.f32(), ast::StorageClass::kPrivate);
   Global("a", ty.f32(), ast::StorageClass::kPrivate, Expr("initializer"));
 
   WrapInFunction(Expr("a"));