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"
@@ -38,23 +40,21 @@ RootShadowNode::Unshared ReanimatedCommitHook::shadowTreeWillCommit(
   // ShadowTree not commited by Reanimated, apply updates from PropsRegistry
 
   auto rootNode = newRootShadowNode->ShadowNode::clone(ShadowNodeFragment{});
+  std::unordered_map<
+      const ShadowNodeFamily *,
+      std::vector<std::shared_ptr<RawProps>>>
+      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].push_back(std::make_shared<RawProps>(props));
         });
+
+    rootNode = std::static_pointer_cast<RootShadowNode>(
+        cloneShadowTreeWithNewProps(rootNode, propsMap));
   }
 
   // If the commit comes from React Native then skip one commit from Reanimated

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

@@ -6,60 +6,67 @@
 
 namespace reanimated {
 
-ShadowNode::Unshared cloneShadowTreeWithNewProps(
-    const ShadowNode::Shared &oldRootNode,
-    const ShadowNodeFamily &family,
-    RawProps &&rawProps) {
-  // adapted from ShadowNode::cloneTree
+ShadowNode::Unshared cloneShadowTreeWithNewPropsRecursive(
+    std::unordered_map<const ShadowNodeFamily *, std::vector<int>> &childrenMap,
+    const ShadowNode::Shared &shadowNode,
+    std::unordered_map<
+        const ShadowNodeFamily *,
+        std::vector<std::shared_ptr<RawProps>>> &propsMap) {
+  auto family = &shadowNode->getFamily();
+  auto children = shadowNode->getChildren();
+  auto &affectedChildren = childrenMap[family];
+
+  for (auto index : affectedChildren) {
+    children[index] = cloneShadowTreeWithNewPropsRecursive(
+        childrenMap, children[index], propsMap);
+  }
 
-  auto ancestors = family.getAncestors(*oldRootNode);
+  Props::Shared newProps = nullptr;
 
-  if (ancestors.empty()) {
-    return ShadowNode::Unshared{nullptr};
+  if (propsMap.contains(family)) {
+    PropsParserContext propsParserContext{
+        shadowNode->getSurfaceId(), *shadowNode->getContextContainer()};
+    newProps = shadowNode->getProps();
+    for (auto &props : propsMap[family]) {
+      newProps = shadowNode->getComponentDescriptor().cloneProps(
+          propsParserContext, newProps, std::move(*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);
-    }
+  auto result = shadowNode->clone(
+      {newProps ? newProps : ShadowNodeFragment::propsPlaceholder(),
+       std::make_shared<ShadowNode::ListOfShared>(children),
+       shadowNode->getState()});
 
-    children[childIndex] = newChildNode;
+  return result;
+}
 
-    newChildNode = parentNode.clone(
-        {ShadowNodeFragment::propsPlaceholder(),
-         std::make_shared<ShadowNode::ListOfShared>(children),
-         parentNode.getState()});
+ShadowNode::Unshared cloneShadowTreeWithNewProps(
+    const ShadowNode::Shared &oldRootNode,
+    std::unordered_map<
+        const ShadowNodeFamily *,
+        std::vector<std::shared_ptr<RawProps>>> &propsMap) {
+  std::unordered_map<const ShadowNodeFamily *, std::vector<int>> childrenMap;
+
+  for (auto &[family, _] : propsMap) {
+    auto ancestors = family->getAncestors(*oldRootNode);
+
+    for (auto it = ancestors.rbegin(); it != ancestors.rend(); ++it) {
+      auto &parentNode = it->first.get();
+      auto index = it->second;
+      auto parentFamily = &parentNode.getFamily();
+      auto &affectedChildren = childrenMap[parentFamily];
+
+      affectedChildren.push_back(index);
+
+      if (affectedChildren.size() > 1) {
+        break;
+      }
+    }
   }
 
-  return std::const_pointer_cast<ShadowNode>(newChildNode);
+  return cloneShadowTreeWithNewPropsRecursive(
+      childrenMap, oldRootNode, propsMap);
 }
 
 } // namespace reanimated

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

@@ -5,7 +5,8 @@
 #include <react/renderer/uimanager/UIManager.h>
 
 #include <memory>
-#include <set>
+#include <unordered_map>
+#include <vector>
 
 using namespace facebook;
 using namespace react;
@@ -14,8 +15,9 @@ namespace reanimated {
 
 ShadowNode::Unshared cloneShadowTreeWithNewProps(
     const ShadowNode::Shared &oldRootNode,
-    const ShadowNodeFamily &family,
-    RawProps &&rawProps);
+    std::unordered_map<
+        const ShadowNodeFamily *,
+        std::vector<std::shared_ptr<RawProps>>> &propsMap);
 
 } // namespace reanimated
 

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

@@ -720,12 +720,15 @@ void NativeReanimatedModule::performOperations() {
     shadowTree.commit(
         [&](RootShadowNode const &oldRootShadowNode)
             -> RootShadowNode::Unshared {
+          std::unordered_map<
+              const ShadowNodeFamily *,
+              std::vector<std::shared_ptr<RawProps>>>
+              propsMap;
           auto rootNode =
               oldRootShadowNode.ShadowNode::clone(ShadowNodeFragment{});
-
           for (const auto &[shadowNode, props] : copiedOperationsQueue) {
-            const ShadowNodeFamily &family = shadowNode->getFamily();
-            react_native_assert(family.getSurfaceId() == surfaceId_);
+            auto family = &shadowNode->getFamily();
+            propsMap[family].push_back(std::make_shared<RawProps>(rt, *props));
 
 #if REACT_NATIVE_MINOR_VERSION >= 73
             // Fix for catching nullptr returned from commit hook was
@@ -736,22 +739,9 @@ 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 std::static_pointer_cast<RootShadowNode>(
+              cloneShadowTreeWithNewProps(rootNode, propsMap));
         },
         { /* .enableStateReconciliation = */
           false,