ast: Clone symbols instead of ever-growing

The interface for cloning a module was made significantly more complex in https://dawn-review.googlesource.com/c/tint/+/35502/ as some of the transforms required constructing symbols before the clone. This was temporary solution in 35502 was to copy the symbol table, then construct the new types, then perform the clone. This lead to a really messy callback interface, that was extremely error prone (e.g. lamda-capturing stack variables from the initializer callback that had been unwound).

Instead, clone the symbols as they're encountered. This may produce an entirely different set of identifiers, but no longer ever-grows the symbol list, and keeps the interface clean.

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

src/ast/clone_context.cc

@@ -14,11 +14,22 @@
 
 #include "src/ast/clone_context.h"
 
+#include "src/ast/module.h"
+
 namespace tint {
 namespace ast {
 
-CloneContext::CloneContext(Module* m) : mod(m) {}
+CloneContext::CloneContext(Module* to, Module const* from)
+    : mod(to), src(from) {}
 CloneContext::~CloneContext() = default;
 
+Symbol CloneContext::Clone(const Symbol& s) const {
+  return mod->RegisterSymbol(src->SymbolToName(s));
+}
+
+void CloneContext::Clone() {
+  src->Clone(this);
+}
+
 }  // namespace ast
 }  // namespace tint

src/ast/clone_context.h

@@ -22,6 +22,7 @@
 #include "src/ast/traits.h"
 #include "src/castable.h"
 #include "src/source.h"
