From 57be2ff2eb03eba86fb4b3e0b1bd9cc1869969fa Mon Sep 17 00:00:00 2001
From: Ben Clayton <bclayton@google.com>
Date: Thu, 2 Mar 2023 17:45:21 +0000
Subject: [PATCH] tint: Change CloneContext::Replace() to replace the map entry

instead of silently being a no-op on second call.

This allows for repeated updates of a node's mapping, which is useful
for immediate-mode style cloning.

Change-Id: I6bd2851cbe67d3886d37cefc3ef9de488d4b0878
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/122300
Reviewed-by: James Price <jrprice@google.com>
Commit-Queue: Ben Clayton <bclayton@google.com>
Kokoro: Ben Clayton <bclayton@google.com>
Kokoro: Kokoro <noreply+kokoro@google.com>
Commit-Queue: Ben Clayton <bclayton@chromium.org>
---
 src/tint/clone_context.h       |  4 +-
 src/tint/clone_context_test.cc | 78 ++++++++++++++++++++++++++++++++++
 2 files changed, 80 insertions(+), 2 deletions(-)

diff --git a/src/tint/clone_context.h b/src/tint/clone_context.h
index 188ee253ac..8e0e2ff297 100644
--- a/src/tint/clone_context.h
+++ b/src/tint/clone_context.h
@@ -359,7 +359,7 @@ class CloneContext {
     CloneContext& Replace(const WHAT* what, const WITH* with) {
         TINT_ASSERT_PROGRAM_IDS_EQUAL_IF_VALID(Clone, src, what);
         TINT_ASSERT_PROGRAM_IDS_EQUAL_IF_VALID(Clone, dst, with);
-        replacements_.Add(what, [with]() -> const Cloneable* { return with; });
+        replacements_.Replace(what, [with]() -> const Cloneable* { return with; });
         return *this;
     }
 
@@ -379,7 +379,7 @@ class CloneContext {
     template <typename WHAT, typename WITH, typename = std::invoke_result_t<WITH>>
     CloneContext& Replace(const WHAT* what, WITH&& with) {
         TINT_ASSERT_PROGRAM_IDS_EQUAL_IF_VALID(Clone, src, what);
-        replacements_.Add(what, with);
+        replacements_.Replace(what, with);
         return *this;
     }
 
diff --git a/src/tint/clone_context_test.cc b/src/tint/clone_context_test.cc
index f2ed24779c..59c47c0cb4 100644
--- a/src/tint/clone_context_test.cc
+++ b/src/tint/clone_context_test.cc
@@ -339,6 +339,45 @@ TEST_F(CloneContextNodeTest, CloneWithReplacePointer) {
     EXPECT_EQ(cloned_root->c->name, cloned.Symbols().Get("c"));
 }
 
+TEST_F(CloneContextNodeTest, CloneWithRepeatedImmediateReplacePointer) {
+    Allocator a;
+
+    ProgramBuilder builder;
+    auto* original_root = a.Create<Node>(builder.Symbols().New("root"));
+    original_root->a = a.Create<Node>(builder.Symbols().New("a"));
+    original_root->b = a.Create<Node>(builder.Symbols().New("b"));
+    original_root->c = a.Create<Node>(builder.Symbols().New("c"));
+    Program original(std::move(builder));
+
+    ProgramBuilder cloned;
+
+    CloneContext ctx(&cloned, &original);
+
+    // Demonstrate that ctx.Replace() can be called multiple times to update the replacement of a
+    // node.
+
+    auto* replacement_x =
+        a.Create<Node>(cloned.Symbols().New("replacement_x"), ctx.Clone(original_root->b));
+    ctx.Replace(original_root->b, replacement_x);
+
+    auto* replacement_y =
+        a.Create<Node>(cloned.Symbols().New("replacement_y"), ctx.Clone(original_root->b));
+    ctx.Replace(original_root->b, replacement_y);
+
+    auto* replacement_z =
+        a.Create<Node>(cloned.Symbols().New("replacement_z"), ctx.Clone(original_root->b));
+    ctx.Replace(original_root->b, replacement_z);
+
+    auto* cloned_root = ctx.Clone(original_root);
+
+    EXPECT_NE(cloned_root->a, replacement_z);
+    EXPECT_EQ(cloned_root->b, replacement_z);
+    EXPECT_NE(cloned_root->c, replacement_z);
+
+    EXPECT_EQ(replacement_z->a, replacement_y);
+    EXPECT_EQ(replacement_y->a, replacement_x);
+}
+
 TEST_F(CloneContextNodeTest, CloneWithReplaceFunction) {
     Allocator a;
 
@@ -371,6 +410,45 @@ TEST_F(CloneContextNodeTest, CloneWithReplaceFunction) {
     EXPECT_EQ(cloned_root->c->name, cloned.Symbols().Get("c"));
 }
 
+TEST_F(CloneContextNodeTest, CloneWithRepeatedImmediateReplaceFunction) {
+    Allocator a;
+
+    ProgramBuilder builder;
+    auto* original_root = a.Create<Node>(builder.Symbols().New("root"));
+    original_root->a = a.Create<Node>(builder.Symbols().New("a"));
+    original_root->b = a.Create<Node>(builder.Symbols().New("b"));
+    original_root->c = a.Create<Node>(builder.Symbols().New("c"));
+    Program original(std::move(builder));
+
+    ProgramBuilder cloned;
+
+    CloneContext ctx(&cloned, &original);
+
+    // Demonstrate that ctx.Replace() can be called multiple times to update the replacement of a
+    // node.
+
+    Node* replacement_x =
+        a.Create<Node>(cloned.Symbols().New("replacement_x"), ctx.Clone(original_root->b));
+    ctx.Replace(original_root->b, [&] { return replacement_x; });
+
+    Node* replacement_y =
+        a.Create<Node>(cloned.Symbols().New("replacement_y"), ctx.Clone(original_root->b));
+    ctx.Replace(original_root->b, [&] { return replacement_y; });
+
+    Node* replacement_z =
+        a.Create<Node>(cloned.Symbols().New("replacement_z"), ctx.Clone(original_root->b));
+    ctx.Replace(original_root->b, [&] { return replacement_z; });
+
+    auto* cloned_root = ctx.Clone(original_root);
+
+    EXPECT_NE(cloned_root->a, replacement_z);
+    EXPECT_EQ(cloned_root->b, replacement_z);
+    EXPECT_NE(cloned_root->c, replacement_z);
+
+    EXPECT_EQ(replacement_z->a, replacement_y);
+    EXPECT_EQ(replacement_y->a, replacement_x);
+}
+
 TEST_F(CloneContextNodeTest, CloneWithRemove) {
     Allocator a;