[Fabric] Optimize cloning ShadowTree when animating layout props (sof…

…tware-mansion#3369)

## Description

This PR improves the implementation of ShadowTree cloning algorithm in `NativeReanimatedModule::performOperations`.

Previously, `ShadowNode::cloneTree` was called for each update, so effectively the root node was cloned $n$ times.

In this PR I've proposed an algorithm for solving this task more optimally. The main idea is that unsealed ShadowNodes can be still updated. In particular, their children can be replaced using `ShadowNode::replaceChild`. Also, freshly cloned ShadowNodes are unsealed, so each ShadowNode can be cloned only once and then updated.

One important thing to mention is that `replaceChild` does not update `yogaChild_` field in ShadowNode, so we need to invoke `ShadowNode::updateYogaChildren` manually. This function has $O(n)$ time complexity where $n$ denotes the number of children. In order to avoid calling `updateYogaChildren` multiple number of times for the same ShadowNode, the calls are batched in `std::set`. Since in most use-cases there is no more than ~50 props updates, I've decided to use ordered set implementation with $O(\log{n})$ time complexity of insertion in order to save memory as well as avoid the overhead of calculating hash function.

| Before | After |
|:-:|:-:|
| <img width="180" alt="Before" src="https://user-images.githubusercontent.com/20516055/177736298-0ed230dc-edf2-40a5-9a1d-8400c9f58c39.png"> | <img width="159" alt="After" src="https://user-images.githubusercontent.com/20516055/177736518-fd9eeda6-025c-45f1-a636-b4a0665b0c23.png"> |

<details>
<summary><strong>Execution time measurements</strong></summary>

| nested view levels | chessboard size | old algorithm [μs] | new algorithm [μs] | speedup [x] |
|:-:|:-:|:-:|:-:|:-:|
| 0 | 10x10 | 4008 | 1326 | 3,0 |
| 0 | 20x20 | 15962 | 3620 | 4,4 |
| 0 | 30x30 | 50287 | 5939 | 8,5 |
| 0 | 40x40 | 131410 | 10669 | 12,3 |
| 10 | 10x10 | 105792 | 4058 | 26,1 |
| 10 | 20x20 | 388404 | 5935 | 65,4 |
| 10 | 30x30 | 879948 | 9334 | 94,3 |
| 10 | 40x40 | 1598547 | 15477 | 103,3 |
| 50 | 10x10 | 524744 | 11750 | 44,7 |
| 50 | 20x20 | 2056805 | 17257 | 119,2 |
| 50 | 30x30 | 4657397 | 24588 | 189,4 |
| 50 | 40x40 | 8305191 | 35413 | 234,5 |

</details>

<!--
Description and motivation for this PR.

Inlude Fixes #<number> if this is fixing some issue.

Fixes # .
-->

## Changes

- Implemented `ShadowTreeCloner` for optimal cloning ShadowTree when animating layout props on Fabric
- Used `ShadowTreeCloner` in `NativeReanimatedModule::performOperations`
- Moved creation of `PropsParserContext` to `ShadowTreeCloner`

<!--
Please describe things you've changed here, make a **high level** overview, if change is simple you can omit this section.

For example:

- Added `foo` method which add bouncing animation
- Updated `about.md` docs
- Added caching in CI builds

-->

<!--

## Screenshots / GIFs

Here you can add screenshots / GIFs documenting your change.

You can add before / after section if you're changing some behavior.

### Before

### After

-->

## Test code and steps to reproduce

<!--
Please include code that can be used to test this change and short description how this example should work.
This snippet should be as minimal as possible and ready to be pasted into editor (don't exclude exports or remove "not important" parts of reproduction example)
-->

The following examples can be used to test the algorithm:
- `WidthExample`
- `RefExample`
- `ChessboardExample`
- `NewestShadowNodesRegistryRemoveExample`

## Checklist

- [ ] Included code example that can be used to test this change
- [ ] Updated TS types
- [ ] Added TS types tests
- [ ] Added unit / integration tests
- [ ] Updated documentation
- [ ] Ensured that CI passes

Common/cpp/Fabric/ShadowTreeCloner.cpp

@@ -0,0 +1,92 @@
+#ifdef RCT_NEW_ARCH_ENABLED
+
+#include "ShadowTreeCloner.h"
+#include "FabricUtils.h"
+
+namespace reanimated {
+
+ShadowTreeCloner::ShadowTreeCloner(
+    std::shared_ptr<NewestShadowNodesRegistry> newestShadowNodesRegistry,
+    std::shared_ptr<UIManager> uiManager,
+    SurfaceId surfaceId)
+    : newestShadowNodesRegistry_{newestShadowNodesRegistry},
+      propsParserContext_{
+          surfaceId,
+          *getContextContainerFromUIManager(&*uiManager)} {}
+
+ShadowTreeCloner::~ShadowTreeCloner() {
+#ifdef DEBUG
+  react_native_assert(
+      yogaChildrenUpdates_.empty() &&
+      "Deallocating `ShadowTreeCloner` without calling `updateYogaChildren`.");
+#endif
+}
+
+ShadowNode::Unshared ShadowTreeCloner::cloneWithNewProps(
+    const ShadowNode::Shared &oldRootNode,
+    const ShadowNodeFamily &family,
+    RawProps &&rawProps) {
+  // adapted from ShadowNode::cloneTree
+
+  auto ancestors = family.getAncestors(*oldRootNode);
+
+  if (ancestors.empty()) {
+    return ShadowNode::Unshared{nullptr};
+  }
+
+  auto &parent = ancestors.back();
+  auto &oldShadowNode = parent.first.get().getChildren().at(parent.second);
+
+  const auto newest = newestShadowNodesRegistry_->get(oldShadowNode->getTag());
+
+  const auto &source = newest == nullptr ? oldShadowNode : newest;
+
+  const auto props = source->getComponentDescriptor().cloneProps(
+      propsParserContext_, source->getProps(), rawProps);
+
+  auto newChildNode = source->clone({/* .props = */ props});
+
+  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));
+
+    newestShadowNodesRegistry_->set(newChildNode, parentNode.getTag());
+
+    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);
+      yogaChildrenUpdates_.insert(&parentNodeNonConst);
+      return std::const_pointer_cast<ShadowNode>(oldRootNode);
+    }
+
+    children[childIndex] = newChildNode;
+
+    newChildNode = parentNode.clone({
+        ShadowNodeFragment::propsPlaceholder(),
+        std::make_shared<SharedShadowNodeList>(children),
+    });
+  }
+
+  return std::const_pointer_cast<ShadowNode>(newChildNode);
+}
+
+void ShadowTreeCloner::updateYogaChildren() {
+  // Unfortunately, `replaceChild` does not update Yoga nodes, so we need to
+  // update them manually here.
+  for (ShadowNode *shadowNode : yogaChildrenUpdates_) {
+    static_cast<YogaLayoutableShadowNode *>(shadowNode)->updateYogaChildren();
+  }
+#ifdef DEBUG
+  yogaChildrenUpdates_.clear();
+#endif
+}
+
+} // namespace reanimated
+
+#endif // RCT_NEW_ARCH_ENABLED

