Optimize Shadow Tree cloning

packages/react-native-reanimated/Common/cpp/Fabric/ReanimatedCommitHook.cpp

@@ -1,6 +1,8 @@
 #ifdef RCT_NEW_ARCH_ENABLED
 
 #include <react/renderer/core/ComponentDescriptor.h>
+#include <unordered_map>
+#include <vector>
 
 #include "ReanimatedCommitHook.h"
 #include "ReanimatedCommitMarker.h"
@@ -37,24 +39,18 @@ RootShadowNode::Unshared ReanimatedCommitHook::shadowTreeWillCommit(
 
   // ShadowTree not commited by Reanimated, apply updates from PropsRegistry
 
-  auto rootNode = newRootShadowNode->ShadowNode::clone(ShadowNodeFragment{});
+  RootShadowNode::Unshared rootNode = newRootShadowNode;
+  PropsMap propsMap;
 
   {
     auto lock = propsRegistry_->createLock();
 
     propsRegistry_->for_each(
         [&](const ShadowNodeFamily &family, const folly::dynamic &props) {
-          auto newRootNode =
-              cloneShadowTreeWithNewProps(rootNode, family, RawProps(props));
-
-          if (newRootNode == nullptr) {
-            // this happens when React removed the component but Reanimated
-            // still tries to animate it, let's skip update for this specific
-            // component
-            return;
-          }
-          rootNode = newRootNode;
+          propsMap[&family].emplace_back(props);
         });
+
+    rootNode = cloneShadowTreeWithNewProps(*rootNode, propsMap);
   }
 
   // If the commit comes from React Native then skip one commit from Reanimated
@@ -63,7 +59,7 @@ RootShadowNode::Unshared ReanimatedCommitHook::shadowTreeWillCommit(
   // applied in ReanimatedCommitHook by iterating over PropsRegistry.
   propsRegistry_->pleaseSkipReanimatedCommit();
 
-  return std::static_pointer_cast<RootShadowNode>(rootNode);
+  return rootNode;
 }
 
 } // namespace reanimated

packages/react-native-reanimated/Common/cpp/Fabric/ShadowTreeCloner.cpp

@@ -1,65 +1,73 @@
 #ifdef RCT_NEW_ARCH_ENABLED
 
