diff --git a/ReactCommon/fabric/components/root/RootShadowNode.cpp b/ReactCommon/fabric/components/root/RootShadowNode.cpp
index 363b383f4a4..599b41138d5 100644
--- a/ReactCommon/fabric/components/root/RootShadowNode.cpp
+++ b/ReactCommon/fabric/components/root/RootShadowNode.cpp
@@ -49,14 +49,21 @@ RootShadowNode::Unshared RootShadowNode::clone(
 }
 
 RootShadowNode::Unshared RootShadowNode::clone(
-    ShadowNode const &oldShadowNode,
-    ShadowNode::Shared const &newShadowNode) const {
-  auto ancestors = oldShadowNode.getAncestors(*this);
+    ShadowNode const &shadowNode,
+    std::function<ShadowNode::Unshared(ShadowNode const &oldShadowNode)>
+        callback) const {
+  auto ancestors = shadowNode.getAncestors(*this);
 
   if (ancestors.size() == 0) {
     return RootShadowNode::Unshared{nullptr};
   }
 
+  auto &parent = ancestors.back();
+  auto &oldShadowNode = parent.first.get().getChildren().at(parent.second);
+
+  assert(ShadowNode::sameFamily(shadowNode, *oldShadowNode));
+  auto newShadowNode = callback(*oldShadowNode);
+
   auto childNode = newShadowNode;
 
   for (auto it = ancestors.rbegin(); it != ancestors.rend(); ++it) {
diff --git a/ReactCommon/fabric/components/root/RootShadowNode.h b/ReactCommon/fabric/components/root/RootShadowNode.h
index c880f30da45..9899a6e170e 100644
--- a/ReactCommon/fabric/components/root/RootShadowNode.h
+++ b/ReactCommon/fabric/components/root/RootShadowNode.h
@@ -47,14 +47,17 @@ class RootShadowNode final
       LayoutContext const &layoutContext) const;
 
   /*
-   * Clones the node replacing a given old shadow node with a new one in the
-   * tree by cloning all nodes on the path to the root node and then complete
-   * the tree. Returns `nullptr` if the operation cannot be finished
-   * successfully.
+   * Clones the node (and partially the tree starting from the node) by
+   * replacing a `oldShadowNode` (which corresponds to a given `shadowNode`)
+   * with a node that `callback` returns. `oldShadowNode` might not be the same
+   * as `shadowNode` but they must share the same family.
+   *
+   * Returns `nullptr` if the operation cannot be performed successfully.
    */
   RootShadowNode::Unshared clone(
-      ShadowNode const &oldShadowNode,
-      ShadowNode::Shared const &newShadowNode) const;
+      ShadowNode const &shadowNode,
+      std::function<ShadowNode::Unshared(ShadowNode const &oldShadowNode)>
+          callback) const;
 
  private:
   using YogaLayoutableShadowNode::layout;
diff --git a/ReactCommon/fabric/uimanager/UIManager.cpp b/ReactCommon/fabric/uimanager/UIManager.cpp
index b7a41c9dd64..a0d156079ae 100644
--- a/ReactCommon/fabric/uimanager/UIManager.cpp
+++ b/ReactCommon/fabric/uimanager/UIManager.cpp
@@ -131,17 +131,20 @@ void UIManager::setNativeProps(
 
   auto &componentDescriptor = shadowNode.getComponentDescriptor();
   auto props = componentDescriptor.cloneProps(shadowNode.getProps(), rawProps);
-  auto newShadowNode = shadowNode.clone({
-      /* .tag = */ ShadowNodeFragment::tagPlaceholder(),
-      /* .surfaceId = */ ShadowNodeFragment::surfaceIdPlaceholder(),
-      /* .props = */ props,
-  });
 
   shadowTreeRegistry_.visit(
       shadowNode.getSurfaceId(), [&](ShadowTree const &shadowTree) {
         shadowTree.tryCommit(
             [&](RootShadowNode::Shared const &oldRootShadowNode) {
-              return oldRootShadowNode->clone(shadowNode, newShadowNode);
+              return oldRootShadowNode->clone(
+                  shadowNode, [&](ShadowNode const &oldShadowNode) {
+                    return oldShadowNode.clone({
+                        /* .tag = */ ShadowNodeFragment::tagPlaceholder(),
+                        /* .surfaceId = */
+                        ShadowNodeFragment::surfaceIdPlaceholder(),
+                        /* .props = */ props,
+                    });
+                  });
             });
       });
 }
@@ -181,22 +184,27 @@ void UIManager::updateState(
   auto &componentDescriptor = shadowNode.getComponentDescriptor();
   auto state =
       componentDescriptor.createState(shadowNode.getState(), rawStateData);
-  auto newShadowNode = shadowNode.clone({
-      /* .tag = */ ShadowNodeFragment::tagPlaceholder(),
-      /* .surfaceId = */ ShadowNodeFragment::surfaceIdPlaceholder(),
-      /* .props = */ ShadowNodeFragment::propsPlaceholder(),
-      /* .eventEmitter = */ ShadowNodeFragment::eventEmitterPlaceholder(),
-      /* .children = */ ShadowNodeFragment::childrenPlaceholder(),
-      /* .localData = */ ShadowNodeFragment::localDataPlaceholder(),
-      /* .state = */ state,
-  });
 
   shadowTreeRegistry_.visit(
-      shadowNode.getSurfaceId(), [&](const ShadowTree &shadowTree) {
-        shadowTree.tryCommit(
-            [&](RootShadowNode::Shared const &oldRootShadowNode) {
-              return oldRootShadowNode->clone(shadowNode, newShadowNode);
-            });
+      shadowNode.getSurfaceId(), [&](ShadowTree const &shadowTree) {
+        shadowTree.tryCommit([&](RootShadowNode::Shared const
+                                     &oldRootShadowNode) {
+          return oldRootShadowNode->clone(
+              shadowNode, [&](ShadowNode const &oldShadowNode) {
+                return oldShadowNode.clone({
+                    /* .tag = */ ShadowNodeFragment::tagPlaceholder(),
+                    /* .surfaceId = */
+                    ShadowNodeFragment::surfaceIdPlaceholder(),
+                    /* .props = */ ShadowNodeFragment::propsPlaceholder(),
+                    /* .eventEmitter = */
+                    ShadowNodeFragment::eventEmitterPlaceholder(),
+                    /* .children = */ ShadowNodeFragment::childrenPlaceholder(),
+                    /* .localData = */
+                    ShadowNodeFragment::localDataPlaceholder(),
+                    /* .state = */ state,
+                });
+              });
+        });
       });
 }