CloneContext: allow insert before/after of a node marked for removal

Change-Id: I48f232c9ba54e2b8457571bea8e6942bed001156
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/51402
Commit-Queue: Antonio Maiorano <amaiorano@google.com>
Kokoro: Kokoro <noreply+kokoro@google.com>
Reviewed-by: Ben Clayton <bclayton@google.com>
diff --git a/src/clone_context.h b/src/clone_context.h
index b00e547..f547e17 100644
--- a/src/clone_context.h
+++ b/src/clone_context.h
@@ -249,16 +249,15 @@
     if (list_transform_it != list_transforms_.end()) {
       const auto& transforms = list_transform_it->second;
       for (auto& el : v) {
-        if (transforms.remove_.count(el)) {
-          continue;
-        }
         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));
           }
         }
-        out.emplace_back(Clone(el));
+        if (transforms.remove_.count(el) == 0) {
+          out.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) {
diff --git a/src/clone_context_test.cc b/src/clone_context_test.cc
index 4ec9585..9a4a9f7 100644
--- a/src/clone_context_test.cc
+++ b/src/clone_context_test.cc
@@ -510,6 +510,50 @@
   EXPECT_EQ(cloned_root->vec[3]->name, cloned.Symbols().Get("c"));
 }
 
+TYPED_TEST(CloneContextNodeTest, CloneWithInsertBeforeAndAfterRemoved) {
+  using Node = typename TestFixture::Node;
+  constexpr bool is_unique = TestFixture::is_unique;
+
+  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};
+  Program original(std::move(builder));
+
+  ProgramBuilder cloned;
+  auto* insertion_before =
+      a.Create<Node>(cloned.Symbols().New("insertion_before"));
+  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);
+
+  EXPECT_EQ(cloned_root->vec.size(), 4u);
+  if (is_unique) {
+    EXPECT_NE(cloned_root->vec[0], cloned_root->a);
+    EXPECT_NE(cloned_root->vec[3], cloned_root->c);
+  } else {
+    EXPECT_EQ(cloned_root->vec[0], cloned_root->a);
+    EXPECT_EQ(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_before"));
+  EXPECT_EQ(cloned_root->vec[2]->name, cloned.Symbols().Get("insertion_after"));
+  EXPECT_EQ(cloned_root->vec[3]->name, cloned.Symbols().Get("c"));
+}
+
 TYPED_TEST(CloneContextNodeTest, CloneIntoSameBuilder) {
   using Node = typename TestFixture::Node;
   constexpr bool is_unique = TestFixture::is_unique;