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

@@ -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();
 

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);