Fix issues with skipping reanimated commit

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

@@ -25,30 +25,33 @@ class PropsRegistry {
 
   void remove(const Tag tag);
 
-  void pleaseSkipReanimatedCommit() {
-    shouldReanimatedSkipCommit_ = true;
+  void pauseReanimatedCommits() {
+    isPaused_ = true;
   }
 
   bool shouldReanimatedSkipCommit() {
-#if REACT_NATIVE_MINOR_VERSION >= 73
-    // In RN 0.73+ we have a mount hook that will properly unset this flag
-    // after a non-Reanimated commit.
-    return shouldReanimatedSkipCommit_;
-#else
-    return shouldReanimatedSkipCommit_.exchange(false);
-#endif
+    return isPaused_;
   }
 
-  void resetReanimatedSkipCommitFlag() {
-    shouldReanimatedSkipCommit_ = false;
+  void unpauseReanimatedCommits() {
+    isPaused_ = false;
+  }
+
+  void pleaseCommitAfterPause() {
+    shouldCommitAfterPause_ = true;
+  }
+
+  bool shouldCommitAfterPause() {
+    return shouldCommitAfterPause_.exchange(false);
   }
 
  private:
   std::unordered_map<Tag, std::pair<ShadowNode::Shared, folly::dynamic>> map_;
 
   mutable std::mutex mutex_; // Protects `map_`.
 
-  std::atomic<bool> shouldReanimatedSkipCommit_;
+  std::atomic<bool> isPaused_;
+  std::atomic<bool> shouldCommitAfterPause_;
 };
 
 } // namespace reanimated

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

@@ -26,12 +26,7 @@ ReanimatedCommitHook::~ReanimatedCommitHook() noexcept {
 RootShadowNode::Unshared ReanimatedCommitHook::shadowTreeWillCommit(
     ShadowTree const &,
     RootShadowNode::Shared const &,
-#if REACT_NATIVE_MINOR_VERSION >= 73
     RootShadowNode::Unshared const &newRootShadowNode) noexcept {
-#else
-    RootShadowNode::Unshared const &newRootShadowNode) const noexcept {
-#endif
-
   auto reaShadowNode =
       std::reinterpret_pointer_cast<ReanimatedCommitShadowNode>(
           newRootShadowNode);
@@ -40,11 +35,12 @@ RootShadowNode::Unshared ReanimatedCommitHook::shadowTreeWillCommit(
     // ShadowTree commited by Reanimated, no need to apply updates from
     // PropsRegistry
     reaShadowNode->unsetReanimatedCommitTrait();
+    reaShadowNode->setReanimatedMountTrait();
     return newRootShadowNode;
   }
 
   // ShadowTree not commited by Reanimated, apply updates from PropsRegistry
-
+  reaShadowNode->unsetReanimatedMountTrait();
   RootShadowNode::Unshared rootNode = newRootShadowNode;
   PropsMap propsMap;
 
@@ -58,12 +54,15 @@ RootShadowNode::Unshared ReanimatedCommitHook::shadowTreeWillCommit(
 
     rootNode = cloneShadowTreeWithNewProps(*rootNode, propsMap);
 
-    // If the commit comes from React Native then skip one commit from
+    // If the commit comes from React Native then pause commits from
     // Reanimated since the ShadowTree to be committed by Reanimated may not
     // include the new changes from React Native yet and all changes of animated
     // props will be applied in ReanimatedCommitHook by iterating over
     // PropsRegistry.
-    propsRegistry_->pleaseSkipReanimatedCommit();
+    // This is very important, since if we didn't pause Reanimated commits,
+    // it could lead to RN commits being delayed until the animation is finished
+    // (very bad).
+    propsRegistry_->pauseReanimatedCommits();
   }
 
   return rootNode;

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

@@ -9,8 +9,14 @@ namespace reanimated {
 // We need this information to skip unnecessary updates in
 // the commit hook.
 // Currently RN traits go up to 10, so hopefully
-// the arbitrarily chosen number 27 will be safe :)
+// the arbitrarily chosen numbers 27 and 28 will be safe :)
+
+// We have to use 2 traits, because we want to distinguish reanimated
+// commits both in the commit hook and mount hook. If we only had one trait
+// and didn't remove it in the commit hook, then any node that would clone
+// this node would also have our commit trait, rendering this trait useless.
 constexpr ShadowNodeTraits::Trait ReanimatedCommitTrait{1 << 27};
+constexpr ShadowNodeTraits::Trait ReanimatedMountTrait{1 << 28};
 
 class ReanimatedCommitShadowNode : public ShadowNode {
  public:
@@ -23,6 +29,15 @@ class ReanimatedCommitShadowNode : public ShadowNode {
   inline bool hasReanimatedCommitTrait() {
     return traits_.check(ReanimatedCommitTrait);
   }
+  inline void setReanimatedMountTrait() {
+    traits_.set(ReanimatedMountTrait);
+  }
+  inline void unsetReanimatedMountTrait() {
+    traits_.unset(ReanimatedMountTrait);
+  }
+  inline bool hasReanimatedMountTrait() {
+    return traits_.check(ReanimatedMountTrait);
+  }
 };
 
 } // namespace reanimated

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

@@ -0,0 +1,82 @@
+#ifdef RCT_NEW_ARCH_ENABLED
+
+#include "ReanimatedMountHook.h"
+#include "ReanimatedCommitShadowNode.h"
+
+namespace reanimated {
+
+ReanimatedMountHook::ReanimatedMountHook(
+    const std::shared_ptr<PropsRegistry> &propsRegistry,
+    const std::shared_ptr<UIManager> &uiManager)
+    : propsRegistry_(propsRegistry), uiManager_(uiManager) {
+  uiManager_->registerMountHook(*this);
+}
+
+ReanimatedMountHook::~ReanimatedMountHook() noexcept {
+  uiManager_->unregisterMountHook(*this);
+}
+
+void ReanimatedMountHook::shadowTreeDidMount(
+    RootShadowNode::Shared const &rootShadowNode,
+    double) noexcept {
+  auto reaShadowNode =
+      std::reinterpret_pointer_cast<ReanimatedCommitShadowNode>(
+          std::const_pointer_cast<RootShadowNode>(rootShadowNode));
+
+  if (reaShadowNode->hasReanimatedMountTrait()) {
+    // We mark reanimated commits with ReanimatedMountTrait. We don't want other
+    // shadow nodes to use this trait, but since this rootShadowNode is Shared,
+    // we don't have that guarantee. That's why we also unset this trait in the
+    // commit hook. We remove it here mainly for the sake of cleanliness.
+    reaShadowNode->unsetReanimatedMountTrait();
+    return;
+  }
+
+  // When commit from React Native has finished, we reset the skip commit flag
+  // in order to allow Reanimated to commit its tree
+  propsRegistry_->unpauseReanimatedCommits();
+  if (!propsRegistry_->shouldCommitAfterPause()) {
+    return;
+  }
+
+  const auto &shadowTreeRegistry = uiManager_->getShadowTreeRegistry();
+  shadowTreeRegistry.visit(
+      rootShadowNode->getSurfaceId(), [&](ShadowTree const &shadowTree) {
+        shadowTree.commit(
+            [&](RootShadowNode const &oldRootShadowNode)
+                -> RootShadowNode::Unshared {
+              PropsMap propsMap;
+
+              RootShadowNode::Unshared rootNode =
+                  std::static_pointer_cast<RootShadowNode>(
+                      oldRootShadowNode.ShadowNode::clone({}));
+              {
+                auto lock = propsRegistry_->createLock();
+
+                propsRegistry_->for_each([&](const ShadowNodeFamily &family,
+                                             const folly::dynamic &props) {
+                  propsMap[&family].emplace_back(props);
+                });
+
+                rootNode =
+                    cloneShadowTreeWithNewProps(oldRootShadowNode, propsMap);
+              }
+
+              // Mark the commit as Reanimated commit so that we can
+              // distinguish it in ReanimatedCommitHook.
+              auto reaShadowNode =
+                  std::reinterpret_pointer_cast<ReanimatedCommitShadowNode>(
+                      rootNode);
+              reaShadowNode->setReanimatedCommitTrait();
+
+              return rootNode;
+            },
+            {/* .enableStateReconciliation = */
+             false,
+             /* .mountSynchronously = */ true});
+      });
+}
+
+} // namespace reanimated
+
+#endif // RCT_NEW_ARCH_ENABLED

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

@@ -0,0 +1,33 @@
+#pragma once
+#ifdef RCT_NEW_ARCH_ENABLED
+
+#include "PropsRegistry.h"
+
+#include <react/renderer/uimanager/UIManagerMountHook.h>
+#include "ShadowTreeCloner.h"
+
+#include <memory>
+
+namespace reanimated {
+
+using namespace facebook::react;
+
+class ReanimatedMountHook : public UIManagerMountHook {
+ public:
+  ReanimatedMountHook(
+      const std::shared_ptr<PropsRegistry> &propsRegistry,
+      const std::shared_ptr<UIManager> &uiManager);
+  ~ReanimatedMountHook() noexcept override;
+
+  void shadowTreeDidMount(
+      RootShadowNode::Shared const &rootShadowNode,
+      double mountTime) noexcept override;
+
+ private:
+  const std::shared_ptr<PropsRegistry> propsRegistry_;
+  const std::shared_ptr<UIManager> uiManager_;
+};
+
+} // namespace reanimated
+
+#endif // RCT_NEW_ARCH_ENABLED

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

@@ -648,8 +648,9 @@ void NativeReanimatedModule::performOperations() {
   {
     auto lock = propsRegistry_->createLock();
 
-    if (copiedOperationsQueue.size() > 0) {
-      propsRegistry_->resetReanimatedSkipCommitFlag();
+    if (copiedOperationsQueue.size() > 0 &&
+        propsRegistry_->shouldReanimatedSkipCommit()) {
+      propsRegistry_->pleaseCommitAfterPause();
     }
 
     // remove recently unmounted ShadowNodes from PropsRegistry
@@ -726,15 +727,9 @@ void NativeReanimatedModule::performOperations() {
             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
-            // introduced in 0.72.4 but we have only check for minor version
-            // of React Native so enable that optimization in React Native >=
-            // 0.73
             if (propsRegistry_->shouldReanimatedSkipCommit()) {
               return nullptr;
             }
-#endif
           }
 
           auto rootNode =
@@ -750,15 +745,12 @@ void NativeReanimatedModule::performOperations() {
 
           return rootNode;
         },
-        { /* .enableStateReconciliation = */
-          false,
-#if REACT_NATIVE_MINOR_VERSION >= 72
-              /* .mountSynchronously = */ true,
-#endif
-              /* .shouldYield = */ [this]() {
-                return propsRegistry_->shouldReanimatedSkipCommit();
-              }
-        });
+        {/* .enableStateReconciliation = */
+         false,
+         /* .mountSynchronously = */ true,
+         /* .shouldYield = */ [this]() {
+           return propsRegistry_->shouldReanimatedSkipCommit();
+         }});
   });
 }
 
@@ -844,6 +836,8 @@ void NativeReanimatedModule::initializeFabric(
 
   initializeLayoutAnimations();
 
+  mountHook_ =
+      std::make_shared<ReanimatedMountHook>(propsRegistry_, uiManager_);
   commitHook_ =
       std::make_shared<ReanimatedCommitHook>(propsRegistry_, uiManager_);
 }

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

@@ -23,6 +23,7 @@
 #include "LayoutAnimationsProxy.h"
 #include "PropsRegistry.h"
 #include "ReanimatedCommitHook.h"
+#include "ReanimatedMountHook.h"
 #endif
 
 namespace reanimated {
@@ -224,6 +225,7 @@ class NativeReanimatedModule : public NativeReanimatedModuleSpec {
 
   std::shared_ptr<PropsRegistry> propsRegistry_;
   std::shared_ptr<ReanimatedCommitHook> commitHook_;
+  std::shared_ptr<ReanimatedMountHook> mountHook_;
 
   std::vector<Tag> tagsToRemove_; // from `propsRegistry_`
 #else