Optimize Shadow Tree cloning

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

@@ -0,0 +1,45 @@
+#pragma once
+#ifdef RCT_NEW_ARCH_ENABLED
+
+#include <react/renderer/uimanager/UIManager.h>
+#include <memory>
+#include <utility>
+
+using namespace facebook;
+using namespace react;
+
+namespace reanimated {
+
+class PropsWrapper {
+ public:
+  virtual inline RawProps getRawProps() = 0;
+  virtual ~PropsWrapper() = default;
+};
+
+class JsiValuePropsWrapper : public PropsWrapper {
+  jsi::Runtime &runtime;
+  std::unique_ptr<jsi::Value> value;
+
+ public:
+  JsiValuePropsWrapper(
+      jsi::Runtime &runtime,
+      std::unique_ptr<jsi::Value> &&value)
+      : runtime(runtime), value(std::move(value)) {}
+  inline RawProps getRawProps() override {
+    return RawProps(runtime, *value);
+  }
+};
+
+struct FollyDynamicPropsWrapper : public PropsWrapper {
+  const folly::dynamic &value;
+
+ public:
+  explicit FollyDynamicPropsWrapper(const folly::dynamic &value) : value(value) {}
+  inline RawProps getRawProps() override {
+    return RawProps(value);
+  }
+};
+
+} // namespace reanimated
+
+#endif // RCT_NEW_ARCH_ENABLED

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,18 @@ RootShadowNode::Unshared ReanimatedCommitHook::shadowTreeWillCommit(
   // ShadowTree not commited by Reanimated, apply updates from PropsRegistry
 
   auto rootNode = newRootShadowNode->ShadowNode::clone(ShadowNodeFragment{});
+  PropsMap<FollyDynamicPropsWrapper> 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(std::make_unique<FollyDynamicPropsWrapper>(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,62 +6,78 @@
 
 namespace reanimated {
 
-ShadowNode::Unshared cloneShadowTreeWithNewProps(
-    const ShadowNode::Shared &oldRootNode,
-    const ShadowNodeFamily &family,
-    RawProps &&rawProps) {
-  // adapted from ShadowNode::cloneTree
+template<Derived T>
+ShadowNode::Unshared cloneShadowTreeWithNewPropsRecursive(
+    const ChildrenMap &childrenMap,
+    const ShadowNode::Shared &shadowNode,
+    const PropsMap<T> &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(
+          childrenMap, children[index], 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, props->getRawProps());
+    }
   }
 
-  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;
+}
+
+template<Derived T>
+ShadowNode::Unshared cloneShadowTreeWithNewProps(
+    const ShadowNode::Shared &oldRootNode,
+    const PropsMap<T> &propsMap) {
+  ChildrenMap childrenMap;
+
+  for (auto &[family, _] : propsMap) {
+    const auto ancestors = family->getAncestors(*oldRootNode);
 
-    children[childIndex] = newChildNode;
+    for (auto it = ancestors.rbegin(); it != ancestors.rend(); ++it) {
+      const auto &parentNode = it->first.get();
+      const auto index = it->second;
+      const auto parentFamily = &parentNode.getFamily();
+      auto &affectedChildren = childrenMap[parentFamily];
 
-    newChildNode = parentNode.clone(
-        {ShadowNodeFragment::propsPlaceholder(),
-         std::make_shared<ShadowNode::ListOfShared>(children),
-         parentNode.getState()});
+      affectedChildren.push_back(index);
+
+      if (affectedChildren.size() > 1) {
+        break;
+      }
+    }
   }
 
-  return std::const_pointer_cast<ShadowNode>(newChildNode);
+  return cloneShadowTreeWithNewPropsRecursive(
+      childrenMap, oldRootNode, propsMap);
 }
 
+template ShadowNode::Unshared cloneShadowTreeWithNewProps(
+    const ShadowNode::Shared &oldRootNode,
+    const PropsMap<JsiValuePropsWrapper> &propsMap);
+
+template ShadowNode::Unshared cloneShadowTreeWithNewProps(
+    const ShadowNode::Shared &oldRootNode,
+    const PropsMap<FollyDynamicPropsWrapper> &propsMap);
+
 } // namespace reanimated
 
 #endif // RCT_NEW_ARCH_ENABLED

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

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

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

@@ -720,12 +720,12 @@ void NativeReanimatedModule::performOperations() {
     shadowTree.commit(
         [&](RootShadowNode const &oldRootShadowNode)
             -> RootShadowNode::Unshared {
+          PropsMap<JsiValuePropsWrapper> propsMap;
           auto rootNode =
               oldRootShadowNode.ShadowNode::clone(ShadowNodeFragment{});
-
-          for (const auto &[shadowNode, props] : copiedOperationsQueue) {
-            const ShadowNodeFamily &family = shadowNode->getFamily();
-            react_native_assert(family.getSurfaceId() == surfaceId_);
+          for (auto &[shadowNode, props] : copiedOperationsQueue) {
+            auto family = &shadowNode->getFamily();
+            propsMap[family].emplace_back(std::make_unique<JsiValuePropsWrapper>(rt, std::move(props)));
 
 #if REACT_NATIVE_MINOR_VERSION >= 73
             // Fix for catching nullptr returned from commit hook was
@@ -736,22 +736,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,