CloneContext: fix InsertAfter and InsertBack not working if done while cloning a node in a vector

Also make tests more idiomatic by removing diamonds in the Node
hierarchy.

Bug: tint:1300
Change-Id: I681f4251bd8d9bdef169dcdf5de345792c927436
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/79680
Reviewed-by: Ben Clayton <bclayton@google.com>
Kokoro: Kokoro <noreply+kokoro@google.com>
Commit-Queue: Antonio Maiorano <amaiorano@google.com>
diff --git a/src/clone_context.h b/src/clone_context.h
index 9b55a1f..974ee90 100644
--- a/src/clone_context.h
+++ b/src/clone_context.h
@@ -222,6 +222,29 @@
     } else {
       for (auto& el : from) {
         to.emplace_back(Clone(el));
+
+        // Clone(el) may have inserted after
+        list_transform_it = list_transforms_.find(&from);
+        if (list_transform_it != list_transforms_.end()) {
+          const auto& transforms = list_transform_it->second;
+
+          auto insert_after_it = transforms.insert_after_.find(el);
+          if (insert_after_it != transforms.insert_after_.end()) {
+            for (auto insert : insert_after_it->second) {
+              to.emplace_back(CheckedCast<T>(insert));
+            }
+          }
+        }
+      }
+
+      // Clone(el)s may have inserted back
+      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_back_) {
+          to.emplace_back(CheckedCast<T>(o));
+        }
       }
     }
   }
diff --git a/src/clone_context_test.cc b/src/clone_context_test.cc
index d51d4ae..a8e45ef 100644
--- a/src/clone_context_test.cc
+++ b/src/clone_context_test.cc
@@ -381,15 +381,16 @@
 
   ProgramBuilder builder;
   auto* original_root = a.Create<Node>(builder.Symbols().Register("root"));
-  original_root->a = a.Create<Node>(builder.Symbols().Register("a"));
-  original_root->b = a.Create<Node>(builder.Symbols().Register("b"));
-  original_root->c = a.Create<Node>(builder.Symbols().Register("c"));
-  original_root->vec = {original_root->a, original_root->b, original_root->c};
+  original_root->vec = {
+      a.Create<Node>(builder.Symbols().Register("a")),
+      a.Create<Node>(builder.Symbols().Register("b")),
+      a.Create<Node>(builder.Symbols().Register("c")),
+  };
   Program original(std::move(builder));
 
   ProgramBuilder cloned;
   auto* cloned_root = CloneContext(&cloned, &original)
-                          .Remove(original_root->vec, original_root->b)
+                          .Remove(original_root->vec, original_root->vec[1])
                           .Clone(original_root);
 
   EXPECT_EQ(cloned_root->vec.size(), 2u);
@@ -407,10 +408,11 @@
 
   ProgramBuilder builder;
   auto* original_root = a.Create<Node>(builder.Symbols().Register("root"));
-  original_root->a = a.Create<Node>(builder.Symbols().Register("a"));
-  original_root->b = a.Create<Node>(builder.Symbols().Register("b"));
-  original_root->c = a.Create<Node>(builder.Symbols().Register("c"));
-  original_root->vec = {original_root->a, original_root->b, original_root->c};
+  original_root->vec = {
+      a.Create<Node>(builder.Symbols().Register("a")),
+      a.Create<Node>(builder.Symbols().Register("b")),
+      a.Create<Node>(builder.Symbols().Register("c")),
+  };
   Program original(std::move(builder));
 
   ProgramBuilder cloned;
@@ -459,10 +461,11 @@
 
   ProgramBuilder builder;
   auto* original_root = a.Create<Node>(builder.Symbols().Register("root"));
-  original_root->a = a.Create<Node>(builder.Symbols().Register("a"));
-  original_root->b = a.Create<Node>(builder.Symbols().Register("b"));
-  original_root->c = a.Create<Node>(builder.Symbols().Register("c"));
-  original_root->vec = {original_root->a, original_root->b, original_root->c};
+  original_root->vec = {
+      a.Create<Node>(builder.Symbols().Register("a")),
+      a.Create<Node>(builder.Symbols().Register("b")),
+      a.Create<Node>(builder.Symbols().Register("c")),
+  };
   Program original(std::move(builder));
 
   ProgramBuilder cloned;
@@ -532,10 +535,11 @@
 
   ProgramBuilder builder;
   auto* original_root = a.Create<Node>(builder.Symbols().Register("root"));