+#include "src/symbol.h"
 
 namespace tint {
 namespace ast {
@@ -32,19 +33,22 @@ class Module;
 class CloneContext {
  public:
   /// Constructor
-  /// @param m the target module to clone into
+  /// @param to the target module to clone into
-  explicit CloneContext(Module* m);
+  /// @param from the source module to clone from
+  CloneContext(Module* to, Module const* from);
 
   /// Destructor
   ~CloneContext();
 
-  /// Clones the `Node` or `type::Type` `a` into the module #mod if `a` is not
+  /// Clones the Node or type::Type `a` into the module #mod if `a` is not
   /// null. If `a` is null, then Clone() returns null. If `a` has been cloned
   /// already by this CloneContext then the same cloned pointer is returned.
   ///
   /// Clone() may use a function registered with ReplaceAll() to create a
   /// transformed version of the object. See ReplaceAll() for more information.
   ///
+  /// The Node or type::Type `a` must be owned by the module #src.
+  ///
   /// @note Semantic information such as resolved expression type and intrinsic
   /// information is not cloned.
   /// @param a the `Node` or `type::Type` to clone
@@ -70,15 +74,26 @@ class CloneContext {
     return static_cast<T*>(c);
   }
 
-  /// Clones the `Source` `s` into `mod`
+  /// Clones the Source `s` into `mod`
   /// TODO(bclayton) - Currently this 'clone' is a shallow copy. If/when
   /// `Source.File`s are owned by the `Module` this should make a copy of the
   /// file.
   /// @param s the `Source` to clone
   /// @return the cloned source
-  Source Clone(const Source& s) { return s; }
+  Source Clone(const Source& s) const { return s; }
+
+  /// Clones the Symbol `s` into `mod`
+  ///
+  /// The Symbol `s` must be owned by the module #src.
+  ///
+  /// @param s the Symbol to clone
+  /// @return the cloned source
+  Symbol Clone(const Symbol& s) const;
 
   /// Clones each of the elements of the vector `v` into the module #mod->
+  ///
+  /// All the elements of the vector `v` must be owned by the module #src.
+  ///
   /// @param v the vector to clone
   /// @return the cloned vector
   template <typename T>
@@ -92,11 +107,12 @@ class CloneContext {
   }
 
   /// ReplaceAll() registers `replacer` to be called whenever the Clone() method
-  /// is called with a type that matches (or derives from) the type of the first
+  /// is called with a type that matches (or derives from) the type of the
-  /// parameter of `replacer`.
+  /// second parameter of `replacer`.
   ///
-  /// `replacer` must be function-like with the signature: `T* (T*)`, where `T`
+  /// `replacer` must be function-like with the signature:
-  /// is a type deriving from CastableBase.
+  ///   `T* (CloneContext*, T*)`
+  ///  where `T` is a type deriving from CastableBase.
   ///
   /// If `replacer` returns a nullptr then Clone() will attempt the next
   /// registered replacer function that matches the object type. If no replacers
@@ -107,32 +123,43 @@ class CloneContext {
   ///
   /// ```
   ///   // Replace all ast::UintLiterals with the number 42
-  ///   CloneCtx ctx(mod);
+  ///   CloneCtx ctx(&out, in)
-  ///   ctx.ReplaceAll([&] (ast::UintLiteral* in) {
+  ///     .ReplaceAll([&] (CloneContext* ctx, ast::UintLiteral* l) {
-  ///     return ctx.mod->create<ast::UintLiteral>(ctx.Clone(in->source()),
+  ///       return ctx->mod->create<ast::UintLiteral>(ctx->Clone(l->source()),
-  ///                                              ctx.Clone(in->type()), 42);
+  ///                                                 ctx->Clone(l->type()),
-  ///   });
+  ///                                                 42);
-  ///   auto* out = ctx.Clone(tree);
+  ///     }).Clone();
   /// ```
   ///
   /// @param replacer a function or function-like object with the signature
-  ///        `T* (T*)`, where `T` derives from CastableBase
+  ///        `T* (CloneContext*, T*)`, where `T` derives from CastableBase
+  /// @returns this CloneContext so calls can be chained
   template <typename F>
-  void ReplaceAll(F replacer) {
+  CloneContext& ReplaceAll(F replacer) {
-    using TPtr = traits::ParamTypeT<F, 0>;
+    using TPtr = traits::ParamTypeT<F, 1>;
     using T = typename std::remove_pointer<TPtr>::type;
     transforms_.emplace_back([=](CastableBase* in) {
       auto* in_as_t = in->As<T>();
-      return in_as_t != nullptr ? replacer(in_as_t) : nullptr;
+      return in_as_t != nullptr ? replacer(this, in_as_t) : nullptr;
     });
+    return *this;
   }
 
+  /// Clone performs the clone of the entire module #src to #mod.
+  void Clone();
+
   /// The target module to clone into.
   Module* const mod;
 
+  /// The source module to clone from.
+  Module const* const src;
+
  private:
   using Transform = std::function<CastableBase*(CastableBase*)>;
 
+  CloneContext(const CloneContext&) = delete;
+  CloneContext& operator=(const CloneContext&) = delete;
+
   /// LookupOrTransform is the template-independent logic of Clone().
   /// This is outside of Clone() to reduce the amount of template-instantiated
   /// code.

src/ast/clone_context_test.cc

@@ -65,8 +65,7 @@ TEST(CloneContext, Clone) {
   // C: Clonable
 
   ast::Module cloned;
-  CloneContext ctx(&cloned);
+  auto* cloned_root = CloneContext(&cloned, &original).Clone(original_root);
-  auto* cloned_root = original_root->Clone(&ctx);
 
   EXPECT_NE(cloned_root->a, nullptr);
   EXPECT_EQ(cloned_root->a->a, nullptr);
@@ -110,14 +109,14 @@ TEST(CloneContext, CloneWithReplacements) {
   // R: Replaceable
 
   ast::Module cloned;
-  CloneContext ctx(&cloned);
+  auto* cloned_root = CloneContext(&cloned, &original)
-  ctx.ReplaceAll([&](Replaceable* in) {
+                          .ReplaceAll([&](CloneContext* ctx, Replaceable* in) {
-    auto* out = cloned.create<Replacement>();
+                            auto* out = cloned.create<Replacement>();
-    out->b = cloned.create<Cloneable>();
+                            out->b = cloned.create<Cloneable>();
-    out->c = ctx.Clone(in->a);
+                            out->c = ctx->Clone(in->a);
-    return out;
+                            return out;
-  });
+                          })
-  auto* cloned_root = original_root->Clone(&ctx);
+                          .Clone(original_root);
 
   //                         root
   //        ╭─────────────────┼──────────────────╮

src/ast/function.cc

@@ -228,7 +228,7 @@ const Statement* Function::get_last_statement() const {
 
 Function* Function::Clone(CloneContext* ctx) const {
   return ctx->mod->create<Function>(
-      ctx->Clone(source()), symbol_, name_, ctx->Clone(params_),
+      ctx->Clone(source()), ctx->Clone(symbol()), name_, ctx->Clone(params_),
       ctx->Clone(return_type_), ctx->Clone(body_), ctx->Clone(decorations_));
 }
 

src/ast/identifier_expression.cc

@@ -32,8 +32,8 @@ IdentifierExpression::IdentifierExpression(IdentifierExpression&&) = default;
 IdentifierExpression::~IdentifierExpression() = default;
 
 IdentifierExpression* IdentifierExpression::Clone(CloneContext* ctx) const {
-  return ctx->mod->create<IdentifierExpression>(ctx->Clone(source()), sym_,
+  return ctx->mod->create<IdentifierExpression>(ctx->Clone(source()),
-                                                name_);
+                                                ctx->Clone(symbol()), name_);
 }
 
 bool IdentifierExpression::IsValid() const {

src/ast/module.cc

@@ -30,33 +30,13 @@ Module& Module::operator=(Module&& rhs) = default;
 
 Module::~Module() = default;
 
-Module Module::Clone() {
+Module Module::Clone() const {
   Module out;
-  CloneContext ctx(&out);
+  CloneContext(&out, this).Clone();
-
-  // Symbol table must be cloned first so that the resulting module has the
-  // symbols before we start the tree mutations.
-  ctx.mod->symbol_table_ = symbol_table_;
-
-  CloneUsing(&ctx);
   return out;
 }
 
-Module Module::Clone(const std::function<void(CloneContext* ctx)>& init) {
+void Module::Clone(CloneContext* ctx) const {
-  Module out;
-  CloneContext ctx(&out);
-
-  // Symbol table must be cloned first so that the resulting module has the
-  // symbols before we start the tree mutations.
-  ctx.mod->symbol_table_ = symbol_table_;
-
-  init(&ctx);
-
-  CloneUsing(&ctx);
-  return out;
-}
-
-void Module::CloneUsing(CloneContext* ctx) {
   for (auto* ty : constructed_types_) {
     ctx->mod->constructed_types_.emplace_back(ctx->Clone(ty));
   }

src/ast/module.h

@@ -51,13 +51,15 @@ class Module {
   ~Module();
 
   /// @return a deep copy of this module
-  Module Clone();
+  Module Clone() const;
 
-  /// @param init a callback function to configure the CloneContex before
+  /// Clone this module into `ctx->mod` using the provided CloneContext
-  /// cloning any of the module's state
+  /// The module will be cloned in this order:
-  /// @return a deep copy of this module, calling `init` to first initialize the
+  /// * Constructed types
-  /// context.
+  /// * Global variables
-  Module Clone(const std::function<void(CloneContext* ctx)>& init);
+  /// * Functions
+  /// @param ctx the clone context
+  void Clone(CloneContext* ctx) const;
 
   /// Add a global variable to the module
   /// @param var the variable to add
@@ -177,14 +179,6 @@ class Module {
  private:
   Module(const Module&) = delete;
 
-  /// Clone this module into `ctx->mod` using the provided CloneContext
-  /// The module will be cloned in this order:
-  /// * Constructed types
-  /// * Global variables
-  /// * Functions
-  /// @param ctx the clone context
-  void CloneUsing(CloneContext* ctx);
-
   SymbolTable symbol_table_;
   VariableList global_variables_;
   // The constructed types are owned by the type manager

src/ast/module_clone_test.cc

@@ -16,6 +16,7 @@
 
 #include "gtest/gtest.h"
 #include "src/ast/case_statement.h"
+#include "src/demangler.h"
 #include "src/reader/wgsl/parser.h"
 #include "src/writer/wgsl/generator.h"
 
@@ -116,7 +117,9 @@ fn main() -> void {
   auto dst = src.Clone();
 
   // Expect the AST printed with to_str() to match
-  EXPECT_EQ(src.to_str(), dst.to_str());
+  Demangler demanger;
+  EXPECT_EQ(demanger.Demangle(src, src.to_str()),
+            demanger.Demangle(dst, dst.to_str()));
 
   // Check that none of the AST nodes or type pointers in dst are found in src
   std::unordered_set<ast::Node*> src_nodes;

src/ast/type/alias_type.cc

@@ -47,7 +47,8 @@ uint64_t Alias::BaseAlignment(MemoryLayout mem_layout) const {
 }
 
 Alias* Alias::Clone(CloneContext* ctx) const {
-  return ctx->mod->create<Alias>(symbol_, name_, ctx->Clone(subtype_));
+  return ctx->mod->create<Alias>(ctx->Clone(symbol()), name_,
+                                 ctx->Clone(subtype_));
 }
 
 }  // namespace type

src/ast/type/struct_type.cc

@@ -84,7 +84,8 @@ uint64_t Struct::BaseAlignment(MemoryLayout mem_layout) const {
 }
 
 Struct* Struct::Clone(CloneContext* ctx) const {
-  return ctx->mod->create<Struct>(symbol_, name_, ctx->Clone(struct_));
+  return ctx->mod->create<Struct>(ctx->Clone(symbol()), name_,
+                                  ctx->Clone(struct_));
 }
 
 }  // namespace type

src/transform/bound_array_accessors.cc

@@ -53,13 +53,14 @@ namespace transform {
 BoundArrayAccessors::BoundArrayAccessors() = default;
 BoundArrayAccessors::~BoundArrayAccessors() = default;
 
-Transform::Output BoundArrayAccessors::Run(ast::Module* mod) {
+Transform::Output BoundArrayAccessors::Run(ast::Module* in) {
   Output out;
-  out.module = mod->Clone([&](ast::CloneContext* ctx) {
+  ast::CloneContext(&out.module, in)
-    ctx->ReplaceAll([&, ctx](ast::ArrayAccessorExpression* expr) {
+      .ReplaceAll(
-      return Transform(expr, ctx, &out.diagnostics);
+          [&](ast::CloneContext* ctx, ast::ArrayAccessorExpression* expr) {
-    });
+            return Transform(expr, ctx, &out.diagnostics);
-  });
+          })
+      .Clone();
   return out;
 }
 

src/transform/emit_vertex_point_size.cc

@@ -47,49 +47,41 @@ Transform::Output EmitVertexPointSize::Run(ast::Module* in) {
     return out;
   }
 
-  tint::ast::AssignmentStatement* pointsize_assign = nullptr;
+  auto* f32 = out.module.create<ast::type::F32>();
-  auto get_pointsize_assign = [&pointsize_assign](ast::Module* mod) {
-    if (pointsize_assign != nullptr) {
-      return pointsize_assign;
-    }
 
-    auto* f32 = mod->create<ast::type::F32>();
+  // Declare the pointsize builtin output variable.
+  auto* pointsize_var = out.module.create<ast::Variable>(
+      Source{},                    // source
+      kPointSizeVar,               // name
+      ast::StorageClass::kOutput,  // storage_class
+      f32,                         // type
+      false,                       // is_const
+      nullptr,                     // constructor
+      ast::VariableDecorationList{
+          // decorations
+          out.module.create<ast::BuiltinDecoration>(Source{},
+                                                    ast::Builtin::kPointSize),
+      });
+  out.module.AddGlobalVariable(pointsize_var);
 
-    // Declare the pointsize builtin output variable.
+  // Build the AST expression & statement for assigning pointsize one.
-    auto* pointsize_var =
+  auto* one = out.module.create<ast::ScalarConstructorExpression>(
-        mod->create<ast::Variable>(Source{},                    // source
+      Source{}, out.module.create<ast::FloatLiteral>(Source{}, f32, 1.0f));
-                                   kPointSizeVar,               // name
+  auto* pointsize_ident = out.module.create<ast::IdentifierExpression>(
-                                   ast::StorageClass::kOutput,  // storage_class
+      Source{}, out.module.RegisterSymbol(kPointSizeVar), kPointSizeVar);
-                                   f32,                         // type
+  auto* pointsize_assign = out.module.create<ast::AssignmentStatement>(
-                                   false,                       // is_const
+      Source{}, pointsize_ident, one);
-                                   nullptr,                     // constructor
-                                   ast::VariableDecorationList{
-                                       // decorations
-                                       mod->create<ast::BuiltinDecoration>(
-                                           Source{}, ast::Builtin::kPointSize),
-                                   });
-    mod->AddGlobalVariable(pointsize_var);
-
-    // Build the AST expression & statement for assigning pointsize one.
-    auto* one = mod->create<ast::ScalarConstructorExpression>(
-        Source{}, mod->create<ast::FloatLiteral>(Source{}, f32, 1.0f));
-    auto* pointsize_ident = mod->create<ast::IdentifierExpression>(
-        Source{}, mod->RegisterSymbol(kPointSizeVar), kPointSizeVar);
-    pointsize_assign =
-        mod->create<ast::AssignmentStatement>(Source{}, pointsize_ident, one);
-    return pointsize_assign;
-  };
 
   // Add the pointsize assignment statement to the front of all vertex stages.
-  out.module = in->Clone([&](ast::CloneContext* ctx) {
+  ast::CloneContext(&out.module, in)
-    ctx->ReplaceAll([&, ctx](ast::Function* func) -> ast::Function* {
+      .ReplaceAll(
-      if (func->pipeline_stage() != ast::PipelineStage::kVertex) {
+          [&](ast::CloneContext* ctx, ast::Function* func) -> ast::Function* {
-        return nullptr;  // Just clone func
+            if (func->pipeline_stage() != ast::PipelineStage::kVertex) {
-      }
+              return nullptr;  // Just clone func
-      return CloneWithStatementsAtStart(ctx, func,
+            }
-                                        {get_pointsize_assign(ctx->mod)});
+            return CloneWithStatementsAtStart(ctx, func, {pointsize_assign});
-    });
+          })
-  });
+      .Clone();
 
   return out;
 }

src/transform/first_index_offset.cc

@@ -126,48 +126,50 @@ Transform::Output FirstIndexOffset::Run(ast::Module* in) {
   // these builtins.
 
   Output out;
-  out.module = in->Clone([&](ast::CloneContext* ctx) {
+  ast::CloneContext(&out.module, in)
-    ctx->ReplaceAll([&, ctx](ast::Variable* var) -> ast::Variable* {
+      .ReplaceAll(
-      for (ast::VariableDecoration* dec : var->decorations()) {
+          [&](ast::CloneContext* ctx, ast::Variable* var) -> ast::Variable* {
-        if (auto* blt_dec = dec->As<ast::BuiltinDecoration>()) {
+            for (ast::VariableDecoration* dec : var->decorations()) {
-          ast::Builtin blt_type = blt_dec->value();
+              if (auto* blt_dec = dec->As<ast::BuiltinDecoration>()) {
-          if (blt_type == ast::Builtin::kVertexIdx) {
+                ast::Builtin blt_type = blt_dec->value();
-            vertex_index_name = var->name();
+                if (blt_type == ast::Builtin::kVertexIdx) {
-            has_vertex_index_ = true;
+                  vertex_index_name = var->name();
-            return clone_variable_with_new_name(
+                  has_vertex_index_ = true;
-                ctx, var, kIndexOffsetPrefix + var->name());
+                  return clone_variable_with_new_name(
-          } else if (blt_type == ast::Builtin::kInstanceIdx) {
+                      ctx, var, kIndexOffsetPrefix + var->name());
-            instance_index_name = var->name();
+                } else if (blt_type == ast::Builtin::kInstanceIdx) {
-            has_instance_index_ = true;
+                  instance_index_name = var->name();
-            return clone_variable_with_new_name(
+                  has_instance_index_ = true;
-                ctx, var, kIndexOffsetPrefix + var->name());
+                  return clone_variable_with_new_name(
-          }
+                      ctx, var, kIndexOffsetPrefix + var->name());
-        }
+                }
-      }
+              }
-      return nullptr;  // Just clone var
-    });
-    ctx->ReplaceAll(  // Note: This happens in the same pass as the rename above
-                      // which determines the original builtin variable names,
-                      // but this should be fine, as variables are cloned first.
-        [&, ctx](ast::Function* func) -> ast::Function* {
-          maybe_create_buffer_var(ctx->mod);
-          if (buffer_var == nullptr) {
-            return nullptr;  // no transform need, just clone func
-          }
-          ast::StatementList statements;
-          for (const auto& data : func->local_referenced_builtin_variables()) {
-            if (data.second->value() == ast::Builtin::kVertexIdx) {
-              statements.emplace_back(CreateFirstIndexOffset(
-                  vertex_index_name, kFirstVertexName, buffer_var, ctx->mod));
-            } else if (data.second->value() == ast::Builtin::kInstanceIdx) {
-              statements.emplace_back(CreateFirstIndexOffset(
-                  instance_index_name, kFirstInstanceName, buffer_var,
-                  ctx->mod));
             }
-          }
+            return nullptr;  // Just clone var
-          return CloneWithStatementsAtStart(ctx, func, statements);
+          })
-        });
+      .ReplaceAll(  // Note: This happens in the same pass as the rename above
-  });
+                    // which determines the original builtin variable names,
+                    // but this should be fine, as variables are cloned first.
+          [&](ast::CloneContext* ctx, ast::Function* func) -> ast::Function* {
+            maybe_create_buffer_var(ctx->mod);
+            if (buffer_var == nullptr) {
+              return nullptr;  // no transform need, just clone func
+            }
+            ast::StatementList statements;
+            for (const auto& data :
+                 func->local_referenced_builtin_variables()) {
+              if (data.second->value() == ast::Builtin::kVertexIdx) {
+                statements.emplace_back(CreateFirstIndexOffset(
+                    vertex_index_name, kFirstVertexName, buffer_var, ctx->mod));
+              } else if (data.second->value() == ast::Builtin::kInstanceIdx) {
+                statements.emplace_back(CreateFirstIndexOffset(
+                    instance_index_name, kFirstInstanceName, buffer_var,
+                    ctx->mod));
+              }
+            }
+            return CloneWithStatementsAtStart(ctx, func, statements);
+          })
+      .Clone();
 
   return out;
 }

src/transform/transform.cc

@@ -32,7 +32,7 @@ ast::Function* Transform::CloneWithStatementsAtStart(
     statements.emplace_back(ctx->Clone(s));
   }
   return ctx->mod->create<ast::Function>(
-      ctx->Clone(in->source()), in->symbol(), in->name(),
+      ctx->Clone(in->source()), ctx->Clone(in->symbol()), in->name(),
       ctx->Clone(in->params()), ctx->Clone(in->return_type()),
       ctx->mod->create<ast::BlockStatement>(ctx->Clone(in->body()->source()),
                                             statements),

src/transform/vertex_pulling.cc

@@ -98,21 +98,23 @@ Transform::Output VertexPulling::Run(ast::Module* in) {
   // TODO(idanr): Make sure we covered all error cases, to guarantee the
   // following stages will pass
   Output out;
-  out.module = in->Clone([&](ast::CloneContext* ctx) {
-    State state{in, ctx->mod, cfg};
-    state.FindOrInsertVertexIndexIfUsed();
-    state.FindOrInsertInstanceIndexIfUsed();
-    state.ConvertVertexInputVariablesToPrivate();
-    state.AddVertexStorageBuffers();
 
-    ctx->ReplaceAll([func, ctx, state](ast::Function* f) -> ast::Function* {
+  State state{in, &out.module, cfg};
-      if (f == func) {
+  state.FindOrInsertVertexIndexIfUsed();
-        return CloneWithStatementsAtStart(
+  state.FindOrInsertInstanceIndexIfUsed();
-            ctx, f, {state.CreateVertexPullingPreamble()});
+  state.ConvertVertexInputVariablesToPrivate();
-      }
+  state.AddVertexStorageBuffers();
-      return nullptr;  // Just clone func
+
-    });
+  ast::CloneContext(&out.module, in)
-  });
+      .ReplaceAll(
+          [&](ast::CloneContext* ctx, ast::Function* f) -> ast::Function* {
+            if (f == func) {
+              return CloneWithStatementsAtStart(
+                  ctx, f, {state.CreateVertexPullingPreamble()});
+            }
+            return nullptr;  // Just clone func
+          })
+      .Clone();
 
   return out;
 }