Fix dropped React revision when merging commit branches

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

@@ -428,7 +428,9 @@ CommitStatus ShadowTree::tryCommit(
       return CommitStatus::Failed;
     }
 
-    auto newRevisionNumber = currentRevision_.number + 1;
+    auto newRevisionNumber = isReactBranch
+        ? oldRevisionForStateProgression.number + 1
+        : currentRevision_.number + 1;
 
     {
       std::scoped_lock dispatchLock(EventEmitter::DispatchMutex());
@@ -513,11 +515,12 @@ void ShadowTree::mergeReactRevision() const {
     }
   }
 
-  ShadowTreeRevision::Number lastMergedRevisionNumber;
+  RootShadowNode::Shared lastMergedRootShadowNode;
+  bool lastMergeSucceeded = false;
 
   if (isPropsUpdatesAccumulationGuaranteed()) {
-    lastMergedRevisionNumber = promotedRevision.number;
-    this->commit(
+    lastMergedRootShadowNode = promotedRevision.rootShadowNode;
+    auto status = this->commit(
         [revision = std::move(promotedRevision)](
             const RootShadowNode& /*oldRootShadowNode*/) {
           return std::make_shared<RootShadowNode>(
@@ -528,13 +531,14 @@ void ShadowTree::mergeReactRevision() const {
             .mountSynchronously = true,
             .source = CommitSource::ReactRevisionMerge,
         });
+    lastMergeSucceeded = status == CommitStatus::Succeeded;
   } else {
     for (size_t i = 0; i < promotedRevisions.size(); ++i) {
       auto& revision = promotedRevisions[i];
       bool isLast = i == promotedRevisions.size() - 1;
-      lastMergedRevisionNumber = revision.number;
+      lastMergedRootShadowNode = revision.rootShadowNode;
 
-      this->commit(
+      auto status = this->commit(
           [revision = std::move(revision)](
               const RootShadowNode& /*oldRootShadowNode*/) {
             return std::make_shared<RootShadowNode>(
@@ -545,6 +549,7 @@ void ShadowTree::mergeReactRevision() const {
               .mountSynchronously = true,
               .source = CommitSource::ReactRevisionMerge,
           });
+      lastMergeSucceeded = status == CommitStatus::Succeeded;
     }
   }
 
@@ -554,8 +559,9 @@ void ShadowTree::mergeReactRevision() const {
 
     // If the current react revision is the same as the one that was just
     // merged, clear it.
-    if (currentReactRevision_.has_value() &&
-        lastMergedRevisionNumber == currentReactRevision_.value().number) {
+    if (lastMergeSucceeded && currentReactRevision_.has_value() &&
+        currentReactRevision_.value().rootShadowNode ==
+            lastMergedRootShadowNode) {
       currentReactRevision_.reset();
     }
   }

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

@@ -0,0 +1,141 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <memory>
+
+#include <gtest/gtest.h>
+
+#include <react/featureflags/ReactNativeFeatureFlags.h>
+#include <react/featureflags/ReactNativeFeatureFlagsDefaults.h>
+#include <react/renderer/components/view/ViewComponentDescriptor.h>
+#include <react/renderer/element/ComponentBuilder.h>
+#include <react/renderer/element/Element.h>
+#include <react/renderer/element/testUtils.h>
+#include <react/renderer/mounting/MountingCoordinator.h>
+#include <react/renderer/mounting/ShadowTree.h>
+#include <react/renderer/mounting/ShadowTreeDelegate.h>
+
+using namespace facebook::react;
+
+namespace {
+
+class BranchingEnabledFlags : public ReactNativeFeatureFlagsDefaults {
+ public:
+  bool enableFabricCommitBranching() override {
+    return true;
+  }
+};
+
+class DummyShadowTreeDelegate : public ShadowTreeDelegate {
+ public:
+  RootShadowNode::Unshared shadowTreeWillCommit(
+      const ShadowTree& /*shadowTree*/,
+      const RootShadowNode::Shared& /*oldRootShadowNode*/,
+      const RootShadowNode::Unshared& newRootShadowNode,
+      const ShadowTree::CommitOptions& /*commitOptions*/) const override {
+    return newRootShadowNode;
+  }
+
+  void shadowTreeDidFinishTransaction(
+      std::shared_ptr<const MountingCoordinator> /*mountingCoordinator*/,
+      bool /*mountSynchronously*/) const override {}
+
+  void shadowTreeDidFinishReactCommit(
+      const ShadowTree& /*shadowTree*/) const override {}
+
+  void shadowTreeDidPromoteReactRevision(
+      const ShadowTree& /*shadowTree*/) const override {}
+};
+
+} // namespace
+
+class ShadowTreeReactBranchingTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    // Must override before any component construction reads a flag.
+    ReactNativeFeatureFlags::dangerouslyReset();
+    ReactNativeFeatureFlags::override(
+        std::make_unique<BranchingEnabledFlags>());
+  }
+
+  void TearDown() override {
+    ReactNativeFeatureFlags::dangerouslyReset();
+  }
+};
+
+// Two React revisions committed before a merge share a revision number.
+// Merging the first must not clear the second: the old number-based comparison
+// cleared it (dropping a pending update); the fix compares root node identity.
+TEST_F(ShadowTreeReactBranchingTest, mergeDoesNotDropNewerReactRevision) {
+  // clang-format off
+  auto element =
+      Element<RootShadowNode>()
+        .children({
+          Element<ViewShadowNode>()
+        });
+  // clang-format on
+
+  ContextContainer contextContainer{};
+  auto builder = simpleComponentBuilder();
+  auto initialRootShadowNode = builder.build(element);
+  auto shadowTreeDelegate = DummyShadowTreeDelegate{};
+
+  ShadowTree shadowTree{
+      SurfaceId{11},
+      LayoutConstraints{},
+      LayoutContext{},
+      shadowTreeDelegate,
+      contextContainer};
+
+  // Initial (non-React) commit establishes `currentRevision_`.
+  shadowTree.commit(
+      [&](const RootShadowNode& /*oldRootShadowNode*/) {
+        return std::static_pointer_cast<RootShadowNode>(initialRootShadowNode);
+      },
+      {.enableStateReconciliation = false});
+
+  // First React revision (R1).
+  auto rootR1 = std::static_pointer_cast<RootShadowNode>(
+      initialRootShadowNode->ShadowNode::clone({}));
+  shadowTree.commit(
+      [&](const RootShadowNode& /*oldRootShadowNode*/) { return rootR1; },
+      {.enableStateReconciliation = false,
+       .mountSynchronously = false,
+       .source = ShadowTreeCommitSource::React});
+
+  auto revisionAfterR1 = shadowTree.getCurrentReactRevision();
+  ASSERT_TRUE(revisionAfterR1.has_value());
+
+  // Promote R1 (end-of-tick), but do not merge it yet.
+  shadowTree.promoteReactRevision();
+
+  // Second React revision (R2) lands before the merge. The main revision has
+  // not advanced, so R2 gets R1's revision number but is a distinct tree.
+  auto rootR2 = std::static_pointer_cast<RootShadowNode>(
+      initialRootShadowNode->ShadowNode::clone({}));
+  shadowTree.commit(
+      [&](const RootShadowNode& /*oldRootShadowNode*/) { return rootR2; },
+      {.enableStateReconciliation = false,
+       .mountSynchronously = false,
+       .source = ShadowTreeCommitSource::React});
+
+  auto revisionAfterR2 = shadowTree.getCurrentReactRevision();
+  ASSERT_TRUE(revisionAfterR2.has_value());
+
+  // Precondition for the bug: same number, different revisions.
+  EXPECT_EQ(revisionAfterR1->number, revisionAfterR2->number);
+  EXPECT_NE(revisionAfterR2->rootShadowNode, revisionAfterR1->rootShadowNode);
+
+  // Merging R1 must not clear R2.
+  shadowTree.mergeReactRevision();
+
+  auto revisionAfterMerge = shadowTree.getCurrentReactRevision();
+  ASSERT_TRUE(revisionAfterMerge.has_value())
+      << "Newer React revision (R2) was incorrectly dropped while merging R1";
+  EXPECT_EQ(
+      revisionAfterMerge->rootShadowNode, revisionAfterR2->rootShadowNode);
+}