src/transform: Reimplement tests in WGSL

Easier to read and write, and ensures that the tests exercise valid AST instead of synthetic structures that can never exist.

Change-Id: I5d361ef96383c71943a424f5765952f21d740042
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/36422
Reviewed-by: dan sinclair <dsinclair@chromium.org>
Commit-Queue: Ben Clayton <bclayton@google.com>
Auto-Submit: Ben Clayton <bclayton@google.com>

BUILD.gn

@@ -831,6 +831,7 @@ source_set("tint_unittests_core_src") {
     "src/transform/bound_array_accessors_test.cc",
     "src/transform/emit_vertex_point_size_test.cc",
     "src/transform/first_index_offset_test.cc",
+    "src/transform/test_helper.h",
     "src/transform/vertex_pulling_test.cc",
     "src/type_determiner_test.cc",
     "src/validator/validator_control_block_test.cc",

src/CMakeLists.txt

@@ -465,10 +465,6 @@ if(${TINT_BUILD_TESTS})
     scope_stack_test.cc
     symbol_table_test.cc
     symbol_test.cc
-    transform/emit_vertex_point_size_test.cc
-    transform/bound_array_accessors_test.cc
-    transform/first_index_offset_test.cc
-    transform/vertex_pulling_test.cc
     type_determiner_test.cc
     validator/validator_control_block_test.cc
     validator/validator_function_test.cc
@@ -656,6 +652,16 @@ if(${TINT_BUILD_TESTS})
     )
   endif()
 