+#include <ranges>
 #include <utility>
 
 #include "ShadowTreeCloner.h"
 
 namespace reanimated {
 
-ShadowNode::Unshared cloneShadowTreeWithNewProps(
-    const ShadowNode::Shared &oldRootNode,
-    const ShadowNodeFamily &family,
-    RawProps &&rawProps) {
-  // adapted from ShadowNode::cloneTree
+ShadowNode::Unshared cloneShadowTreeWithNewPropsRecursive(
+    const ShadowNode &shadowNode,
+    const ChildrenMap &childrenMap,
+    const PropsMap &propsMap) {
+  const auto family = &shadowNode.getFamily();
+  const auto affectedChildrenIt = childrenMap.find(family);
+  const auto propsIt = propsMap.find(family);
+  auto children = shadowNode.getChildren();
+
+  if (affectedChildrenIt != childrenMap.end()) {
+    for (const auto index : affectedChildrenIt->second) {
+      children[index] = cloneShadowTreeWithNewPropsRecursive(
+          *children[index], childrenMap, propsMap);
+    }
+  }
 
-  auto ancestors = family.getAncestors(*oldRootNode);
+  Props::Shared newProps = nullptr;
 
-  if (ancestors.empty()) {
-    return ShadowNode::Unshared{nullptr};
+  if (propsIt != propsMap.end()) {
+    PropsParserContext propsParserContext{
+        shadowNode.getSurfaceId(), *shadowNode.getContextContainer()};
+    newProps = shadowNode.getProps();
+    for (const auto &props : propsIt->second) {
+      newProps = shadowNode.getComponentDescriptor().cloneProps(
+          propsParserContext, newProps, RawProps(props));
+    }
   }
 
-  auto &parent = ancestors.back();
-  auto &source = parent.first.get().getChildren().at(parent.second);
-
-  PropsParserContext propsParserContext{
-      source->getSurfaceId(), *source->getContextContainer()};
-  const auto props = source->getComponentDescriptor().cloneProps(
-      propsParserContext, source->getProps(), std::move(rawProps));
-
-  auto newChildNode = source->clone(
-      {/* .props = */ props,
-       ShadowNodeFragment::childrenPlaceholder(),
-       source->getState()});
-
-  for (auto it = ancestors.rbegin(); it != ancestors.rend(); ++it) {
-    auto &parentNode = it->first.get();
-    auto childIndex = it->second;
-
-    auto children = parentNode.getChildren();
-    const auto &oldChildNode = *children.at(childIndex);
-    react_native_assert(ShadowNode::sameFamily(oldChildNode, *newChildNode));
-
-    if (!parentNode.getSealed()) {
-      // Optimization: if a ShadowNode is unsealed, we can directly update its
-      // children instead of cloning the whole path to the root node.
-      auto &parentNodeNonConst = const_cast<ShadowNode &>(parentNode);
-      parentNodeNonConst.replaceChild(oldChildNode, newChildNode, childIndex);
-      // Unfortunately, `replaceChild` does not update Yoga nodes, so we need to
-      // update them manually here.
-      static_cast<YogaLayoutableShadowNode *>(&parentNodeNonConst)
-          ->updateYogaChildren();
-      return std::const_pointer_cast<ShadowNode>(oldRootNode);
-    }
+  const auto result = shadowNode.clone(
+      {newProps ? newProps : ShadowNodeFragment::propsPlaceholder(),
+       std::make_shared<ShadowNode::ListOfShared>(children),
+       shadowNode.getState()});
+
+  return result;
+}
+
+RootShadowNode::Unshared cloneShadowTreeWithNewProps(
+    const RootShadowNode &oldRootNode,
+    const PropsMap &propsMap) {
+  ChildrenMap childrenMap;
+
+  for (auto &[family, _] : propsMap) {
+    const auto ancestors = family->getAncestors(oldRootNode);
 
-    children[childIndex] = newChildNode;
+    for (const auto &[parentNode, index] :
+         std::ranges::reverse_view(ancestors)) {
+      const auto parentFamily = &parentNode.get().getFamily();
+      auto &affectedChildren = childrenMap[parentFamily];
 
-    newChildNode = parentNode.clone(
-        {ShadowNodeFragment::propsPlaceholder(),
-         std::make_shared<ShadowNode::ListOfShared>(children),
-         parentNode.getState()});
+      if (affectedChildren.contains(index)) {
+        continue;
+      }
+
+      affectedChildren.insert(index);
+    }
   }
 
-  return std::const_pointer_cast<ShadowNode>(newChildNode);
+  // This cast is safe, because this function returns a clone
+  // of the oldRootNode, which is an instance of RootShadowNode
+  return std::static_pointer_cast<RootShadowNode>(
+      cloneShadowTreeWithNewPropsRecursive(oldRootNode, childrenMap, propsMap));
 }
 
 } // namespace reanimated

packages/react-native-reanimated/Common/cpp/Fabric/ShadowTreeCloner.h

@@ -4,18 +4,23 @@
 #include <react/renderer/core/PropsParserContext.h>
 #include <react/renderer/uimanager/UIManager.h>
 
+#include <type_traits>
 #include <memory>
-#include <set>
+#include <unordered_map>
+#include <unordered_set>
+#include <vector>
 
 using namespace facebook;
 using namespace react;
 
 namespace reanimated {
 
-ShadowNode::Unshared cloneShadowTreeWithNewProps(
-    const ShadowNode::Shared &oldRootNode,
-    const ShadowNodeFamily &family,
-    RawProps &&rawProps);
+using PropsMap = std::unordered_map<const ShadowNodeFamily *, std::vector<RawProps>>;
+using ChildrenMap = std::unordered_map<const ShadowNodeFamily *, std::unordered_set<int>>;
+
+RootShadowNode::Unshared cloneShadowTreeWithNewProps(
+    const RootShadowNode &oldRootNode,
+    const PropsMap &propsMap);
 
 } // namespace reanimated
 

packages/react-native-reanimated/Common/cpp/NativeModules/NativeReanimatedModule.cpp

@@ -720,12 +720,11 @@ void NativeReanimatedModule::performOperations() {
     shadowTree.commit(
         [&](RootShadowNode const &oldRootShadowNode)
             -> RootShadowNode::Unshared {
-          auto rootNode =
-              oldRootShadowNode.ShadowNode::clone(ShadowNodeFragment{});
-
-          for (const auto &[shadowNode, props] : copiedOperationsQueue) {
-            const ShadowNodeFamily &family = shadowNode->getFamily();
-            react_native_assert(family.getSurfaceId() == surfaceId_);
+          PropsMap propsMap;
+          for (auto &[shadowNode, props] : copiedOperationsQueue) {
+            auto family = &shadowNode->getFamily();
+            react_native_assert(family->getSurfaceId() == surfaceId_);
+            propsMap[family].emplace_back(rt, std::move(*props));
 
 #if REACT_NATIVE_MINOR_VERSION >= 73
             // Fix for catching nullptr returned from commit hook was
@@ -736,22 +735,8 @@ void NativeReanimatedModule::performOperations() {
               return nullptr;
             }
 #endif
-
-            auto newRootNode = cloneShadowTreeWithNewProps(
-                rootNode, family, RawProps(rt, *props));
-
-            if (newRootNode == nullptr) {
-              // this happens when React removed the component but Reanimated
-              // still tries to animate it, let's skip update for this
-              // specific component
-              continue;
-            }
-            rootNode = newRootNode;
           }
-
-          auto newRoot = std::static_pointer_cast<RootShadowNode>(rootNode);
-
-          return newRoot;
+          return cloneShadowTreeWithNewProps(oldRootShadowNode, propsMap);
         },
         { /* .enableStateReconciliation = */
           false,