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>
diff --git a/src/ast/clone_context.cc b/src/ast/clone_context.cc
index 9ec4bfe..fe9e849 100644
--- a/src/ast/clone_context.cc
+++ b/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
diff --git a/src/ast/clone_context.h b/src/ast/clone_context.h
index a9e303b..64a7b11 100644
--- a/src/ast/clone_context.h
+++ b/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 CloneContext {
  public:
   /// Constructor
-  /// @param m the target module to clone into
-  explicit CloneContext(Module* m);
+  /// @param to the target module to clone into
+  /// @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 @@
     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 @@
   }
 
   /// 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
-  /// parameter of `replacer`.
+  /// is called with a type that matches (or derives from) the type of the
+  /// second parameter of `replacer`.
   ///
-  /// `replacer` must be function-like with the signature: `T* (T*)`, where `T`
-  /// is a type deriving from CastableBase.
+  /// `replacer` must be function-like with the signature:
+  ///   `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 @@
   ///
   /// ```
   ///   // Replace all ast::UintLiterals with the number 42
-  ///   CloneCtx ctx(mod);
-  ///   ctx.ReplaceAll([&] (ast::UintLiteral* in) {
-  ///     return ctx.mod->create<ast::UintLiteral>(ctx.Clone(in->source()),
-  ///                                              ctx.Clone(in->type()), 42);
-  ///   });
-  ///   auto* out = ctx.Clone(tree);
+  ///   CloneCtx ctx(&out, in)
+  ///     .ReplaceAll([&] (CloneContext* ctx, ast::UintLiteral* l) {
+  ///       return ctx->mod->create<ast::UintLiteral>(ctx->Clone(l->source()),
+  ///                                                 ctx->Clone(l->type()),
+  ///                                                 42);
+  ///     }).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) {
-    using TPtr = traits::ParamTypeT<F, 0>;
+  CloneContext& ReplaceAll(F replacer) {
+    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.
diff --git a/src/ast/clone_context_test.cc b/src/ast/clone_context_test.cc
index 386a1ae..d8890a2 100644
--- a/src/ast/clone_context_test.cc
+++ b/src/ast/clone_context_test.cc
@@ -65,8 +65,7 @@
   // C: Clonable
 
   ast::Module cloned;
-  CloneContext ctx(&cloned);
-  auto* cloned_root = original_root->Clone(&ctx);
+  auto* cloned_root = CloneContext(&cloned, &original).Clone(original_root);
 
   EXPECT_NE(cloned_root->a, nullptr);
   EXPECT_EQ(cloned_root->a->a, nullptr);
@@ -110,14 +109,14 @@
   // R: Replaceable
 
   ast::Module cloned;
-  CloneContext ctx(&cloned);
-  ctx.ReplaceAll([&](Replaceable* in) {
-    auto* out = cloned.create<Replacement>();
-    out->b = cloned.create<Cloneable>();
-    out->c = ctx.Clone(in->a);
-    return out;
-  });
-  auto* cloned_root = original_root->Clone(&ctx);
+  auto* cloned_root = CloneContext(&cloned, &original)
+                          .ReplaceAll([&](CloneContext* ctx, Replaceable* in) {
+                            auto* out = cloned.create<Replacement>();
+                            out->b = cloned.create<Cloneable>();
+                            out->c = ctx->Clone(in->a);
+                            return out;
+                          })
+                          .Clone(original_root);
 
   //                         root
   //        ╭─────────────────┼──────────────────╮
diff --git a/src/ast/function.cc b/src/ast/function.cc
index ae69287..1c22b9a 100644
--- a/src/ast/function.cc
+++ b/src/ast/function.cc
@@ -228,7 +228,7 @@
 
 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_));
 }
 
