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

src/clone_context.h

@@ -210,7 +210,7 @@ class CloneContext {
     return out;
   }
 
-  /// Clones each of the elements of the vector `v` into the ProgramBuilder
+  /// Clones each of the elements of the vector `v` using the ProgramBuilder
   /// #dst, inserting any additional elements into the list that were registered
   /// with calls to InsertBefore().
   ///
@@ -221,40 +221,53 @@ class CloneContext {
   template <typename T>
   std::vector<T*> Clone(const std::vector<T*>& v) {
     std::vector<T*> out;
-    out.reserve(v.size());
+    Clone(out, v);
+    return out;
+  }
 
-    auto list_transform_it = list_transforms_.find(&v);
+  /// Clones each of the elements of the vector `from` into the vector `to`,
+  /// inserting any additional elements into the list that were registered with
+  /// calls to InsertBefore().
+  ///
+  /// All the elements of the vector `from` must be owned by the Program #src.
+  ///
+  /// @param from the vector to clone
+  /// @param to the cloned result
+  template <typename T>
+  void Clone(std::vector<T*>& to, const std::vector<T*>& from) {
+    to.reserve(from.size());
+
+    auto list_transform_it = list_transforms_.find(&from);
     if (list_transform_it != list_transforms_.end()) {
       const auto& transforms = list_transform_it->second;
       for (auto* o : transforms.insert_front_) {
-        out.emplace_back(CheckedCast<T>(o));
+        to.emplace_back(CheckedCast<T>(o));
       }
-      for (auto& el : v) {
+      for (auto& el : from) {
         auto insert_before_it = transforms.insert_before_.find(el);
         if (insert_before_it != transforms.insert_before_.end()) {
           for (auto insert : insert_before_it->second) {
-            out.emplace_back(CheckedCast<T>(insert));
+            to.emplace_back(CheckedCast<T>(insert));
           }
         }
         if (transforms.remove_.count(el) == 0) {
-          out.emplace_back(Clone(el));
+          to.emplace_back(Clone(el));
         }
         auto insert_after_it = transforms.insert_after_.find(el);
         if (insert_after_it != transforms.insert_after_.end()) {
           for (auto insert : insert_after_it->second) {
-            out.emplace_back(CheckedCast<T>(insert));
+            to.emplace_back(CheckedCast<T>(insert));
           }
         }
       }
       for (auto* o : transforms.insert_back_) {
-        out.emplace_back(CheckedCast<T>(o));
+        to.emplace_back(CheckedCast<T>(o));
       }
     } else {
-      for (auto& el : v) {
-        out.emplace_back(Clone(el));
+      for (auto& el : from) {
+        to.emplace_back(Clone(el));
       }
     }
-    return out;
   }
 
   /// Clones each of the elements of the vector `v` into the ProgramBuilder