From 5d9b1c38c23b806e39bae15cef92a5bc8975d2ca Mon Sep 17 00:00:00 2001
From: Ben Clayton <bclayton@google.com>
Date: Tue, 26 Jan 2021 16:57:10 +0000
Subject: [PATCH] Add CloneContext::Replace()

Instructs the clone context to replace a single object instance with a given replacement.
Will be used to fix brokenness in transforms.

Bug: tint:390
Change-Id: I17bf1cdf7549f697281ca7c286bdb5771e5a6f1a
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/38553
Reviewed-by: dan sinclair <dsinclair@chromium.org>
---
 src/clone_context.h       | 13 +++++++++++++
 src/clone_context_test.cc | 26 +++++++++++++++++++++++++-
 2 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/src/clone_context.h b/src/clone_context.h
index fc391b015b..7d8253203c 100644
--- a/src/clone_context.h
+++ b/src/clone_context.h
@@ -154,6 +154,19 @@ class CloneContext {
     return *this;
   }
 
+  /// Replace replaces all occurrences of `what` in #src with `with` in #dst
+  /// when calling Clone().
+  /// @param what a pointer to the object in #src that will be replaced with
+  /// `with`
+  /// @param with a pointer to the replacement object that will be used when
+  /// cloning into #dst
+  /// @returns this CloneContext so calls can be chained
+  template <typename T>
+  CloneContext& Replace(T* what, T* with) {
+    cloned_.emplace(what, with);
+    return *this;
+  }
+
   /// Clone performs the clone of the entire program #src to #dst.
   void Clone();
 
diff --git a/src/clone_context_test.cc b/src/clone_context_test.cc
index 5762cf3247..ff6a7814d9 100644
--- a/src/clone_context_test.cc
+++ b/src/clone_context_test.cc
@@ -87,7 +87,7 @@ TEST(CloneContext, Clone) {
   EXPECT_EQ(cloned_root->c, cloned_root->b);     // Aliased
 }
 
-TEST(CloneContext, CloneWithReplacements) {
+TEST(CloneContext, CloneWithReplaceAll) {
   Program original;
   auto* original_root = original.create<Cloneable>();
   original_root->a = original.create<Cloneable>();
@@ -160,6 +160,30 @@ TEST(CloneContext, CloneWithReplacements) {
   EXPECT_FALSE(cloned_root->b->b->Is<Replacement>());
 }
 
+TEST(CloneContext, CloneWithReplace) {
+  Program original;
+  auto* original_root = original.create<Cloneable>();
+  original_root->a = original.create<Cloneable>();
+  original_root->b = original.create<Cloneable>();
+  original_root->c = original.create<Cloneable>();
+
+  //                          root
+  //        ╭──────────────────┼──────────────────╮
+  //       (a)                (b)                (c)
+  //                        Replaced
+
+  Program cloned;
+  auto* replacement = cloned.create<Cloneable>();
+
+  auto* cloned_root = CloneContext(&cloned, &original)
+                          .Replace(original_root->b, replacement)
+                          .Clone(original_root);
+
+  EXPECT_NE(cloned_root->a, replacement);
+  EXPECT_EQ(cloned_root->b, replacement);
+  EXPECT_NE(cloned_root->c, replacement);
+}
+
 }  // namespace
 
 TINT_INSTANTIATE_CLASS_ID(Cloneable);