+  if(${TINT_BUILD_WGSL_READER} AND ${TINT_BUILD_WGSL_WRITER})
+    list(APPEND TINT_TEST_SRCS
+      transform/bound_array_accessors_test.cc
+      transform/emit_vertex_point_size_test.cc
+      transform/first_index_offset_test.cc
+      transform/test_helper.h
+      transform/vertex_pulling_test.cc
+    )
+  endif()
+
   if(${TINT_BUILD_MSL_WRITER})
     list(APPEND TINT_TEST_SRCS
       writer/msl/generator_impl_alias_type_test.cc

src/transform/emit_vertex_point_size_test.cc

@@ -14,203 +14,106 @@
 
 #include "src/transform/emit_vertex_point_size.h"
 
-#include <memory>
-#include <utility>
-
-#include "gtest/gtest.h"
-#include "src/ast/builder.h"
-#include "src/ast/stage_decoration.h"
-#include "src/ast/variable_decl_statement.h"
-#include "src/demangler.h"
-#include "src/diagnostic/formatter.h"
-#include "src/transform/manager.h"
+#include "src/transform/test_helper.h"
 
 namespace tint {
 namespace transform {
 namespace {
 
-class EmitVertexPointSizeTest : public testing::Test {
- public:
-  Transform::Output GetTransform(ast::Module in) {
-    Manager manager;
-    manager.append(std::make_unique<EmitVertexPointSize>());
-    return manager.Run(&in);
-  }
-};
-
-struct ModuleBuilder : public ast::BuilderWithModule {
-  ModuleBuilder() {}
-
-  ast::Module Module() {
-    Build();
-    return std::move(*mod);
-  }
-
- protected:
-  virtual void Build() = 0;
-};
+using EmitVertexPointSizeTest = TransformTest;
 
 TEST_F(EmitVertexPointSizeTest, VertexStageBasic) {
-  struct Builder : ModuleBuilder {
-    void Build() override {
-      mod->AddFunction(Func("non_entry_a", ast::VariableList{}, ty.void_,
-                            ast::StatementList{},
-                            ast::FunctionDecorationList{}));
+  auto* src = R"(
+fn non_entry_a() -> void {
+}
 
-      auto* entry =
-          Func("entry", ast::VariableList{}, ty.void_,
-               ast::StatementList{
-                   create<ast::VariableDeclStatement>(
-                       Var("builtin_assignments_should_happen_before_this",
-                           tint::ast::StorageClass::kFunction, ty.f32)),
-               },
-               ast::FunctionDecorationList{
-                   create<ast::StageDecoration>(ast::PipelineStage::kVertex),
-               });
-      mod->AddFunction(entry);
+[[stage(vertex)]]
+fn entry() -> void {
+  var builtin_assignments_should_happen_before_this : f32;
+}
 
-      mod->AddFunction(Func("non_entry_b", ast::VariableList{}, ty.void_,
-                            ast::StatementList{},
-                            ast::FunctionDecorationList{}));
-    }
-  };
-
-  auto result = GetTransform(Builder{}.Module());
-  ASSERT_FALSE(result.diagnostics.contains_errors())
-      << diag::Formatter().format(result.diagnostics);
-
-  auto* expected = R"(Module{
-  Variable{
-    Decorations{
-      BuiltinDecoration{pointsize}
-    }
-    tint_pointsize
-    out
-    __f32
-  }
-  Function non_entry_a -> __void
-  ()
-  {
-  }
-  Function entry -> __void
-  StageDecoration{vertex}
-  ()
-  {
-    Assignment{
-      Identifier[__ptr_out__f32]{tint_pointsize}
-      ScalarConstructor[__f32]{1.000000}
-    }
-    VariableDeclStatement{
-      Variable{
-        builtin_assignments_should_happen_before_this
-        function
-        __f32
-      }
-    }
-  }
-  Function non_entry_b -> __void
-  ()
-  {
-  }
+fn non_entry_b() -> void {
 }
 )";
-  EXPECT_EQ(expected,
-            Demangler().Demangle(result.module, result.module.to_str()));
+
+  auto* expect = R"(
+[[builtin(pointsize)]] var<out> tint_pointsize : f32;
+
+fn non_entry_a() -> void {
+}
+
+[[stage(vertex)]]
+fn entry() -> void {
+  tint_pointsize = 1.0;
+  var builtin_assignments_should_happen_before_this : f32;
+}
+
+fn non_entry_b() -> void {
+}
+)";
+
+  auto got = Transform<EmitVertexPointSize>(src);
+
+  EXPECT_EQ(expect, got);
 }
 
 TEST_F(EmitVertexPointSizeTest, VertexStageEmpty) {
-  struct Builder : ModuleBuilder {
-    void Build() override {
-      mod->AddFunction(Func("non_entry_a", ast::VariableList{}, ty.void_,
-                            ast::StatementList{},
-                            ast::FunctionDecorationList{}));
+  auto* src = R"(
+fn non_entry_a() -> void {
+}
 
-      mod->AddFunction(
-          Func("entry", ast::VariableList{}, ty.void_, ast::StatementList{},
-               ast::FunctionDecorationList{
-                   create<ast::StageDecoration>(ast::PipelineStage::kVertex),
-               }));
+[[stage(vertex)]]
+fn entry() -> void {
+}
 
-      mod->AddFunction(Func("non_entry_b", ast::VariableList{}, ty.void_,
-                            ast::StatementList{},
-                            ast::FunctionDecorationList{}));
-    }
-  };
-
-  auto result = GetTransform(Builder{}.Module());
-  ASSERT_FALSE(result.diagnostics.contains_errors())
-      << diag::Formatter().format(result.diagnostics);
-
-  auto* expected = R"(Module{
-  Variable{
-    Decorations{
-      BuiltinDecoration{pointsize}
-    }
-    tint_pointsize
-    out
-    __f32
-  }
-  Function non_entry_a -> __void
-  ()
-  {
-  }
-  Function entry -> __void
-  StageDecoration{vertex}
-  ()
-  {
-    Assignment{
-      Identifier[__ptr_out__f32]{tint_pointsize}
-      ScalarConstructor[__f32]{1.000000}
-    }
-  }
-  Function non_entry_b -> __void
-  ()
-  {
-  }
+fn non_entry_b() -> void {
 }
 )";
-  EXPECT_EQ(expected,
-            Demangler().Demangle(result.module, result.module.to_str()));
+
+  auto* expect = R"(
+[[builtin(pointsize)]] var<out> tint_pointsize : f32;
+
+fn non_entry_a() -> void {
+}
+
+[[stage(vertex)]]
+fn entry() -> void {
+  tint_pointsize = 1.0;
+}
+
+fn non_entry_b() -> void {
+}
+)";
+
+  auto got = Transform<EmitVertexPointSize>(src);
+
+  EXPECT_EQ(expect, got);
 }
 
 TEST_F(EmitVertexPointSizeTest, NonVertexStage) {
-  struct Builder : ModuleBuilder {
-    void Build() override {
-      auto* fragment_entry = Func(
-          "fragment_entry", ast::VariableList{}, ty.void_, ast::StatementList{},
-          ast::FunctionDecorationList{
-              create<ast::StageDecoration>(ast::PipelineStage::kFragment),
-          });
-      mod->AddFunction(fragment_entry);
+  auto* src = R"(
+[[stage(fragment)]]
+fn fragment_entry() -> void {
+}
 
-      auto* compute_entry = Func(
-          "compute_entry", ast::VariableList{}, ty.void_, ast::StatementList{},
-          ast::FunctionDecorationList{
-              create<ast::StageDecoration>(ast::PipelineStage::kCompute),
-          });
-      mod->AddFunction(compute_entry);
-    }
-  };
-
-  auto result = GetTransform(Builder{}.Module());
-  ASSERT_FALSE(result.diagnostics.contains_errors())
-      << diag::Formatter().format(result.diagnostics);
-
-  auto* expected = R"(Module{
-  Function fragment_entry -> __void
-  StageDecoration{fragment}
-  ()
-  {
-  }
-  Function compute_entry -> __void
-  StageDecoration{compute}
-  ()
-  {
-  }
+[[stage(compute)]]
+fn compute_entry() -> void {
 }
 )";
