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/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) {
-    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;
 }
 

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 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;
 }

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) {
-    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;
 }

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* {
-      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;
 }