-  original_root->a = a.Create<Node>(builder.Symbols().Register("a"));
-  original_root->b = a.Create<Node>(builder.Symbols().Register("b"));
-  original_root->c = a.Create<Node>(builder.Symbols().Register("c"));
-  original_root->vec = {original_root->a, original_root->b, original_root->c};
+  original_root->vec = {
+      a.Create<Node>(builder.Symbols().Register("a")),
+      a.Create<Node>(builder.Symbols().Register("b")),
+      a.Create<Node>(builder.Symbols().Register("c")),
+  };
   Program original(std::move(builder));
 
   ProgramBuilder cloned;
@@ -543,15 +547,11 @@
 
   auto* cloned_root =
       CloneContext(&cloned, &original)
-          .InsertBefore(original_root->vec, original_root->b, insertion)
+          .InsertBefore(original_root->vec, original_root->vec[1], insertion)
           .Clone(original_root);
 
   EXPECT_EQ(cloned_root->vec.size(), 4u);
 
-  EXPECT_NE(cloned_root->vec[0], cloned_root->a);
-  EXPECT_NE(cloned_root->vec[2], cloned_root->b);
-  EXPECT_NE(cloned_root->vec[3], cloned_root->c);
-
   EXPECT_EQ(cloned_root->name, cloned.Symbols().Get("root"));
   EXPECT_EQ(cloned_root->vec[0]->name, cloned.Symbols().Get("a"));
   EXPECT_EQ(cloned_root->vec[1]->name, cloned.Symbols().Get("insertion"));
@@ -564,10 +564,11 @@
 
   ProgramBuilder builder;
   auto* original_root = a.Create<Node>(builder.Symbols().Register("root"));
-  original_root->a = a.Create<Node>(builder.Symbols().Register("a"));
-  original_root->b = a.Create<Node>(builder.Symbols().Register("b"));
-  original_root->c = a.Create<Node>(builder.Symbols().Register("c"));
-  original_root->vec = {original_root->a, original_root->b, original_root->c};
+  original_root->vec = {
+      a.Create<Node>(builder.Symbols().Register("a")),
+      a.Create<Node>(builder.Symbols().Register("b")),
+      a.Create<Node>(builder.Symbols().Register("c")),
+  };
   Program original(std::move(builder));
 
   ProgramBuilder cloned;
@@ -575,15 +576,11 @@
 
   auto* cloned_root =
       CloneContext(&cloned, &original)
-          .InsertAfter(original_root->vec, original_root->b, insertion)
+          .InsertAfter(original_root->vec, original_root->vec[1], insertion)
           .Clone(original_root);
 
   EXPECT_EQ(cloned_root->vec.size(), 4u);
 
-  EXPECT_NE(cloned_root->vec[0], cloned_root->a);
-  EXPECT_NE(cloned_root->vec[1], cloned_root->b);
-  EXPECT_NE(cloned_root->vec[3], cloned_root->c);
-
   EXPECT_EQ(cloned_root->name, cloned.Symbols().Get("root"));
   EXPECT_EQ(cloned_root->vec[0]->name, cloned.Symbols().Get("a"));
   EXPECT_EQ(cloned_root->vec[1]->name, cloned.Symbols().Get("b"));
@@ -591,15 +588,80 @@
   EXPECT_EQ(cloned_root->vec[3]->name, cloned.Symbols().Get("c"));
 }
 
