ast: Tweak cloning rules for module-scope decls

Previously the Clone() of the AST would clone all the functions, globals
and type declarations in a temporary vector, then assign this to the
ast::Module. This meant that adding new module-scope declarations inside
callbacks of ReplaceAll() would place them right at the top, before any
of the cloned declarations.

As top-level declarations are not statements, ensuring that a new object
comes before the current ReplaceAll() declaration is surprisingly
tricky.

With this change, we can now safely assume that calling
ProgramBuilder::Var(), ProgramBuilder::Func(), ProgramBuilder::Alias()
or ProgramBuilder::Structure() inside a ReplaceAll() will add that
module-scoped declaration before the currently processed top-level
declaration.

Change-Id: I52772fdc85940c8ac8d941fbd53374a4dd64a9f4
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/54320
Kokoro: Kokoro <noreply+kokoro@google.com>
Reviewed-by: James Price <jrprice@google.com>

src/ast/module.cc

@@ -69,6 +69,7 @@ void Module::AddGlobalVariable(ast::Variable* var) {
 
 void Module::AddTypeDecl(ast::TypeDecl* type) {
   TINT_ASSERT(type);
+  TINT_ASSERT_PROGRAM_IDS_EQUAL_IF_VALID(type, program_id());
   type_decls_.push_back(type);
   global_declarations_.push_back(type);
 }
@@ -87,17 +88,21 @@ Module* Module::Clone(CloneContext* ctx) const {
 }
 
 void Module::Copy(CloneContext* ctx, const Module* src) {
-  for (auto* decl : ctx->Clone(src->global_declarations_)) {
+  ctx->Clone(global_declarations_, src->global_declarations_);
+  for (auto* decl : global_declarations_) {
     if (!decl) {
       TINT_ICE(ctx->dst->Diagnostics()) << "src global declaration was nullptr";
       continue;
     }
-    if (auto* ty = decl->As<ast::TypeDecl>()) {
-      AddTypeDecl(ty);
+    if (auto* type = decl->As<ast::TypeDecl>()) {
+      TINT_ASSERT_PROGRAM_IDS_EQUAL_IF_VALID(type, program_id());
+      type_decls_.push_back(type);
     } else if (auto* func = decl->As<Function>()) {
-      AddFunction(func);
+      TINT_ASSERT_PROGRAM_IDS_EQUAL_IF_VALID(func, program_id());
+      functions_.push_back(func);
     } else if (auto* var = decl->As<Variable>()) {
-      AddGlobalVariable(var);
+      TINT_ASSERT_PROGRAM_IDS_EQUAL_IF_VALID(var, program_id());
+      global_variables_.push_back(var);
     } else {
       TINT_ICE(ctx->dst->Diagnostics()) << "Unknown global declaration type";
     }

src/ast/module_test.cc

@@ -14,6 +14,7 @@
 
 #include "gtest/gtest-spi.h"
 #include "src/ast/test_helper.h"
+#include "src/clone_context.h"
 
 namespace tint {
 namespace ast {
@@ -96,6 +97,53 @@ TEST_F(ModuleTest, Assert_Null_Function) {
       "internal compiler error");
 }
 
+TEST_F(ModuleTest, CloneOrder) {
+  // Create a program with a function, alias decl and var decl.
+  Program p = [] {
+    ProgramBuilder b;
+    b.Func("F", {}, b.ty.void_(), {});
+    b.Alias("A", b.ty.u32());
+    b.Global("V", b.ty.i32(), ast::StorageClass::kPrivate);
+    return Program(std::move(b));
+  }();
+
+  // Clone the program, using ReplaceAll() to create new module-scope
+  // declarations. We want to test that these are added just before the
+  // declaration that triggered the ReplaceAll().
+  ProgramBuilder cloned;
+  CloneContext ctx(&cloned, &p);
+  ctx.ReplaceAll([&](ast::Function*) -> ast::Function* {
+    ctx.dst->Alias("inserted_before_F", cloned.ty.u32());
+    return nullptr;
+  });
+  ctx.ReplaceAll([&](ast::Alias*) -> ast::Alias* {
+    ctx.dst->Alias("inserted_before_A", cloned.ty.u32());
+    return nullptr;
+  });
+  ctx.ReplaceAll([&](ast::Variable*) -> ast::Variable* {
+    ctx.dst->Alias("inserted_before_V", cloned.ty.u32());
+    return nullptr;
+  });
+  ctx.Clone();
+
+  auto& decls = cloned.AST().GlobalDeclarations();
+  ASSERT_EQ(decls.size(), 6u);
+  EXPECT_TRUE(decls[1]->Is<ast::Function>());
+  EXPECT_TRUE(decls[3]->Is<ast::Alias>());
+  EXPECT_TRUE(decls[5]->Is<ast::Variable>());
+
+  ASSERT_TRUE(decls[0]->Is<ast::Alias>());
+  ASSERT_TRUE(decls[2]->Is<ast::Alias>());
+  ASSERT_TRUE(decls[4]->Is<ast::Alias>());
+
+  ASSERT_EQ(cloned.Symbols().NameFor(decls[0]->As<ast::Alias>()->name()),
+            "inserted_before_F");
+  ASSERT_EQ(cloned.Symbols().NameFor(decls[2]->As<ast::Alias>()->name()),
+            "inserted_before_A");
+  ASSERT_EQ(cloned.Symbols().NameFor(decls[4]->As<ast::Alias>()->name()),
+            "inserted_before_V");
+}
+
 }  // namespace
 }  // namespace ast
 }  // namespace tint