Common/cpp/NativeModules/NativeReanimatedModule.cpp

@@ -13,6 +13,7 @@
 #include "FabricUtils.h"
 #include "NewestShadowNodesRegistry.h"
 #include "ReanimatedUIManagerBinding.h"
+#include "ShadowTreeCloner.h"
 #endif
 
 #include "EventHandlerRegistry.h"
@@ -507,56 +508,42 @@ void NativeReanimatedModule::performOperations() {
 
   react_native_assert(uiManager_ != nullptr);
   const auto &shadowTreeRegistry = uiManager_->getShadowTreeRegistry();
-  auto contextContainer = getContextContainerFromUIManager(
-      &*uiManager_); // TODO: use Scheduler::getContextContainer
-  PropsParserContext propsParserContext{surfaceId_, *contextContainer};
   jsi::Runtime &rt = *runtime.get();
 
   shadowTreeRegistry.visit(surfaceId_, [&](ShadowTree const &shadowTree) {
     shadowTree.commit([&](RootShadowNode const &oldRootShadowNode) {
-      // lock once due to performance reasons
-      auto lock = newestShadowNodesRegistry_->createLock();
-
       auto rootNode = oldRootShadowNode.ShadowNode::clone(ShadowNodeFragment{});
 
-      for (const auto &pair : copiedOperationsQueue) {
-        const ShadowNodeFamily &family = pair.first->getFamily();
-        react_native_assert(family.getSurfaceId() == surfaceId_);
-
-        auto newRootNode =
-            rootNode->cloneTree(family, [&](ShadowNode const &oldShadowNode) {
-              const auto newest =
-                  newestShadowNodesRegistry_->get(oldShadowNode.getTag());
+      ShadowTreeCloner shadowTreeCloner{
+          newestShadowNodesRegistry_, uiManager_, surfaceId_};
 
-              const auto &source = newest == nullptr ? oldShadowNode : *newest;
+      {
+        // lock once due to performance reasons
+        auto lock = newestShadowNodesRegistry_->createLock();
 
-              const auto newProps = source.getComponentDescriptor().cloneProps(
-                  propsParserContext,
-                  source.getProps(),
-                  RawProps(rt, *pair.second));
+        for (const auto &pair : copiedOperationsQueue) {
+          const ShadowNodeFamily &family = pair.first->getFamily();
+          react_native_assert(family.getSurfaceId() == surfaceId_);
 
-              return source.clone({/* .props = */ newProps});
-            });
+          auto newRootNode = shadowTreeCloner.cloneWithNewProps(
+              rootNode, family, RawProps(rt, *pair.second));
 
-        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;
+          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;
         }
-        rootNode = newRootNode;
 
-        auto ancestors = family.getAncestors(*rootNode);
-        for (const auto &pair : ancestors) {
-          const auto &parent = pair.first.get();
-          const auto &child = parent.getChildren().at(pair.second);
-          newestShadowNodesRegistry_->set(child, parent.getTag());
+        // remove ShadowNodes and its ancestors from NewestShadowNodesRegistry
+        for (auto tag : copiedTagsToRemove) {
+          newestShadowNodesRegistry_->remove(tag);
         }
       }
 
-      // remove ShadowNodes and its ancestors from NewestShadowNodesRegistry
-      for (auto tag : copiedTagsToRemove) {
-        newestShadowNodesRegistry_->remove(tag);
-      }
+      shadowTreeCloner.updateYogaChildren();
 
       return std::static_pointer_cast<RootShadowNode>(rootNode);
     });

Common/cpp/headers/Fabric/ShadowTreeCloner.h

@@ -0,0 +1,41 @@
+#pragma once
+#ifdef RCT_NEW_ARCH_ENABLED
+
+#include <react/renderer/core/PropsParserContext.h>
+#include <react/renderer/uimanager/UIManager.h>
+
+#include <memory>
+#include <set>
+
+#include "NewestShadowNodesRegistry.h"
+
+using namespace facebook;
+using namespace react;
+
+namespace reanimated {
+
+class ShadowTreeCloner {
+ public:
+  ShadowTreeCloner(
+      std::shared_ptr<NewestShadowNodesRegistry> newestShadowNodesRegistry,
+      std::shared_ptr<UIManager> uiManager,
+      SurfaceId surfaceId);
+
+  ~ShadowTreeCloner();
+
+  ShadowNode::Unshared cloneWithNewProps(
+      const ShadowNode::Shared &oldRootNode,
+      const ShadowNodeFamily &family,
+      RawProps &&rawProps);
+
+  void updateYogaChildren();
+
+ private:
+  PropsParserContext propsParserContext_;
+  std::shared_ptr<NewestShadowNodesRegistry> newestShadowNodesRegistry_;
+  std::set<ShadowNode *> yogaChildrenUpdates_;
+};
+
+} // namespace reanimated
+
+#endif // RCT_NEW_ARCH_ENABLED

android/CMakeLists.txt

@@ -212,6 +212,12 @@ if(${IS_NEW_ARCHITECTURE_ENABLED})
             PATHS ${LIBRN_DIR}
             NO_CMAKE_FIND_ROOT_PATH
     )
+    find_library(
+            RRC_VIEW
+            rrc_view
+            PATHS ${LIBRN_DIR}
+            NO_CMAKE_FIND_ROOT_PATH
+    )
     find_library(
             REACT_RENDER_SCHEDULER
             react_render_scheduler
@@ -307,6 +313,7 @@ if(${IS_NEW_ARCHITECTURE_ENABLED})
             ${REACT_DEBUG}
             ${REACT_RENDER_DEBUG}
             ${RRC_ROOT}
+            ${RRC_VIEW}
             ${FABRICJNI}
             ${REACT_RENDER_SCHEDULER}
     )