-  EXPECT_EQ(expected,
-            Demangler().Demangle(result.module, result.module.to_str()));
+
+  auto* expect = R"(
+[[stage(fragment)]]
+fn fragment_entry() -> void {
+}
+
+[[stage(compute)]]
+fn compute_entry() -> void {
+}
+)";
+
+  auto got = Transform<EmitVertexPointSize>(src);
+
+  EXPECT_EQ(expect, got);
 }
 
 }  // namespace

src/transform/first_index_offset_test.cc

@@ -15,413 +15,223 @@
 #include "src/transform/first_index_offset.h"
 
 #include <memory>
-#include <string>
 #include <utility>
+#include <vector>
 
-#include "gtest/gtest.h"
-#include "src/ast/block_statement.h"
-#include "src/ast/builder.h"
-#include "src/ast/builtin.h"
-#include "src/ast/builtin_decoration.h"
-#include "src/ast/call_expression.h"
-#include "src/ast/call_statement.h"
-#include "src/ast/function.h"
-#include "src/ast/identifier_expression.h"
-#include "src/ast/module.h"
-#include "src/ast/return_statement.h"
-#include "src/ast/storage_class.h"
-#include "src/ast/type/u32_type.h"
-#include "src/ast/variable.h"
-#include "src/ast/variable_decoration.h"
-#include "src/demangler.h"
-#include "src/diagnostic/formatter.h"
-#include "src/source.h"
-#include "src/transform/manager.h"
+#include "src/transform/test_helper.h"
 
 namespace tint {
 namespace transform {
 namespace {
 
-class FirstIndexOffsetTest : public testing::Test {};
-
-struct ModuleBuilder : public ast::BuilderWithModule {
-  ast::Module Module() {
-    Build();
-    return std::move(*mod);
-  }
-
- protected:
-  void AddBuiltinInput(const std::string& name, ast::Builtin builtin) {
-    mod->AddGlobalVariable(Var(name, ast::StorageClass::kInput, ty.u32, nullptr,
-                               {create<ast::BuiltinDecoration>(builtin)}));
-  }
-
-  ast::Function* AddFunction(const std::string& name,
-                             ast::StatementList stmts) {
-    auto* func = Func(name, ast::VariableList{}, ty.u32, stmts,
-                      ast::FunctionDecorationList{});
-    mod->AddFunction(func);
-    return func;
-  }
-
-  virtual void Build() = 0;
-};
+using FirstIndexOffsetTest = TransformTest;
 
 TEST_F(FirstIndexOffsetTest, Error_AlreadyTransformed) {
-  struct Builder : public ModuleBuilder {
-    void Build() override {
-      AddBuiltinInput("vert_idx", ast::Builtin::kVertexIdx);
-      AddFunction("test", {create<ast::ReturnStatement>(Expr("vert_idx"))});
-    }
-  };
+  auto* src = R"(
+[[builtin(vertex_idx)]] var<in> vert_idx : u32;
 
-  Manager manager;
-  manager.append(std::make_unique<FirstIndexOffset>(0, 0));
-  manager.append(std::make_unique<FirstIndexOffset>(1, 1));
+fn test() -> u32 {
+  return vert_idx;
+}
 
-  auto module = Builder{}.Module();
-  auto result = manager.Run(&module);
+[[stage(vertex)]]
+fn entry() -> void {
+  test();
+}
+)";
 
-  // Release the source module to ensure there's no uncloned data in result
-  { auto tmp = std::move(module); }
+  auto* expect = R"(manager().Run() errored:
+error: First index offset transform has already been applied.)";
 
-  ASSERT_EQ(diag::Formatter().format(result.diagnostics),
-            "error: First index offset transform has already been applied.");
+  std::vector<std::unique_ptr<transform::Transform>> transforms;
+  transforms.emplace_back(std::make_unique<FirstIndexOffset>(0, 0));
+  transforms.emplace_back(std::make_unique<FirstIndexOffset>(1, 1));
+
+  auto got = Transform(src, std::move(transforms));
+
+  EXPECT_EQ(expect, got);
 }
 
 TEST_F(FirstIndexOffsetTest, EmptyModule) {
-  Manager manager;
-  manager.append(std::make_unique<FirstIndexOffset>(0, 0));
+  auto* src = "";
+  auto* expect = "";
 
-  ast::Module module;
-  auto result = manager.Run(&module);
+  auto got = Transform<FirstIndexOffset>(src, 0, 0);
 
-  // Release the source module to ensure there's no uncloned data in result
-  { auto tmp = std::move(module); }
-
-  ASSERT_FALSE(result.diagnostics.contains_errors())
-      << diag::Formatter().format(result.diagnostics);
-
-  auto got = result.module.to_str();
-  auto* expected = "Module{\n}\n";
-  EXPECT_EQ(got, expected);
+  EXPECT_EQ(expect, got);
 }
 
 TEST_F(FirstIndexOffsetTest, BasicModuleVertexIndex) {
-  struct Builder : public ModuleBuilder {
-    void Build() override {
-      AddBuiltinInput("vert_idx", ast::Builtin::kVertexIdx);
-      AddFunction("test", {create<ast::ReturnStatement>(Expr("vert_idx"))});
-    }
-  };
+  auto* src = R"(
+[[builtin(vertex_idx)]] var<in> vert_idx : u32;
 
-  Manager manager;
-  manager.append(std::make_unique<FirstIndexOffset>(1, 2));
+fn test() -> u32 {
+  return vert_idx;
+}
 
-  auto module = Builder{}.Module();
-  auto result = manager.Run(&module);
-
-  // Release the source module to ensure there's no uncloned data in result
-  { auto tmp = std::move(module); }
-
-  ASSERT_FALSE(result.diagnostics.contains_errors())
-      << diag::Formatter().format(result.diagnostics);
-
-  auto got = result.module.to_str();
-  auto* expected =
-      R"(Module{
-  TintFirstIndexOffsetData Struct{
-    [[block]]
-    StructMember{[[ offset 0 ]] tint_first_vertex_index: __u32}
-  }
-  Variable{
-    Decorations{
-      BuiltinDecoration{vertex_idx}
-    }
-    tint_first_index_offset_vert_idx
-    in
-    __u32
-  }
-  Variable{
-    Decorations{
-      BindingDecoration{1}
-      SetDecoration{2}
-    }
-    tint_first_index_data
-    uniform
-    __struct_TintFirstIndexOffsetData
-  }
-  Function test -> __u32
-  ()
-  {
-    VariableDeclStatement{
-      VariableConst{
-        vert_idx
-        none
-        __u32
-        {
-          Binary[__u32]{
-            Identifier[__ptr_in__u32]{tint_first_index_offset_vert_idx}
-            add
-            MemberAccessor[__ptr_uniform__u32]{
-              Identifier[__ptr_uniform__struct_TintFirstIndexOffsetData]{tint_first_index_data}
-              Identifier[not set]{tint_first_vertex_index}
-            }
-          }
-        }
-      }
-    }
-    Return{
-      {
-        Identifier[__u32]{vert_idx}
-      }
-    }
-  }
+[[stage(vertex)]]
+fn entry() -> void {
+  test();
 }
 )";
-  EXPECT_EQ(Demangler().Demangle(result.module, got), expected);
+
+  auto* expect = R"(
+[[block]]
+struct TintFirstIndexOffsetData {
+  [[offset(0)]]
+  tint_first_vertex_index : u32;
+};
+
+[[builtin(vertex_idx)]] var<in> tint_first_index_offset_vert_idx : u32;
+[[binding(1), set(2)]] var<uniform> tint_first_index_data : TintFirstIndexOffsetData;
+
+fn test() -> u32 {
+  const vert_idx : u32 = (tint_first_index_offset_vert_idx + tint_first_index_data.tint_first_vertex_index);
+  return vert_idx;
+}
+
+[[stage(vertex)]]
+fn entry() -> void {
+  test();
+}
+)";
+
+  auto got = Transform<FirstIndexOffset>(src, 1, 2);
+
+  EXPECT_EQ(expect, got);
 }
 
 TEST_F(FirstIndexOffsetTest, BasicModuleInstanceIndex) {
-  struct Builder : public ModuleBuilder {
-    void Build() override {
-      AddBuiltinInput("inst_idx", ast::Builtin::kInstanceIdx);
-      AddFunction("test", {create<ast::ReturnStatement>(Expr("inst_idx"))});
-    }
-  };
+  auto* src = R"(
+[[builtin(instance_idx)]] var<in> inst_idx : u32;
 
-  Manager manager;
-  manager.append(std::make_unique<FirstIndexOffset>(1, 7));
+fn test() -> u32 {
+  return inst_idx;
+}
 
-  auto module = Builder{}.Module();
-  auto result = manager.Run(&module);
-
-  // Release the source module to ensure there's no uncloned data in result
-  { auto tmp = std::move(module); }
-
-  ASSERT_FALSE(result.diagnostics.contains_errors())
-      << diag::Formatter().format(result.diagnostics);
-
-  auto got = result.module.to_str();
-  auto* expected = R"(Module{
-  TintFirstIndexOffsetData Struct{
-    [[block]]
-    StructMember{[[ offset 0 ]] tint_first_instance_index: __u32}
-  }
-  Variable{
-    Decorations{
-      BuiltinDecoration{instance_idx}
-    }
-    tint_first_index_offset_inst_idx
-    in
-    __u32
-  }
-  Variable{
-    Decorations{
-      BindingDecoration{1}
-      SetDecoration{7}
-    }
-    tint_first_index_data
-    uniform
-    __struct_TintFirstIndexOffsetData
-  }
-  Function test -> __u32
-  ()
-  {
-    VariableDeclStatement{
-      VariableConst{
-        inst_idx
-        none
-        __u32
-        {
-          Binary[__u32]{
-            Identifier[__ptr_in__u32]{tint_first_index_offset_inst_idx}
-            add
-            MemberAccessor[__ptr_uniform__u32]{
-              Identifier[__ptr_uniform__struct_TintFirstIndexOffsetData]{tint_first_index_data}
-              Identifier[not set]{tint_first_instance_index}
-            }
-          }
-        }
-      }
-    }
-    Return{
-      {
-        Identifier[__u32]{inst_idx}
-      }
-    }
-  }
+[[stage(vertex)]]
+fn entry() -> void {
+  test();
 }
 )";
