Optimize Shadow Tree cloning (#6214)

## Summary

This PR optimizes the algorithm used to clone the Shadow Tree used in
`ReanimatedCommitHook` and `performOperations`.

## Current approach
The current algorithm works as follows. After receiving a batch of
updates, we iterate over the list and apply the changes one by one. To
apply changes we have to:
1. calculate the path from the affected ShadowNode to the root using
`ShadowNodeFamily::getAncestors`
2. traverse the path upwards cloning all the nodes up to the root

This way we unfortunately clone some ShadowNodes multiple times. For
example for a batch of size `n` we will clone the root node `n` times.

Cloning ShadowNodes is expensive, so we had implemented an optimization
- whenever a node is unsealed we would change it in place instead of
cloning it. Unfortunately this didn't work, since the `getSealed` method
always returns `true` in Production mode. This is not a bug, but the
intended behavior, as sealing is only intended to help finding bugs in
the Debug Mode. This still could be salvaged by memoizing which nodes
were already cloned by us, but this approach still wouldn't be perfect,
as modyfing nodes in place is still a heavy operation.

## New approach
To mitigate those issues we split the process into two phases:
1. calculate the subtree of the ShadowTree that contains all the nodes
that we want to update
2. traverse the ShadowTree and clone nodes (that belong to the subtree)
in the (reversed) topological order

By calculating the subtree first we ensure that in the second phase:
1. we traverse only nodes that absolutely have to be traversed
2. we clone only nodes that absolutely have to be cloned
3. we clone every node at most once

With this approach the second phase is performed in the optimal number
of operations.

## Limitations
The current implementation of phase one (building the subtree) is not
optimal. It is implemented by simply calling `getAncestors` on every
node from the batch. This is fortunately not a huge problem, because
cloning had a much heavier impact on the performance. To optimize this
there will have to be some changes done in RN (because the `parent`
field in `ShadowNodeFamily` is private, so traversing the tree upwards
is only possible through `getAncestors`). I hope to soon open a suitable
PR.

## Some examples
I checked the performance of our heavier examples on some devices in the
Release Mode. For the `BokehExample.tsx` the results are:

| Phone  | Example size | Before [FPS] | After [FPS] | 
| -------- | ------- | -------- | ------- |
| iPhone 12 mini  | 200   | 30-40 | 60 |
| One+ A6 | 100   | 10-20 | 30-40 |
| iPhone 15 Pro  | 250 | 30-40 | 120 |
| Samsung Galaxy S23 | 100 | 55-70 | 120 |

I also tested through Xcode Instruments how much time does the
`performOperations` function take on the same example. Tests were
conducted on the iPhone simulator, but they should give an idea on the
order of the number of operations this function makes (and how fast that
number grows in relation to the example size).

| Example size | Before [ms] | After [ms] | Before/After |
| ------- | -------- | ------- | ------- |
| 1   | 1.95 | 2.1 | 0.92 |
| 20   | 2.4 | 2.1 | 1.14 |
| 100   | 5.3 | 2.3 | 2.3 |
| 250   | 22 | 4 | 5.5 |
| 500 | 77 | 7.7 | 10 |

## Test plan
Check the behavior of examples in the `FabricExample` app. Verify that
heavy examples have improved, while simpler examples have not regressed.

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,