tint: Add callback overloads to CloneContext::Insert*

Change-Id: I2fc0b93d3ea3ba0d9d64fe575364f188db564b8a
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/104041
Commit-Queue: Ben Clayton <bclayton@google.com>
Kokoro: Kokoro <noreply+kokoro@google.com>
Reviewed-by: Antonio Maiorano <amaiorano@google.com>

src/tint/clone_context.h

@@ -203,26 +203,26 @@ class CloneContext {
         auto transforms = list_transforms_.Find(&from);
 
         if (transforms) {
-            for (auto* o : transforms->insert_front_) {
+            for (auto& builder : transforms->insert_front_) {
-                to.Push(CheckedCast<T>(o));
+                to.Push(CheckedCast<T>(builder()));
             }
             for (auto& el : from) {
                 if (auto* insert_before = transforms->insert_before_.Find(el)) {
-                    for (auto insert : *insert_before) {
+                    for (auto& builder : *insert_before) {
-                        to.Push(CheckedCast<T>(insert));
+                        to.Push(CheckedCast<T>(builder()));
                     }
                 }
                 if (!transforms->remove_.Contains(el)) {
                     to.Push(Clone(el));
                 }
                 if (auto* insert_after = transforms->insert_after_.Find(el)) {
-                    for (auto insert : *insert_after) {
+                    for (auto& builder : *insert_after) {
-                        to.Push(CheckedCast<T>(insert));
+                        to.Push(CheckedCast<T>(builder()));
                     }
                 }
             }
-            for (auto* o : transforms->insert_back_) {
+            for (auto& builder : transforms->insert_back_) {
-                to.Push(CheckedCast<T>(o));
+                to.Push(CheckedCast<T>(builder()));
             }
         } else {
             for (auto& el : from) {
@@ -232,8 +232,8 @@ class CloneContext {
                 // transform for `from`.
                 if (transforms) {
                     if (auto* insert_after = transforms->insert_after_.Find(el)) {
-                        for (auto insert : *insert_after) {
+                        for (auto& builder : *insert_after) {
-                            to.Push(CheckedCast<T>(insert));
+                            to.Push(CheckedCast<T>(builder()));
                         }
                     }
                 }
@@ -242,8 +242,8 @@ class CloneContext {
             // Clone(el) may have updated the transformation list, adding an `insert_back_`
             // transform for `from`.
             if (transforms) {
-                for (auto* o : transforms->insert_back_) {
+                for (auto& builder : transforms->insert_back_) {
-                    to.Push(CheckedCast<T>(o));
+                    to.Push(CheckedCast<T>(builder()));
                 }
             }
         }
@@ -374,7 +374,7 @@ class CloneContext {
         return *this;
     }
 
-    /// Removes `object` from the cloned copy of `vector`.
+    /// Removes @p object from the cloned copy of @p vector.
     /// @param vector the vector in #src
     /// @param object a pointer to the object in #src that will be omitted from
     /// the cloned vector.
@@ -392,7 +392,7 @@ class CloneContext {
         return *this;
     }
 
-    /// Inserts `object` before any other objects of `vector`, when it is cloned.
+    /// Inserts @p object before any other objects of @p vector, when the vector is cloned.
     /// @param vector the vector in #src
     /// @param object a pointer to the object in #dst that will be inserted at the
     /// front of the vector
@@ -400,11 +400,21 @@ class CloneContext {
     template <typename T, size_t N, typename OBJECT>
     CloneContext& InsertFront(const utils::Vector<T, N>& vector, OBJECT* object) {
         TINT_ASSERT_PROGRAM_IDS_EQUAL_IF_VALID(Clone, dst, object);
-        list_transforms_.Edit(&vector).insert_front_.Push(object);
+        return InsertFront(vector, [object] { return object; });
+    }
+
+    /// Inserts a lazily built object before any other objects of @p vector, when the vector is
+    /// cloned.
+    /// @param vector the vector in #src
+    /// @param builder a builder of the object that will be inserted at the front of the vector.
+    /// @returns this CloneContext so calls can be chained
+    template <typename T, size_t N, typename BUILDER>
+    CloneContext& InsertFront(const utils::Vector<T, N>& vector, BUILDER&& builder) {
+        list_transforms_.Edit(&vector).insert_front_.Push(std::forward<BUILDER>(builder));
         return *this;
     }
 
-    /// Inserts `object` after any other objects of `vector`, when it is cloned.
+    /// Inserts @p object after any other objects of @p vector, when the vector is cloned.
     /// @param vector the vector in #src
     /// @param object a pointer to the object in #dst that will be inserted at the
     /// end of the vector
@@ -412,15 +422,26 @@ class CloneContext {
     template <typename T, size_t N, typename OBJECT>
     CloneContext& InsertBack(const utils::Vector<T, N>& vector, OBJECT* object) {
         TINT_ASSERT_PROGRAM_IDS_EQUAL_IF_VALID(Clone, dst, object);
-        list_transforms_.Edit(&vector).insert_back_.Push(object);
+        return InsertBack(vector, [object] { return object; });
+    }
+
+    /// Inserts a lazily built object after any other objects of @p vector, when the vector is
+    /// cloned.
+    /// @param vector the vector in #src
+    /// @param builder the builder of the object in #dst that will be inserted at the end of the
+    /// vector.
+    /// @returns this CloneContext so calls can be chained
+    template <typename T, size_t N, typename BUILDER>
+    CloneContext& InsertBack(const utils::Vector<T, N>& vector, BUILDER&& builder) {
+        list_transforms_.Edit(&vector).insert_back_.Push(std::forward<BUILDER>(builder));
         return *this;
     }
 
-    /// Inserts `object` before `before` whenever `vector` is cloned.
+    /// Inserts @p object before @p before whenever @p vector is cloned.
     /// @param vector the vector in #src
     /// @param before a pointer to the object in #src
     /// @param object a pointer to the object in #dst that will be inserted before
-    /// any occurrence of the clone of `before`
+    /// any occurrence of the clone of @p before
     /// @returns this CloneContext so calls can be chained
     template <typename T, size_t N, typename BEFORE, typename OBJECT>
     CloneContext& InsertBefore(const utils::Vector<T, N>& vector,
@@ -434,15 +455,35 @@ class CloneContext {
             return *this;
         }
 
-        list_transforms_.Edit(&vector).insert_before_.GetOrZero(before).Push(object);
+        list_transforms_.Edit(&vector).insert_before_.GetOrZero(before).Push(
+            [object] { return object; });
         return *this;
     }
 
-    /// Inserts `object` after `after` whenever `vector` is cloned.
+    /// Inserts a lazily created object before @p before whenever @p vector is cloned.
+    /// @param vector the vector in #src
+    /// @param before a pointer to the object in #src
+    /// @param builder the builder of the object in #dst that will be inserted before any occurrence
+    /// of the clone of @p before
+    /// @returns this CloneContext so calls can be chained
+    template <typename T,
+              size_t N,
+              typename BEFORE,
+              typename BUILDER,
+              typename _ = std::enable_if_t<!std::is_pointer_v<std::decay_t<BUILDER>>>>
+    CloneContext& InsertBefore(const utils::Vector<T, N>& vector,
+                               const BEFORE* before,
+                               BUILDER&& builder) {
+        list_transforms_.Edit(&vector).insert_before_.GetOrZero(before).Push(
+            std::forward<BUILDER>(builder));
+        return *this;
+    }
+
+    /// Inserts @p object after @p after whenever @p vector is cloned.
     /// @param vector the vector in #src
     /// @param after a pointer to the object in #src
     /// @param object a pointer to the object in #dst that will be inserted after
-    /// any occurrence of the clone of `after`
+    /// any occurrence of the clone of @p after
     /// @returns this CloneContext so calls can be chained
     template <typename T, size_t N, typename AFTER, typename OBJECT>
     CloneContext& InsertAfter(const utils::Vector<T, N>& vector,
@@ -456,7 +497,27 @@ class CloneContext {
             return *this;
         }
 
-        list_transforms_.Edit(&vector).insert_after_.GetOrZero(after).Push(object);
+        list_transforms_.Edit(&vector).insert_after_.GetOrZero(after).Push(
+            [object] { return object; });
+        return *this;
+    }
+
+    /// Inserts a lazily created object after @p after whenever @p vector is cloned.
+    /// @param vector the vector in #src
+    /// @param after a pointer to the object in #src
+    /// @param builder the builder of the object in #dst that will be inserted after any occurrence
+    /// of the clone of @p after
+    /// @returns this CloneContext so calls can be chained
+    template <typename T,
+              size_t N,
+              typename AFTER,
+              typename BUILDER,
+              typename _ = std::enable_if_t<!std::is_pointer_v<std::decay_t<BUILDER>>>>
+    CloneContext& InsertAfter(const utils::Vector<T, N>& vector,
+                              const AFTER* after,
+                              BUILDER&& builder) {
+        list_transforms_.Edit(&vector).insert_after_.GetOrZero(after).Push(
+            std::forward<BUILDER>(builder));
         return *this;
     }
 
@@ -487,7 +548,7 @@ class CloneContext {
     };
 
     /// A vector of const Cloneable*
-    using CloneableList = utils::Vector<const Cloneable*, 4>;
+    using CloneableBuilderList = utils::Vector<std::function<const Cloneable*()>, 4>;
 
     /// Transformations to be applied to a list (vector)
     struct ListTransforms {
@@ -495,20 +556,20 @@ class CloneContext {
         utils::Hashset<const Cloneable*, 4> remove_;
 
         /// A list of objects in #dst to insert before any others when the vector is cloned.
-        CloneableList insert_front_;
+        CloneableBuilderList insert_front_;
 
         /// A list of objects in #dst to insert after all others when the vector is cloned.
-        CloneableList insert_back_;
+        CloneableBuilderList insert_back_;
 
         /// A map of object in #src to the list of cloned objects in #dst.
         /// Clone(const utils::Vector<T*>& v) will use this to insert the map-value
         /// list into the target vector before cloning and inserting the map-key.
-        utils::Hashmap<const Cloneable*, CloneableList, 4> insert_before_;
+        utils::Hashmap<const Cloneable*, CloneableBuilderList, 4> insert_before_;
 
         /// A map of object in #src to the list of cloned objects in #dst.
         /// Clone(const utils::Vector<T*>& v) will use this to insert the map-value
         /// list into the target vector after cloning and inserting the map-key.
-        utils::Hashmap<const Cloneable*, CloneableList, 4> insert_after_;
+        utils::Hashmap<const Cloneable*, CloneableBuilderList, 4> insert_after_;
     };
 
     CloneContext(const CloneContext&) = delete;

src/tint/clone_context_test.cc

@@ -430,6 +430,39 @@ TEST_F(CloneContextNodeTest, CloneWithInsertFront) {
     EXPECT_EQ(cloned_root->vec[3]->name, cloned.Symbols().Get("c"));
 }
 
+TEST_F(CloneContextNodeTest, CloneWithInsertFrontFunction) {
+    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<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)
+            .InsertFront(original_root->vec,
+                         [&] { return a.Create<Node>(cloned.Symbols().New("insertion")); })
+            .Clone(original_root);
+
+    EXPECT_EQ(cloned_root->vec.Length(), 4u);
+
+    EXPECT_NE(cloned_root->vec[0], cloned_root->a);
+    EXPECT_NE(cloned_root->vec[1], cloned_root->b);
+    EXPECT_NE(cloned_root->vec[2], cloned_root->c);
+
+    EXPECT_EQ(cloned_root->name, cloned.Symbols().Get("root"));
+    EXPECT_EQ(cloned_root->vec[0]->name, cloned.Symbols().Get("insertion"));
+    EXPECT_EQ(cloned_root->vec[1]->name, cloned.Symbols().Get("a"));
+    EXPECT_EQ(cloned_root->vec[2]->name, cloned.Symbols().Get("b"));
+    EXPECT_EQ(cloned_root->vec[3]->name, cloned.Symbols().Get("c"));
+}
+
 TEST_F(CloneContextNodeTest, CloneWithInsertFront_Empty) {
     Allocator a;
 
@@ -451,6 +484,28 @@ TEST_F(CloneContextNodeTest, CloneWithInsertFront_Empty) {
     EXPECT_EQ(cloned_root->vec[0]->name, cloned.Symbols().Get("insertion"));
 }
 
+TEST_F(CloneContextNodeTest, CloneWithInsertFront_Empty_Function) {
+    Allocator a;
+
+    ProgramBuilder builder;
+    auto* original_root = a.Create<Node>(builder.Symbols().Register("root"));
+    original_root->vec.Clear();
+    Program original(std::move(builder));
+
+    ProgramBuilder cloned;
+
+    auto* cloned_root =
+        CloneContext(&cloned, &original)
+            .InsertFront(original_root->vec,
+                         [&] { return a.Create<Node>(cloned.Symbols().New("insertion")); })
+            .Clone(original_root);
+
+    EXPECT_EQ(cloned_root->vec.Length(), 1u);
+
+    EXPECT_EQ(cloned_root->name, cloned.Symbols().Get("root"));
+    EXPECT_EQ(cloned_root->vec[0]->name, cloned.Symbols().Get("insertion"));
+}
+
 TEST_F(CloneContextNodeTest, CloneWithInsertBack) {
     Allocator a;
 
@@ -479,6 +534,35 @@ TEST_F(CloneContextNodeTest, CloneWithInsertBack) {
     EXPECT_EQ(cloned_root->vec[3]->name, cloned.Symbols().Get("insertion"));
 }
 
+TEST_F(CloneContextNodeTest, CloneWithInsertBack_Function) {
+    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<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)
+            .InsertBack(original_root->vec,
+                        [&] { return a.Create<Node>(cloned.Symbols().New("insertion")); })
+            .Clone(original_root);
+
+    EXPECT_EQ(cloned_root->vec.Length(), 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, CloneWithInsertBack_Empty) {
     Allocator a;
 
@@ -500,6 +584,28 @@ TEST_F(CloneContextNodeTest, CloneWithInsertBack_Empty) {
     EXPECT_EQ(cloned_root->vec[0]->name, cloned.Symbols().Get("insertion"));
 }
 
+TEST_F(CloneContextNodeTest, CloneWithInsertBack_Empty_Function) {
+    Allocator a;
+
+    ProgramBuilder builder;
+    auto* original_root = a.Create<Node>(builder.Symbols().Register("root"));
+    original_root->vec.Clear();
+    Program original(std::move(builder));
+
+    ProgramBuilder cloned;
+
+    auto* cloned_root =
+        CloneContext(&cloned, &original)
+            .InsertBack(original_root->vec,
+                        [&] { return a.Create<Node>(cloned.Symbols().New("insertion")); })
+            .Clone(original_root);
+
+    EXPECT_EQ(cloned_root->vec.Length(), 1u);
+
+    EXPECT_EQ(cloned_root->name, cloned.Symbols().Get("root"));
+    EXPECT_EQ(cloned_root->vec[0]->name, cloned.Symbols().Get("insertion"));
+}
+
 TEST_F(CloneContextNodeTest, CloneWithInsertFrontAndBack_Empty) {
     Allocator a;
 
@@ -509,8 +615,8 @@ TEST_F(CloneContextNodeTest, CloneWithInsertFrontAndBack_Empty) {
     Program original(std::move(builder));
 
     ProgramBuilder cloned;
-    auto* insertion_front = a.Create<Node>(cloned.Symbols().New("insertion_front"));
     auto* insertion_back = a.Create<Node>(cloned.Symbols().New("insertion_back"));
+    auto* insertion_front = a.Create<Node>(cloned.Symbols().New("insertion_front"));
 
     auto* cloned_root = CloneContext(&cloned, &original)
                             .InsertBack(original_root->vec, insertion_back)
@@ -524,6 +630,31 @@ TEST_F(CloneContextNodeTest, CloneWithInsertFrontAndBack_Empty) {
     EXPECT_EQ(cloned_root->vec[1]->name, cloned.Symbols().Get("insertion_back"));
 }
 
+TEST_F(CloneContextNodeTest, CloneWithInsertFrontAndBack_Empty_Function) {
+    Allocator a;
+
+    ProgramBuilder builder;
+    auto* original_root = a.Create<Node>(builder.Symbols().Register("root"));
+    original_root->vec.Clear();
+    Program original(std::move(builder));
+
+    ProgramBuilder cloned;
+
+    auto* cloned_root =
+        CloneContext(&cloned, &original)
+            .InsertBack(original_root->vec,
+                        [&] { return a.Create<Node>(cloned.Symbols().New("insertion_back")); })
+            .InsertFront(original_root->vec,
+                         [&] { return a.Create<Node>(cloned.Symbols().New("insertion_front")); })
+            .Clone(original_root);
+
+    EXPECT_EQ(cloned_root->vec.Length(), 2u);
+
+    EXPECT_EQ(cloned_root->name, cloned.Symbols().Get("root"));
+    EXPECT_EQ(cloned_root->vec[0]->name, cloned.Symbols().Get("insertion_front"));
+    EXPECT_EQ(cloned_root->vec[1]->name, cloned.Symbols().Get("insertion_back"));
+}
+
 TEST_F(CloneContextNodeTest, CloneWithInsertBefore) {
     Allocator a;
 
@@ -552,6 +683,35 @@ TEST_F(CloneContextNodeTest, CloneWithInsertBefore) {
     EXPECT_EQ(cloned_root->vec[3]->name, cloned.Symbols().Get("c"));
 }
 
+TEST_F(CloneContextNodeTest, CloneWithInsertBefore_Function) {
+    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<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)
+            .InsertBefore(original_root->vec, original_root->vec[1],
+                          [&] { return a.Create<Node>(cloned.Symbols().New("insertion")); })
+            .Clone(original_root);
+
+    EXPECT_EQ(cloned_root->vec.Length(), 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("insertion"));
+    EXPECT_EQ(cloned_root->vec[2]->name, cloned.Symbols().Get("b"));
+    EXPECT_EQ(cloned_root->vec[3]->name, cloned.Symbols().Get("c"));
+}
+
 TEST_F(CloneContextNodeTest, CloneWithInsertAfter) {
     Allocator a;
 
@@ -580,6 +740,35 @@ TEST_F(CloneContextNodeTest, CloneWithInsertAfter) {
     EXPECT_EQ(cloned_root->vec[3]->name, cloned.Symbols().Get("c"));
 }
 
+TEST_F(CloneContextNodeTest, CloneWithInsertAfter_Function) {
+    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<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)
+            .InsertAfter(original_root->vec, original_root->vec[1],
+                         [&] { return a.Create<Node>(cloned.Symbols().New("insertion")); })
+            .Clone(original_root);
+
+    EXPECT_EQ(cloned_root->vec.Length(), 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, CloneWithInsertAfterInVectorNodeClone) {
     Allocator a;
 
@@ -612,6 +801,38 @@ TEST_F(CloneContextNodeTest, CloneWithInsertAfterInVectorNodeClone) {
     EXPECT_EQ(cloned_root->vec[3]->name, cloned.Symbols().Get("c"));
 }
 
+TEST_F(CloneContextNodeTest, CloneWithInsertAfterInVectorNodeClone_Function) {
+    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) {
+        ctx.InsertAfter(original_root->vec, r,
+                        [&] { return a.Create<Node>(cloned.Symbols().New("insertion")); });
+        return nullptr;
+    });
+
+    auto* cloned_root = ctx.Clone(original_root);
+
+    EXPECT_EQ(cloned_root->vec.Length(), 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;
 
@@ -644,6 +865,38 @@ TEST_F(CloneContextNodeTest, CloneWithInsertBackInVectorNodeClone) {
     EXPECT_EQ(cloned_root->vec[3]->name, cloned.Symbols().Get("insertion"));
 }
 
+TEST_F(CloneContextNodeTest, CloneWithInsertBackInVectorNodeClone_Function) {
+    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*/) {
+        ctx.InsertBack(original_root->vec,
+                       [&] { return a.Create<Node>(cloned.Symbols().New("insertion")); });
+        return nullptr;
+    });
+
+    auto* cloned_root = ctx.Clone(original_root);
+
+    EXPECT_EQ(cloned_root->vec.Length(), 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;
 
@@ -676,6 +929,38 @@ TEST_F(CloneContextNodeTest, CloneWithInsertBeforeAndAfterRemoved) {
     EXPECT_EQ(cloned_root->vec[3]->name, cloned.Symbols().Get("c"));
 }
 
+TEST_F(CloneContextNodeTest, CloneWithInsertBeforeAndAfterRemoved_Function) {
+    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<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)
+            .InsertBefore(original_root->vec, original_root->vec[1],
+                          [&] { return a.Create<Node>(cloned.Symbols().New("insertion_before")); })
+            .InsertAfter(original_root->vec, original_root->vec[1],
+                         [&] { return a.Create<Node>(cloned.Symbols().New("insertion_after")); })
+            .Remove(original_root->vec, original_root->vec[1])
+            .Clone(original_root);
+
+    EXPECT_EQ(cloned_root->vec.Length(), 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("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"));
+}
+
 TEST_F(CloneContextNodeTest, CloneIntoSameBuilder) {
     ProgramBuilder builder;
     CloneContext ctx(&builder);