-  EXPECT_EQ(Demangler().Demangle(result.module, got), expected);
+
+  auto* expect = R"(
+[[block]]
+struct TintFirstIndexOffsetData {
+  [[offset(0)]]
+  tint_first_instance_index : u32;
+};
+
+[[builtin(instance_idx)]] var<in> tint_first_index_offset_inst_idx : u32;
+[[binding(1), set(7)]] var<uniform> tint_first_index_data : TintFirstIndexOffsetData;
+
+fn test() -> u32 {
+  const inst_idx : u32 = (tint_first_index_offset_inst_idx + tint_first_index_data.tint_first_instance_index);
+  return inst_idx;
+}
+
+[[stage(vertex)]]
+fn entry() -> void {
+  test();
+}
+)";
+
+  auto got = Transform<FirstIndexOffset>(src, 1, 7);
+
+  EXPECT_EQ(expect, got);
 }
 
 TEST_F(FirstIndexOffsetTest, BasicModuleBothIndex) {
-  struct Builder : public ModuleBuilder {
-    void Build() override {
-      AddBuiltinInput("inst_idx", ast::Builtin::kInstanceIdx);
-      AddBuiltinInput("vert_idx", ast::Builtin::kVertexIdx);
-      AddFunction("test", {
-                              create<ast::ReturnStatement>(Expr(1u)),
-                          });
-    }
-  };
+  auto* src = R"(
+[[builtin(instance_idx)]] var<in> instance_idx : u32;
+[[builtin(vertex_idx)]] var<in> vert_idx : u32;
 
-  auto transform = std::make_unique<FirstIndexOffset>(1, 7);
-  auto* transform_ptr = transform.get();
+fn test() -> u32 {
+  return instance_idx + vert_idx;
+}
 
-  Manager manager;
-  manager.append(std::move(transform));
-
-  auto module = Builder{}.Module();
-  auto result = manager.Run(&module);
-
-  // Release the source module to ensure there's no uncloned data in result
-  { auto tmp = std::move(module); }
-
-  ASSERT_FALSE(result.diagnostics.contains_errors())
-      << diag::Formatter().format(result.diagnostics);
-
-  auto got = result.module.to_str();
-  auto* expected = R"(Module{
-  TintFirstIndexOffsetData Struct{
-    [[block]]
-    StructMember{[[ offset 0 ]] tint_first_vertex_index: __u32}
-    StructMember{[[ offset 4 ]] tint_first_instance_index: __u32}
-  }
-  Variable{
-    Decorations{
-      BuiltinDecoration{instance_idx}
-    }
-    tint_first_index_offset_inst_idx
-    in
-    __u32
-  }
-  Variable{
-    Decorations{
-      BuiltinDecoration{vertex_idx}
-    }
-    tint_first_index_offset_vert_idx
-    in
-    __u32
-  }
-  Variable{
-    Decorations{
-      BindingDecoration{1}
-      SetDecoration{7}
-    }
-    tint_first_index_data
-    uniform
-    __struct_TintFirstIndexOffsetData
-  }
-  Function test -> __u32
-  ()
-  {
-    Return{
-      {
-        ScalarConstructor[__u32]{1}
-      }
-    }
-  }
+[[stage(vertex)]]
+fn entry() -> void {
+  test();
 }
 )";
