transform: Add Data as an input

Like 42264, this is an attempt to keep transforms immutable, and to keep mutable I/O information in Data objects.

Bug: tint:389
Change-Id: Ie417872d3e97f3db3904245a30aa74fdcce76610
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/46260
Commit-Queue: Ben Clayton <bclayton@google.com>
Reviewed-by: James Price <jrprice@google.com>

src/transform/bound_array_accessors.cc

@@ -26,7 +26,7 @@ namespace transform {
 BoundArrayAccessors::BoundArrayAccessors() = default;
 BoundArrayAccessors::~BoundArrayAccessors() = default;
 
-Transform::Output BoundArrayAccessors::Run(const Program* in) {
+Transform::Output BoundArrayAccessors::Run(const Program* in, const DataMap&) {
   ProgramBuilder out;
   CloneContext ctx(&out, in);
   ctx.ReplaceAll([&](ast::ArrayAccessorExpression* expr) {

src/transform/bound_array_accessors.h

@@ -34,8 +34,9 @@ class BoundArrayAccessors : public Transform {
 
   /// Runs the transform on `program`, returning the transformation result.
   /// @param program the source program to transform
+  /// @param data optional extra transform-specific input data
   /// @returns the transformation result
-  Output Run(const Program* program) override;
+  Output Run(const Program* program, const DataMap& data = {}) override;
 
  private:
   ast::ArrayAccessorExpression* Transform(ast::ArrayAccessorExpression* expr,

src/transform/emit_vertex_point_size.cc

@@ -30,7 +30,7 @@ const char kPointSizeVar[] = "tint_pointsize";
 EmitVertexPointSize::EmitVertexPointSize() = default;
 EmitVertexPointSize::~EmitVertexPointSize() = default;
 
-Transform::Output EmitVertexPointSize::Run(const Program* in) {
+Transform::Output EmitVertexPointSize::Run(const Program* in, const DataMap&) {
   if (!in->AST().Functions().HasStage(ast::PipelineStage::kVertex)) {
     // If the module doesn't have any vertex stages, then there's nothing to do.
     return Output(Program(in->Clone()));

src/transform/emit_vertex_point_size.h

@@ -34,8 +34,9 @@ class EmitVertexPointSize : public Transform {
 
   /// Runs the transform on `program`, returning the transformation result.
   /// @param program the source program to transform
+  /// @param data optional extra transform-specific input data
   /// @returns the transformation result
-  Output Run(const Program* program) override;
+  Output Run(const Program* program, const DataMap& data = {}) override;
 };
 
 }  // namespace transform

src/transform/first_index_offset.cc

@@ -67,7 +67,7 @@ FirstIndexOffset::FirstIndexOffset(uint32_t binding, uint32_t group)
 
 FirstIndexOffset::~FirstIndexOffset() = default;
 
-Transform::Output FirstIndexOffset::Run(const Program* in) {
+Transform::Output FirstIndexOffset::Run(const Program* in, const DataMap&) {
   ProgramBuilder out;
 
   // First do a quick check to see if the transform has already been applied.
@@ -191,7 +191,6 @@ ast::Variable* FirstIndexOffset::State::AddUniformBuffer() {
           dst->create<ast::GroupDecoration>(Source{}, group),
       });
 
-
   dst->AST().AddConstructedType(struct_type);
 
   dst->AST().AddGlobalVariable(idx_var);

src/transform/first_index_offset.h

@@ -60,6 +60,7 @@ namespace transform {
 ///
 class FirstIndexOffset : public Transform {
  public:
+  /// Data is outputted by the FirstIndexOffset transform.
   /// Data holds information about shader usage and constant buffer offsets.
   struct Data : public Castable<Data, transform::Data> {
     /// Constructor
@@ -96,8 +97,9 @@ class FirstIndexOffset : public Transform {
 
   /// Runs the transform on `program`, returning the transformation result.
   /// @param program the source program to transform
+  /// @param data optional extra transform-specific input data
   /// @returns the transformation result
-  Output Run(const Program* program) override;
+  Output Run(const Program* program, const DataMap& data = {}) override;
 
  private:
   struct State {

src/transform/first_index_offset_test.cc

@@ -64,7 +64,7 @@ TEST_F(FirstIndexOffsetTest, EmptyModule) {
   auto* src = "";
   auto* expect = "";
 
-  auto got = Run<FirstIndexOffset>(src, 0, 0);
+  auto got = Run(src, std::make_unique<FirstIndexOffset>(0, 0));
 
   EXPECT_EQ(expect, str(got));
 
@@ -112,7 +112,7 @@ fn entry() -> void {
 }
 )";
 
-  auto got = Run<FirstIndexOffset>(src, 1, 2);
+  auto got = Run(src, std::make_unique<FirstIndexOffset>(1, 2));
 
   EXPECT_EQ(expect, str(got));
 
@@ -160,7 +160,7 @@ fn entry() -> void {
 }
 )";
 
-  auto got = Run<FirstIndexOffset>(src, 1, 7);
+  auto got = Run(src, std::make_unique<FirstIndexOffset>(1, 7));
 
   EXPECT_EQ(expect, str(got));
 
@@ -213,7 +213,7 @@ fn entry() -> void {
 }
 )";
 
-  auto got = Run<FirstIndexOffset>(src, 1, 2);
+  auto got = Run(src, std::make_unique<FirstIndexOffset>(1, 2));
 
   EXPECT_EQ(expect, str(got));
 
@@ -269,7 +269,7 @@ fn entry() -> void {
 }
 )";
 
-  auto got = Run<FirstIndexOffset>(src, 1, 2);
+  auto got = Run(src, std::make_unique<FirstIndexOffset>(1, 2));
 
   EXPECT_EQ(expect, str(got));
 

src/transform/hlsl.cc

@@ -29,7 +29,7 @@ namespace transform {
 Hlsl::Hlsl() = default;
 Hlsl::~Hlsl() = default;
 
-Transform::Output Hlsl::Run(const Program* in) {
+Transform::Output Hlsl::Run(const Program* in, const DataMap&) {
   ProgramBuilder out;
   CloneContext ctx(&out, in);
   PromoteArrayInitializerToConstVar(ctx);

src/transform/hlsl.h

@@ -35,8 +35,9 @@ class Hlsl : public Transform {
 
   /// Runs the transform on `program`, returning the transformation result.
   /// @param program the source program to transform
+  /// @param data optional extra transform-specific data
   /// @returns the transformation result
-  Output Run(const Program* program) override;
+  Output Run(const Program* program, const DataMap& data = {}) override;
 
  private:
   /// Hoists the array initializer to a constant variable, declared just before

src/transform/manager.cc

@@ -20,11 +20,11 @@ namespace transform {
 Manager::Manager() = default;
 Manager::~Manager() = default;
 
-Transform::Output Manager::Run(const Program* program) {
+Transform::Output Manager::Run(const Program* program, const DataMap& data) {
   Output out;
   if (!transforms_.empty()) {
     for (auto& transform : transforms_) {
-      auto res = transform->Run(program);
+      auto res = transform->Run(program, data);
       out.program = std::move(res.program);
       out.data.Add(std::move(res.data));
       if (!out.program.IsValid()) {

src/transform/manager.h

@@ -42,8 +42,9 @@ class Manager : public Transform {
 
   /// Runs the transforms on `program`, returning the transformation result.
   /// @param program the source program to transform
+  /// @param data optional extra transform-specific input data
   /// @returns the transformed program and diagnostics
-  Output Run(const Program* program) override;
+  Output Run(const Program* program, const DataMap& data = {}) override;
 
  private:
   std::vector<std::unique_ptr<Transform>> transforms_;

src/transform/msl.cc

@@ -266,7 +266,7 @@ const char* kReservedKeywords[] = {"access",
 Msl::Msl() = default;
 Msl::~Msl() = default;
 
-Transform::Output Msl::Run(const Program* in) {
+Transform::Output Msl::Run(const Program* in, const DataMap&) {
   ProgramBuilder out;
   CloneContext ctx(&out, in);
   RenameReservedKeywords(&ctx, kReservedKeywords);

src/transform/msl.h

@@ -31,8 +31,9 @@ class Msl : public Transform {
 
   /// Runs the transform on `program`, returning the transformation result.
   /// @param program the source program to transform
+  /// @param data optional extra transform-specific input data
   /// @returns the transformation result
-  Output Run(const Program* program) override;
+  Output Run(const Program* program, const DataMap& data = {}) override;
 
  private:
   /// Hoist location-decorated entry point parameters out to struct members.

src/transform/renamer.cc

@@ -39,7 +39,7 @@ Renamer::Renamer(const Config& config) : cfg_(config) {}
 
 Renamer::~Renamer() = default;
 
-Transform::Output Renamer::Run(const Program* in) {
+Transform::Output Renamer::Run(const Program* in, const DataMap&) {
   ProgramBuilder out;
   CloneContext ctx(&out, in);
 

src/transform/renamer.h

@@ -26,6 +26,7 @@ namespace transform {
 /// Renamer is a Transform that renames all the symbols in a program.
 class Renamer : public Transform {
  public:
+  /// Data is outputted by the Renamer transform.
   /// Data holds information about shader usage and constant buffer offsets.
   struct Data : public Castable<Data, transform::Data> {
     /// Remappings is a map of old symbol name to new symbol name
@@ -70,8 +71,9 @@ class Renamer : public Transform {
 
   /// Runs the transform on `program`, returning the transformation result.
   /// @param program the source program to transform
+  /// @param data optional extra transform-specific input data
   /// @returns the transformation result
-  Output Run(const Program* program) override;
+  Output Run(const Program* program, const DataMap& data = {}) override;
 
  private:
   Config const cfg_;

src/transform/spirv.cc

@@ -29,7 +29,7 @@ namespace transform {
 Spirv::Spirv() = default;
 Spirv::~Spirv() = default;
 
-Transform::Output Spirv::Run(const Program* in) {
+Transform::Output Spirv::Run(const Program* in, const DataMap&) {
   ProgramBuilder out;
   CloneContext ctx(&out, in);
   HandleEntryPointIOTypes(ctx);

src/transform/spirv.h

@@ -37,8 +37,9 @@ class Spirv : public Transform {
 
   /// Runs the transform on `program`, returning the transformation result.
   /// @param program the source program to transform
+  /// @param data optional extra transform-specific input data
   /// @returns the transformation result
-  Output Run(const Program* program) override;
+  Output Run(const Program* program, const DataMap& data = {}) override;
 
  private:
   /// Hoist entry point parameters, return values, and struct members out to

src/transform/test_helper.h

@@ -36,10 +36,12 @@ class TransformTestBase : public BASE {
   /// `transforms`.
   /// @param in the input WGSL source
   /// @param transforms the list of transforms to apply
+  /// @param data the optional DataMap to pass to Transform::Run()
   /// @return the transformed output
   Transform::Output Run(
       std::string in,
-      std::vector<std::unique_ptr<transform::Transform>> transforms) {
+      std::vector<std::unique_ptr<transform::Transform>> transforms,
+      DataMap data = {}) {
     auto file = std::make_unique<Source::File>("test", in);
     auto program = reader::wgsl::Parse(file.get());
 
@@ -54,30 +56,31 @@ class TransformTestBase : public BASE {
     for (auto& transform : transforms) {
       manager.append(std::move(transform));
     }
-    return manager.Run(&program);
+    return manager.Run(&program, data);
   }
 
   /// Transforms and returns the WGSL source `in`, transformed using
   /// `transform`.
   /// @param transform the transform to apply
   /// @param in the input WGSL source
+  /// @param data the optional DataMap to pass to Transform::Run()
   /// @return the transformed output
   Transform::Output Run(std::string in,
-                        std::unique_ptr<transform::Transform> transform) {
+                        std::unique_ptr<transform::Transform> transform,
+                        DataMap data = {}) {
     std::vector<std::unique_ptr<transform::Transform>> transforms;
     transforms.emplace_back(std::move(transform));
-    return Run(std::move(in), std::move(transforms));
+    return Run(std::move(in), std::move(transforms), std::move(data));
   }
 
   /// 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
+  /// @param data the optional DataMap to pass to Transform::Run()
   /// @return the transformed output
-  template <typename TRANSFORM, typename... ARGS>
-  Transform::Output Run(std::string in, ARGS&&... args) {
-    return Run(std::move(in),
-               std::make_unique<TRANSFORM>(std::forward<ARGS>(args)...));
+  template <typename TRANSFORM>
+  Transform::Output Run(std::string in, DataMap data = {}) {
+    return Run(std::move(in), std::make_unique<TRANSFORM>(), std::move(data));
   }
 
   /// @param output the output of the transform

src/transform/transform.h

@@ -24,8 +24,8 @@
 namespace tint {
 namespace transform {
 
-/// Data is the base class for transforms that emit extra output information
-/// along with a Program.
+/// Data is the base class for transforms that accept extra input or emit extra
+/// output information along with a Program.
 class Data : public Castable<Data> {
  public:
   /// Constructor
@@ -67,6 +67,14 @@ class DataMap {
     map_[&TypeInfo::Of<T>()] = std::move(data);
   }
 
+  /// Creates the data of type `T` with the provided arguments and adds it into
+  /// DataMap keyed by the ClassID of type T.
+  /// @param args the arguments forwarded to the constructor for type T
+  template <typename T, typename... ARGS>
+  void Add(ARGS&&... args) {
+    Put(std::make_unique<T>(std::forward<ARGS>(args)...));
+  }
+
   /// @returns a pointer to the Data placed into the DataMap with a call to
   /// Put()
   template <typename T>
@@ -137,8 +145,9 @@ class Transform {
 
   /// Runs the transform on `program`, returning the transformation result.
   /// @param program the source program to transform
+  /// @param data optional extra transform-specific input data
   /// @returns the transformation result
-  virtual Output Run(const Program* program) = 0;
+  virtual Output Run(const Program* program, const DataMap& data = {}) = 0;
 
  protected:
   /// Clones the function `in` adding `statements` to the beginning of the

src/transform/vertex_pulling.cc

@@ -40,7 +40,7 @@ VertexPulling::VertexPulling(const Config& config) : cfg(config) {}
 
 VertexPulling::~VertexPulling() = default;
 
-Transform::Output VertexPulling::Run(const Program* in) {
+Transform::Output VertexPulling::Run(const Program* in, const DataMap&) {
   ProgramBuilder out;
 
   // Find entry point

src/transform/vertex_pulling.h

@@ -162,8 +162,9 @@ class VertexPulling : public Transform {
 
   /// Runs the transform on `program`, returning the transformation result.
   /// @param program the source program to transform
+  /// @param data optional extra transform-specific input data
   /// @returns the transformation result
-  Output Run(const Program* program) override;
+  Output Run(const Program* program, const DataMap& data = {}) override;
 
  private:
   Config cfg;