Update branching implementation to default to the JS thread merge (#58615)

Summary:

Changelog: [Internal]

The current behavior, merge on the main thread, is kept behind the new `enableFabricCommitBranchingMergeOnMainThread` flag.

Reviewed By: rubennorte

Differential Revision: D120981616

Author: j-piasecki

packages/react-native/React/Fabric/RCTSurfacePresenter.mm

@@ -101,7 +101,7 @@ - (instancetype)initWithContextContainer:(std::shared_ptr<const ContextContainer
     _mountingManager.contextContainer = contextContainer;
     _mountingManager.delegate = self;
 
-    if (ReactNativeFeatureFlags::enableFabricCommitBranching()) {
+    if (ReactNativeFeatureFlags::enableFabricCommitBranchingMergeOnMainThread()) {
       _mergeRunLoopObserverDelegate = std::make_shared<ReactRevisionMergeRunLoopObserverDelegate>(self);
       _mergeRunLoopObserver = std::make_unique<const MainRunLoopObserver>(
           RunLoopObserver::Activity::BeforeWaiting, _mergeRunLoopObserverDelegate);
@@ -350,7 +350,7 @@ - (void)schedulerShouldRenderTransactions:(std::shared_ptr<const MountingCoordin
 
 - (void)schedulerShouldMergeReactRevision:(SurfaceId)surfaceId
 {
-  if (RCTIsMainQueue()) {
+  if (RCTIsMainQueue() || !ReactNativeFeatureFlags::enableFabricCommitBranchingMergeOnMainThread()) {
     [self _mergeReactRevisionForSurfaceId:surfaceId];
     return;
   }
@@ -369,7 +369,6 @@ - (void)schedulerShouldMergeReactRevision:(SurfaceId)surfaceId
 
 - (void)_mergeReactRevisionForSurfaceId:(SurfaceId)surfaceId
 {
-  RCTAssertMainQueue();
   RCTScheduler *scheduler = [self scheduler];
   if (!scheduler) {
     return;

packages/react-native/ReactAndroid/src/main/java/com/facebook/react/fabric/FabricUIManager.java

@@ -1628,7 +1628,7 @@ public void doFrameGuarded(long frameTimeNanos) {
 
       // Drain pending React revision merges first so that animations,
       // preallocation, and mount items operate against the latest revision.
-      if (ReactNativeFeatureFlags.enableFabricCommitBranching()) {
+      if (ReactNativeFeatureFlags.enableFabricCommitBranchingMergeOnMainThread()) {
         FabricUIManagerBinding binding = mBinding;
         if (binding != null) {
           Integer mergeSurfaceId;

packages/react-native/ReactAndroid/src/main/jni/react/fabric/FabricUIManagerBinding.cpp

@@ -767,17 +767,24 @@ void FabricUIManagerBinding::schedulerShouldRenderTransactions(
 
 void FabricUIManagerBinding::schedulerShouldMergeReactRevision(
     SurfaceId surfaceId) {
-  std::shared_lock lock(installMutex_);
-  auto mountingManager =
-      getMountingManager("schedulerShouldMergeReactRevision");
-  if (mountingManager) {
-    mountingManager->scheduleReactRevisionMerge(surfaceId);
+  if (ReactNativeFeatureFlags::enableFabricCommitBranchingMergeOnMainThread()) {
+    auto mountingManager =
+        getMountingManager("schedulerShouldMergeReactRevision");
+    if (mountingManager) {
+      mountingManager->scheduleReactRevisionMerge(surfaceId);
+    }
+  } else {
+    mergeReactRevision(surfaceId);
   }
 }
 
 void FabricUIManagerBinding::mergeReactRevision(SurfaceId surfaceId) {
-  std::shared_lock lock(installMutex_);
-  scheduler_->getUIManager()->getShadowTreeRegistry().visit(
+  auto scheduler = getScheduler();
+  if (!scheduler) {
+    return;
+  }
+
+  scheduler->getUIManager()->getShadowTreeRegistry().visit(
       surfaceId,
       [](const ShadowTree& shadowTree) { shadowTree.mergeReactRevision(); });
 }

packages/react-native/ReactCommon/react/renderer/mounting/ShadowTree.cpp

@@ -492,6 +492,7 @@ void ShadowTree::mount(ShadowTreeRevision revision, bool mountSynchronously)
 }
 
 void ShadowTree::mergeReactRevision() const {
+  TraceSection s("ShadowTree::mergeReactRevision");
   ShadowTreeRevision promotedRevision;
   std::vector<ShadowTreeRevision> promotedRevisions;
   // If props updates accumulation is guaranteed, we can merge the promoted
@@ -568,7 +569,7 @@ void ShadowTree::mergeReactRevision() const {
   }
 }
 
-void ShadowTree::promoteReactRevision() const {
+bool ShadowTree::promoteReactRevision() const {
   // Promote only when props updates accumulation is guaranteed. Otherwise,
   // queuedReactRevisions_ will be used instead.
   if (isPropsUpdatesAccumulationGuaranteed()) {
@@ -579,7 +580,7 @@ void ShadowTree::promoteReactRevision() const {
       // have more than one promotion in a row. In this case, all but the first
       // one should no-op.
       if (!currentReactRevision_.has_value()) {
-        return;
+        return false;
       }
       currentReactRevision = currentReactRevision_.value();
     }
@@ -592,7 +593,7 @@ void ShadowTree::promoteReactRevision() const {
     UniqueLock lock = uniqueRevisionLock(false);
 
     if (queuedReactRevisions_.empty()) {
-      return;
+      return false;
     }
 
     // Move all queued revisions to the promoted revisions.
@@ -603,7 +604,7 @@ void ShadowTree::promoteReactRevision() const {
     queuedReactRevisions_.clear();
   }
 
-  delegate_.shadowTreeDidPromoteReactRevision(*this);
+  return true;
 }
 
 void ShadowTree::scheduleReactRevisionPromotion() const {

packages/react-native/ReactCommon/react/renderer/mounting/ShadowTree.h

@@ -148,9 +148,9 @@ class ShadowTree final {
 
   /**
    * Promotes the current React revision to be merged into the main branch of the
-   * ShadowTree.
+   * ShadowTree. Returns `true` if a revision was promoted.
    */
-  void promoteReactRevision() const;
+  bool promoteReactRevision() const;
 
   /**
    * Commits the currently promoted React revision to the "main" branch of the

packages/react-native/ReactCommon/react/renderer/mounting/ShadowTreeDelegate.h

@@ -46,12 +46,6 @@ class ShadowTreeDelegate {
    */
   virtual void shadowTreeDidFinishReactCommit(const ShadowTree &shadowTree) const = 0;
 
-  /*
-   * Called right after Shadow Tree promotes a React revision of the tree to
-   * be merged.
-   */
-  virtual void shadowTreeDidPromoteReactRevision(const ShadowTree &shadowTree) const = 0;
-
   /*
    * Called right after a Shadow Tree commits a new tree, reporting the nodes
    * whose layout changed in this commit.

packages/react-native/ReactCommon/react/renderer/mounting/tests/ShadowTreeReactBranchingTest.cpp

@@ -46,9 +46,6 @@ class DummyShadowTreeDelegate : public ShadowTreeDelegate {
 
   void shadowTreeDidFinishReactCommit(
       const ShadowTree& /*shadowTree*/) const override {}
-
-  void shadowTreeDidPromoteReactRevision(
-      const ShadowTree& /*shadowTree*/) const override {}
 };
 
 } // namespace

packages/react-native/ReactCommon/react/renderer/mounting/tests/StateReconciliationTest.cpp

@@ -40,9 +40,6 @@ class DummyShadowTreeDelegate : public ShadowTreeDelegate {
 
   void shadowTreeDidFinishReactCommit(
       const ShadowTree& /*shadowTree*/) const override {}
-
-  void shadowTreeDidPromoteReactRevision(
-      const ShadowTree& /*shadowTree*/) const override {}
 };
 
 namespace {

packages/react-native/ReactCommon/react/renderer/scheduler/Scheduler.cpp

@@ -419,17 +419,17 @@ void Scheduler::uiManagerShouldRemoveEventListener(
 void Scheduler::uiManagerDidFinishReactCommit(const ShadowTree& shadowTree) {
   auto surfaceId = shadowTree.getSurfaceId();
   runtimeScheduler_->scheduleRenderingUpdate(
-      surfaceId, [surfaceId, uiManager = uiManager_]() {
+      surfaceId, [surfaceId, uiManager = uiManager_, this]() {
+        bool promoted = false;
+
         uiManager->getShadowTreeRegistry().visit(
             surfaceId,
-            [](const ShadowTree& tree) { tree.promoteReactRevision(); });
-      });
-}
+            [&](const ShadowTree& tree) { promoted = tree.promoteReactRevision(); });
 
-void Scheduler::uiManagerDidPromoteReactRevision(const ShadowTree& shadowTree) {
-  if (delegate_ != nullptr) {
-    delegate_->schedulerShouldMergeReactRevision(shadowTree.getSurfaceId());
-  }
+        if (promoted && delegate_ != nullptr) {
+          delegate_->schedulerShouldMergeReactRevision(surfaceId);
+        }
+      });
 }
 
 void Scheduler::uiManagerDidStartSurface(const ShadowTree& shadowTree) {

packages/react-native/ReactCommon/react/renderer/scheduler/Scheduler.h

@@ -107,7 +107,6 @@ class Scheduler final : public UIManagerDelegate {
   void uiManagerShouldAddEventListener(std::shared_ptr<const EventListener> listener) final;
   void uiManagerShouldRemoveEventListener(const std::shared_ptr<const EventListener> &listener) final;
   void uiManagerDidFinishReactCommit(const ShadowTree &shadowTree) override;
-  void uiManagerDidPromoteReactRevision(const ShadowTree &shadowTree) override;
   void uiManagerDidStartSurface(const ShadowTree &shadowTree) override;
 
 #pragma mark - ContextContainer