-  EXPECT_EQ(Demangler().Demangle(result.module, got), expected);
 
-  EXPECT_TRUE(transform_ptr->HasVertexIndex());
-  EXPECT_EQ(transform_ptr->GetFirstVertexOffset(), 0u);
+  auto* expect = R"(
+[[block]]
+struct TintFirstIndexOffsetData {
+  [[offset(0)]]
+  tint_first_vertex_index : u32;
+  [[offset(4)]]
+  tint_first_instance_index : u32;
+};
 
-  EXPECT_TRUE(transform_ptr->HasInstanceIndex());
-  EXPECT_EQ(transform_ptr->GetFirstInstanceOffset(), 4u);
+[[builtin(instance_idx)]] var<in> tint_first_index_offset_instance_idx : u32;
+[[builtin(vertex_idx)]] var<in> tint_first_index_offset_vert_idx : u32;
+[[binding(1), set(2)]] var<uniform> tint_first_index_data : TintFirstIndexOffsetData;
+
+fn test() -> u32 {
+  const instance_idx : u32 = (tint_first_index_offset_instance_idx + tint_first_index_data.tint_first_instance_index);
+  const vert_idx : u32 = (tint_first_index_offset_vert_idx + tint_first_index_data.tint_first_vertex_index);
+  return (instance_idx + vert_idx);
+}
+
+[[stage(vertex)]]
+fn entry() -> void {
+  test();
+}
+)";
+
+  auto got = Transform<FirstIndexOffset>(src, 1, 2);
+
+  EXPECT_EQ(expect, got);
 }
 
 TEST_F(FirstIndexOffsetTest, NestedCalls) {
-  struct Builder : public ModuleBuilder {
-    void Build() override {
-      AddBuiltinInput("vert_idx", ast::Builtin::kVertexIdx);
-      AddFunction("func1", {create<ast::ReturnStatement>(Expr("vert_idx"))});
-      AddFunction("func2", {create<ast::ReturnStatement>(Call("func1"))});
-    }
-  };
+  auto* src = R"(
+[[builtin(vertex_idx)]] var<in> vert_idx : u32;
 
-  auto transform = std::make_unique<FirstIndexOffset>(2, 2);
+fn func1() -> u32 {
+  return vert_idx;
+}
 
-  Manager manager;
-  manager.append(std::move(transform));
+fn func2() -> u32 {
+  return func1();
+}
 
-  auto module = Builder{}.Module();
-  auto result = manager.Run(&module);
-
-  // Release the source module to ensure there's no uncloned data in result
-  { auto tmp = std::move(module); }
-
-  ASSERT_FALSE(result.diagnostics.contains_errors())
-      << diag::Formatter().format(result.diagnostics);
-
-  auto got = result.module.to_str();
-  auto* expected = R"(Module{
-  TintFirstIndexOffsetData Struct{
-    [[block]]
-    StructMember{[[ offset 0 ]] tint_first_vertex_index: __u32}
-  }
-  Variable{
-    Decorations{
-      BuiltinDecoration{vertex_idx}
-    }
-    tint_first_index_offset_vert_idx
-    in
-    __u32
-  }
-  Variable{
-    Decorations{
-      BindingDecoration{2}
-      SetDecoration{2}
-    }
-    tint_first_index_data
-    uniform
-    __struct_TintFirstIndexOffsetData
-  }
-  Function func1 -> __u32
-  ()
-  {
-    VariableDeclStatement{
-      VariableConst{
-        vert_idx
-        none
-        __u32
-        {
-          Binary[__u32]{
-            Identifier[__ptr_in__u32]{tint_first_index_offset_vert_idx}
-            add
-            MemberAccessor[__ptr_uniform__u32]{
-              Identifier[__ptr_uniform__struct_TintFirstIndexOffsetData]{tint_first_index_data}
-              Identifier[not set]{tint_first_vertex_index}
-            }
-          }
-        }
-      }
-    }
-    Return{
-      {
-        Identifier[__u32]{vert_idx}
-      }
-    }
-  }
-  Function func2 -> __u32
-  ()
-  {
-    Return{
-      {
-        Call[__u32]{
-          Identifier[__u32]{func1}
-          (
-          )
-        }
-      }
-    }
-  }
+[[stage(vertex)]]
+fn entry() -> void {
+  func2();
 }
 )";