diff --git a/src/ast/identifier_expression.cc b/src/ast/identifier_expression.cc
index a38ea3f..b18ca10 100644
--- a/src/ast/identifier_expression.cc
+++ b/src/ast/identifier_expression.cc
@@ -32,8 +32,8 @@
 IdentifierExpression::~IdentifierExpression() = default;
 
 IdentifierExpression* IdentifierExpression::Clone(CloneContext* ctx) const {
-  return ctx->mod->create<IdentifierExpression>(ctx->Clone(source()), sym_,
-                                                name_);
+  return ctx->mod->create<IdentifierExpression>(ctx->Clone(source()),
+                                                ctx->Clone(symbol()), name_);
 }
 
 bool IdentifierExpression::IsValid() const {
diff --git a/src/ast/module.cc b/src/ast/module.cc
index a357554..c180639 100644
--- a/src/ast/module.cc
+++ b/src/ast/module.cc
@@ -30,33 +30,13 @@
 
 Module::~Module() = default;
 
-Module Module::Clone() {
+Module Module::Clone() 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_;
-
-  CloneUsing(&ctx);
+  CloneContext(&out, this).Clone();
   return out;
 }
 
-Module Module::Clone(const std::function<void(CloneContext* ctx)>& init) {
-  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) {
+void Module::Clone(CloneContext* ctx) const {
   for (auto* ty : constructed_types_) {
     ctx->mod->constructed_types_.emplace_back(ctx->Clone(ty));
   }
diff --git a/src/ast/module.h b/src/ast/module.h
index 3031964..a42bd7f 100644
--- a/src/ast/module.h
+++ b/src/ast/module.h
@@ -51,13 +51,15 @@
   ~Module();
 
   /// @return a deep copy of this module
-  Module Clone();
+  Module Clone() const;
 
-  /// @param init a callback function to configure the CloneContex before
-  /// cloning any of the module's state
-  /// @return a deep copy of this module, calling `init` to first initialize the
-  /// context.
-  Module Clone(const std::function<void(CloneContext* ctx)>& init);
+  /// 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 Clone(CloneContext* ctx) const;
 
   /// Add a global variable to the module
   /// @param var the variable to add
@@ -177,14 +179,6 @@
  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
diff --git a/src/ast/module_clone_test.cc b/src/ast/module_clone_test.cc
index 23be2ad..8405dc1 100644
--- a/src/ast/module_clone_test.cc
+++ b/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 @@
   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;
diff --git a/src/ast/type/alias_type.cc b/src/ast/type/alias_type.cc
index 723bca2..bef935b 100644
--- a/src/ast/type/alias_type.cc
+++ b/src/ast/type/alias_type.cc
@@ -47,7 +47,8 @@
 }
 
 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
diff --git a/src/ast/type/struct_type.cc b/src/ast/type/struct_type.cc
index ddf163b..7a127a6 100644
--- a/src/ast/type/struct_type.cc
+++ b/src/ast/type/struct_type.cc
@@ -84,7 +84,8 @@
 }
 
 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
diff --git a/src/transform/bound_array_accessors.cc b/src/transform/bound_array_accessors.cc
index 3adf6b7..4543a17 100644
--- a/src/transform/bound_array_accessors.cc
+++ b/src/transform/bound_array_accessors.cc
@@ -53,13 +53,14 @@
 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) {
-    ctx->ReplaceAll([&, ctx](ast::ArrayAccessorExpression* expr) {
-      return Transform(expr, ctx, &out.diagnostics);
-    });
-  });
+  ast::CloneContext(&out.module, in)
+      .ReplaceAll(
+          [&](ast::CloneContext* ctx, ast::ArrayAccessorExpression* expr) {
+            return Transform(expr, ctx, &out.diagnostics);
+          })
+      .Clone();
   return out;
 }
 
diff --git a/src/transform/emit_vertex_point_size.cc b/src/transform/emit_vertex_point_size.cc
index c80fab7..374bcd1 100644
--- a/src/transform/emit_vertex_point_size.cc
+++ b/src/transform/emit_vertex_point_size.cc
@@ -47,49 +47,41 @@
     return out;
   }
 