+TEST_F(CloneContextNodeTest, CloneWithInsertAfterInVectorNodeClone) {
+  Allocator a;
+
+  ProgramBuilder builder;
+  auto* original_root = a.Create<Node>(builder.Symbols().Register("root"));
+  original_root->vec = {
+      a.Create<Node>(builder.Symbols().Register("a")),
+      a.Create<Replaceable>(builder.Symbols().Register("b")),
+      a.Create<Node>(builder.Symbols().Register("c")),
+  };
+
+  Program original(std::move(builder));
+
+  ProgramBuilder cloned;
+  CloneContext ctx(&cloned, &original);
+  ctx.ReplaceAll([&](const Replaceable* r) {
+    auto* insertion = a.Create<Node>(cloned.Symbols().New("insertion"));
+    ctx.InsertAfter(original_root->vec, r, insertion);
+    return nullptr;
+  });
+
+  auto* cloned_root = ctx.Clone(original_root);
+
+  EXPECT_EQ(cloned_root->vec.size(), 4u);
+
+  EXPECT_EQ(cloned_root->name, cloned.Symbols().Get("root"));
+  EXPECT_EQ(cloned_root->vec[0]->name, cloned.Symbols().Get("a"));
+  EXPECT_EQ(cloned_root->vec[1]->name, cloned.Symbols().Get("b"));
+  EXPECT_EQ(cloned_root->vec[2]->name, cloned.Symbols().Get("insertion"));
+  EXPECT_EQ(cloned_root->vec[3]->name, cloned.Symbols().Get("c"));
+}
+
+TEST_F(CloneContextNodeTest, CloneWithInsertBackInVectorNodeClone) {
+  Allocator a;
+
+  ProgramBuilder builder;
+  auto* original_root = a.Create<Node>(builder.Symbols().Register("root"));
+  original_root->vec = {
+      a.Create<Node>(builder.Symbols().Register("a")),
+      a.Create<Replaceable>(builder.Symbols().Register("b")),
+      a.Create<Node>(builder.Symbols().Register("c")),
+  };
+
+  Program original(std::move(builder));
+
+  ProgramBuilder cloned;
+  CloneContext ctx(&cloned, &original);
+  ctx.ReplaceAll([&](const Replaceable* /*r*/) {
+    auto* insertion = a.Create<Node>(cloned.Symbols().New("insertion"));
+    ctx.InsertBack(original_root->vec, insertion);
+    return nullptr;
+  });
+
+  auto* cloned_root = ctx.Clone(original_root);
+
+  EXPECT_EQ(cloned_root->vec.size(), 4u);
+
+  EXPECT_EQ(cloned_root->name, cloned.Symbols().Get("root"));
+  EXPECT_EQ(cloned_root->vec[0]->name, cloned.Symbols().Get("a"));
+  EXPECT_EQ(cloned_root->vec[1]->name, cloned.Symbols().Get("b"));
+  EXPECT_EQ(cloned_root->vec[2]->name, cloned.Symbols().Get("c"));
+  EXPECT_EQ(cloned_root->vec[3]->name, cloned.Symbols().Get("insertion"));
+}
+
 TEST_F(CloneContextNodeTest, CloneWithInsertBeforeAndAfterRemoved) {
   Allocator a;
 
   ProgramBuilder builder;
   auto* original_root = a.Create<Node>(builder.Symbols().Register("root"));
-  original_root->a = a.Create<Node>(builder.Symbols().Register("a"));
-  original_root->b = a.Create<Node>(builder.Symbols().Register("b"));
-  original_root->c = a.Create<Node>(builder.Symbols().Register("c"));
-  original_root->vec = {original_root->a, original_root->b, original_root->c};
+  original_root->vec = {
+      a.Create<Node>(builder.Symbols().Register("a")),
+      a.Create<Node>(builder.Symbols().Register("b")),
+      a.Create<Node>(builder.Symbols().Register("c")),
+  };
   Program original(std::move(builder));
 
   ProgramBuilder cloned;
@@ -608,18 +670,16 @@
   auto* insertion_after =
       a.Create<Node>(cloned.Symbols().New("insertion_after"));
 
-  auto* cloned_root =
-      CloneContext(&cloned, &original)
-          .InsertBefore(original_root->vec, original_root->b, insertion_before)
-          .InsertAfter(original_root->vec, original_root->b, insertion_after)
-          .Remove(original_root->vec, original_root->b)
-          .Clone(original_root);
+  auto* cloned_root = CloneContext(&cloned, &original)
+                          .InsertBefore(original_root->vec,
+                                        original_root->vec[1], insertion_before)
+                          .InsertAfter(original_root->vec,
+                                       original_root->vec[1], insertion_after)
+                          .Remove(original_root->vec, original_root->vec[1])
+                          .Clone(original_root);
 
   EXPECT_EQ(cloned_root->vec.size(), 4u);
 
-  EXPECT_NE(cloned_root->vec[0], cloned_root->a);
-  EXPECT_NE(cloned_root->vec[3], cloned_root->c);
-
   EXPECT_EQ(cloned_root->name, cloned.Symbols().Get("root"));
   EXPECT_EQ(cloned_root->vec[0]->name, cloned.Symbols().Get("a"));
   EXPECT_EQ(cloned_root->vec[1]->name,