-  EXPECT_EQ(Demangler().Demangle(result.module, got), expected);
+
+  auto* expect = R"(
+[[block]]
+struct TintFirstIndexOffsetData {
+  [[offset(0)]]
+  tint_first_vertex_index : u32;
+};
+
+[[builtin(vertex_idx)]] var<in> tint_first_index_offset_vert_idx : u32;
+[[binding(1), set(2)]] var<uniform> tint_first_index_data : TintFirstIndexOffsetData;
+
+fn func1() -> u32 {
+  const vert_idx : u32 = (tint_first_index_offset_vert_idx + tint_first_index_data.tint_first_vertex_index);
+  return vert_idx;
+}
+
+fn func2() -> u32 {
+  return func1();
+}
+
+[[stage(vertex)]]
+fn entry() -> void {
+  func2();
+}
+)";
+
+  auto got = Transform<FirstIndexOffset>(src, 1, 2);
+
+  EXPECT_EQ(expect, got);
 }
 
 }  // namespace

src/transform/test_helper.h

@@ -0,0 +1,115 @@
+// Copyright 2021 The Tint Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef SRC_TRANSFORM_TEST_HELPER_H_
+#define SRC_TRANSFORM_TEST_HELPER_H_
+
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "gtest/gtest.h"
+#include "src/reader/wgsl/parser.h"
+#include "src/transform/manager.h"
+#include "src/type_determiner.h"
+#include "src/writer/wgsl/generator.h"
+
+namespace tint {
+namespace transform {
+
+/// Helper class for testing transforms
+class TransformTest : public testing::Test {
+ public:
+  /// Transforms and returns the WGSL source `in`, transformed using
+  /// `transforms`.
+  /// @param in the input WGSL source
+  /// @param transforms the list of transforms to apply
+  /// @return the transformed WGSL output
+  std::string Transform(
+      std::string in,
+      std::vector<std::unique_ptr<transform::Transform>> transforms) {
+    Source::File file("test", in);
+    reader::wgsl::Parser parser(&file);
+    if (!parser.Parse()) {
+      return "WGSL reader failed:\n" + parser.error();
+    }
+
+    auto module = parser.module();
+    TypeDeterminer td(&module);
+    if (!td.Determine()) {
+      return "Type determination failed:\n" + td.error();
+    }
+
+    Manager manager;
+    for (auto& transform : transforms) {
+      manager.append(std::move(transform));
+    }
+    auto result = manager.Run(&module);
+
+    if (result.diagnostics.contains_errors()) {
+      return "manager().Run() errored:\n" +
+             diag::Formatter().format(result.diagnostics);
+    }
+
+    // Release the source module to ensure there's no uncloned data in result
+    { auto tmp = std::move(module); }
+
+    writer::wgsl::Generator generator(std::move(result.module));
+    if (!generator.Generate()) {
+      return "WGSL writer failed:\n" + generator.error();
+    }
+
+    auto res = generator.result();
+    if (res.empty()) {
+      return res;
+    }
+    // The WGSL sometimes has two trailing newlines. Strip them
+    while (res.back() == '\n') {
+      res.pop_back();
+    }
+    if (res.empty()) {
+      return res;
+    }
+    return "\n" + res + "\n";
+  }
+
+  /// Transforms and returns the WGSL source `in`, transformed using
+  /// `transform`.
+  /// @param transform the transform to apply
+  /// @param in the input WGSL source
+  /// @return the transformed WGSL output
+  std::string Transform(std::string in,
+                        std::unique_ptr<transform::Transform> transform) {
+    std::vector<std::unique_ptr<transform::Transform>> transforms;
+    transforms.emplace_back(std::move(transform));
+    return Transform(std::move(in), std::move(transforms));
+  }
+
+  /// Transforms and returns the WGSL source `in`, transformed using
+  /// a transform of type `TRANSFORM`.
+  /// @param in the input WGSL source
+  /// @param args the TRANSFORM constructor arguments
+  /// @return the transformed WGSL output
+  template <typename TRANSFORM, typename... ARGS>
+  std::string Transform(std::string in, ARGS&&... args) {
+    return Transform(std::move(in),
+                     std::make_unique<TRANSFORM>(std::forward<ARGS>(args)...));
+  }
+};
+
+}  // namespace transform
+}  // namespace tint
+
+#endif  // SRC_TRANSFORM_TEST_HELPER_H_