-  tint::ast::AssignmentStatement* pointsize_assign = nullptr;
-  auto get_pointsize_assign = [&pointsize_assign](ast::Module* mod) {
-    if (pointsize_assign != nullptr) {
-      return pointsize_assign;
-    }
+  auto* f32 = out.module.create<ast::type::F32>();
 
-    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.
-    auto* pointsize_var =
-        mod->create<ast::Variable>(Source{},                    // source
-                                   kPointSizeVar,               // name
-                                   ast::StorageClass::kOutput,  // storage_class
-                                   f32,                         // type
-                                   false,                       // is_const
-                                   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;
-  };
+  // Build the AST expression & statement for assigning pointsize one.
+  auto* one = out.module.create<ast::ScalarConstructorExpression>(
+      Source{}, out.module.create<ast::FloatLiteral>(Source{}, f32, 1.0f));
+  auto* pointsize_ident = out.module.create<ast::IdentifierExpression>(
+      Source{}, out.module.RegisterSymbol(kPointSizeVar), kPointSizeVar);
+  auto* pointsize_assign = out.module.create<ast::AssignmentStatement>(
+      Source{}, pointsize_ident, one);
 
   // Add the pointsize assignment statement to the front of all vertex stages.
-  out.module = in->Clone([&](ast::CloneContext* ctx) {
-    ctx->ReplaceAll([&, ctx](ast::Function* func) -> ast::Function* {
-      if (func->pipeline_stage() != ast::PipelineStage::kVertex) {
-        return nullptr;  // Just clone func
-      }
-      return CloneWithStatementsAtStart(ctx, func,
-                                        {get_pointsize_assign(ctx->mod)});
-    });
-  });
+  ast::CloneContext(&out.module, in)
+      .ReplaceAll(
+          [&](ast::CloneContext* ctx, ast::Function* func) -> ast::Function* {
+            if (func->pipeline_stage() != ast::PipelineStage::kVertex) {
+              return nullptr;  // Just clone func
+            }
+            return CloneWithStatementsAtStart(ctx, func, {pointsize_assign});
+          })
+      .Clone();
 
   return out;
 }
diff --git a/src/transform/first_index_offset.cc b/src/transform/first_index_offset.cc
index b89b94f..671076c 100644
--- a/src/transform/first_index_offset.cc
+++ b/src/transform/first_index_offset.cc
@@ -126,48 +126,50 @@
   // these builtins.
 
   Output out;
-  out.module = in->Clone([&](ast::CloneContext* ctx) {
-    ctx->ReplaceAll([&, ctx](ast::Variable* var) -> ast::Variable* {
-      for (ast::VariableDecoration* dec : var->decorations()) {
-        if (auto* blt_dec = dec->As<ast::BuiltinDecoration>()) {
-          ast::Builtin blt_type = blt_dec->value();
-          if (blt_type == ast::Builtin::kVertexIdx) {
-            vertex_index_name = var->name();
-            has_vertex_index_ = true;
-            return clone_variable_with_new_name(
-                ctx, var, kIndexOffsetPrefix + var->name());
-          } else if (blt_type == ast::Builtin::kInstanceIdx) {
-            instance_index_name = var->name();
-            has_instance_index_ = true;
-            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));
+  ast::CloneContext(&out.module, in)
+      .ReplaceAll(
+          [&](ast::CloneContext* ctx, ast::Variable* var) -> ast::Variable* {
+            for (ast::VariableDecoration* dec : var->decorations()) {
+              if (auto* blt_dec = dec->As<ast::BuiltinDecoration>()) {
+                ast::Builtin blt_type = blt_dec->value();
+                if (blt_type == ast::Builtin::kVertexIdx) {
+                  vertex_index_name = var->name();
+                  has_vertex_index_ = true;
+                  return clone_variable_with_new_name(
+                      ctx, var, kIndexOffsetPrefix + var->name());
+                } else if (blt_type == ast::Builtin::kInstanceIdx) {
+                  instance_index_name = var->name();
+                  has_instance_index_ = true;
+                  return clone_variable_with_new_name(
+                      ctx, var, kIndexOffsetPrefix + var->name());
+                }
+              }
             }
-          }
-          return CloneWithStatementsAtStart(ctx, func, statements);
-        });
-  });
+            return nullptr;  // Just clone var
+          })
+      .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;
 }
diff --git a/src/transform/transform.cc b/src/transform/transform.cc
index a03b943..51fe645 100644
--- a/src/transform/transform.cc
+++ b/src/transform/transform.cc
@@ -32,7 +32,7 @@
     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),
diff --git a/src/transform/vertex_pulling.cc b/src/transform/vertex_pulling.cc
index b1fa044..5bbbcb0 100644
--- a/src/transform/vertex_pulling.cc
+++ b/src/transform/vertex_pulling.cc
@@ -98,21 +98,23 @@
   // 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* {
-      if (f == func) {
-        return CloneWithStatementsAtStart(
-            ctx, f, {state.CreateVertexPullingPreamble()});
-      }
-      return nullptr;  // Just clone func
-    });
-  });
+  State state{in, &out.module, cfg};
+  state.FindOrInsertVertexIndexIfUsed();
+  state.FindOrInsertInstanceIndexIfUsed();
+  state.ConvertVertexInputVariablesToPrivate();
+  state.AddVertexStorageBuffers();
+
+  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;
 }