Refactor how we promote shadow tree revisions as latest for JS consistency (#54833)

Summary:

Changelog: [General][Changed] - Refactor how shadow tree revisions are promoted as latest for JS consistency

Changes `LazyShadowTreeRevisionConsistencyManager::updateCurrentRevision` to pull the latest commited revision for a given surface id instead of accepting the new revision as a parameter.

This way, the new revision is read from the same place for every update, which opens up the way to implement commit branching. This change will allow to always read from the JS tree revision, if it exists.

Differential Revision: D88151490

Author: j-piasecki

packages/react-native/ReactCommon/react/renderer/uimanager/UIManager.cpp

@@ -188,8 +188,10 @@ void UIManager::completeSurface(
     ShadowTree::CommitOptions commitOptions) {
   TraceSection s("UIManager::completeSurface", "surfaceId", surfaceId);
 
+  ShadowTree::CommitStatus result;
+
   shadowTreeRegistry_.visit(surfaceId, [&](const ShadowTree& shadowTree) {
-    auto result = shadowTree.commit(
+    result = shadowTree.commit(
         [&](const RootShadowNode& oldRootShadowNode) {
           return std::make_shared<RootShadowNode>(
               oldRootShadowNode,
@@ -199,20 +201,19 @@ void UIManager::completeSurface(
               });
         },
         commitOptions);
+  });
 
-    if (result == ShadowTree::CommitStatus::Succeeded) {
-      // It's safe to update the visible revision of the shadow tree immediately
-      // after we commit a specific one.
-      lazyShadowTreeRevisionConsistencyManager_->updateCurrentRevision(
-          surfaceId, shadowTree.getCurrentRevision().rootShadowNode);
+  if (result == ShadowTree::CommitStatus::Succeeded) {
+    // It's safe to update the visible revision of the shadow tree immediately
+    // after we commit a specific one.
+    lazyShadowTreeRevisionConsistencyManager_->updateCurrentRevision(surfaceId);
 
-      if (ReactNativeFeatureFlags::useSharedAnimatedBackend()) {
-        if (auto animationBackend = animationBackend_.lock()) {
-          animationBackend->clearRegistry(surfaceId);
-        }
+    if (ReactNativeFeatureFlags::useSharedAnimatedBackend()) {
+      if (auto animationBackend = animationBackend_.lock()) {
+        animationBackend->clearRegistry(surfaceId);
       }
     }
-  });
+  }
 }
 
 void UIManager::setIsJSResponder(

packages/react-native/ReactCommon/react/renderer/uimanager/consistency/LazyShadowTreeRevisionConsistencyManager.cpp

@@ -15,22 +15,32 @@ LazyShadowTreeRevisionConsistencyManager::
         ShadowTreeRegistry& shadowTreeRegistry)
     : shadowTreeRegistry_(shadowTreeRegistry) {}
 
-void LazyShadowTreeRevisionConsistencyManager::updateCurrentRevision(
-    SurfaceId surfaceId,
-    RootShadowNode::Shared rootShadowNode) {
+std::shared_ptr<const RootShadowNode>
+LazyShadowTreeRevisionConsistencyManager::updateCurrentRevision(
+    SurfaceId surfaceId) {
+  // This method is only going to be called from JS, so we don't need to protect
+  // the access to the shadow tree registry as well.
+  // If this was multi-threaded, we would need to protect it to avoid capturing
+  // root shadow nodes concurrently.
+  RootShadowNode::Shared rootShadowNode;
+  shadowTreeRegistry_.visit(surfaceId, [&](const ShadowTree& shadowTree) {
+    rootShadowNode = shadowTree.getCurrentRevision().rootShadowNode;
+  });
+
   std::unique_lock lock(capturedRootShadowNodesForConsistencyMutex_);
 
   // We don't need to store the revision if we haven't locked.
   // We can resolve lazily when requested.
   if (lockCount > 0) {
-    capturedRootShadowNodesForConsistency_[surfaceId] =
-        std::move(rootShadowNode);
+    capturedRootShadowNodesForConsistency_[surfaceId] = rootShadowNode;
   }
+
+  return rootShadowNode;
 }
 
 #pragma mark - ShadowTreeRevisionProvider
 
-RootShadowNode::Shared
+std::shared_ptr<const RootShadowNode>
 LazyShadowTreeRevisionConsistencyManager::getCurrentRevision(
     SurfaceId surfaceId) {
   {
@@ -43,23 +53,7 @@ LazyShadowTreeRevisionConsistencyManager::getCurrentRevision(
     }
   }
 
-  // This method is only going to be called from JS, so we don't need to protect
-  // the access to the shadow tree registry as well.
-  // If this was multi-threaded, we would need to protect it to avoid capturing
-  // root shadow nodes concurrently.
-  RootShadowNode::Shared rootShadowNode;
-  shadowTreeRegistry_.visit(surfaceId, [&](const ShadowTree& shadowTree) {
-    rootShadowNode = shadowTree.getCurrentRevision().rootShadowNode;
-  });
-
-  {
-    std::unique_lock lock(capturedRootShadowNodesForConsistencyMutex_);
-    if (lockCount > 0) {
-      capturedRootShadowNodesForConsistency_[surfaceId] = rootShadowNode;
-    }
-  }
-
-  return rootShadowNode;
+  return updateCurrentRevision(surfaceId);
 }
 
 #pragma mark - ConsistentShadowTreeRevisionProvider

packages/react-native/ReactCommon/react/renderer/uimanager/consistency/LazyShadowTreeRevisionConsistencyManager.h

@@ -28,11 +28,11 @@ class LazyShadowTreeRevisionConsistencyManager : public ShadowTreeRevisionConsis
  public:
   explicit LazyShadowTreeRevisionConsistencyManager(ShadowTreeRegistry &shadowTreeRegistry);
 
-  void updateCurrentRevision(SurfaceId surfaceId, RootShadowNode::Shared rootShadowNode);
+  std::shared_ptr<const RootShadowNode> updateCurrentRevision(SurfaceId surfaceId);
 
 #pragma mark - ShadowTreeRevisionProvider
 
-  RootShadowNode::Shared getCurrentRevision(SurfaceId surfaceId) override;
+  std::shared_ptr<const RootShadowNode> getCurrentRevision(SurfaceId surfaceId) override;
 
 #pragma mark - ShadowTreeRevisionConsistencyManager
 

packages/react-native/ReactCommon/react/renderer/uimanager/consistency/tests/LazyShadowTreeRevisionConsistencyManagerTest.cpp

@@ -141,7 +141,7 @@ TEST_F(
 
   EXPECT_EQ(consistencyManager_.getCurrentRevision(0), nullptr);
 
-  consistencyManager_.updateCurrentRevision(0, newRootShadowNode);
+  consistencyManager_.updateCurrentRevision(0);
 
   EXPECT_NE(consistencyManager_.getCurrentRevision(0), nullptr);
   EXPECT_EQ(
@@ -176,7 +176,7 @@ TEST_F(
 
   EXPECT_EQ(consistencyManager_.getCurrentRevision(0), nullptr);
 
-  consistencyManager_.updateCurrentRevision(0, newRootShadowNode);
+  consistencyManager_.updateCurrentRevision(0);
 
   EXPECT_EQ(
       consistencyManager_.getCurrentRevision(0).get(), newRootShadowNode.get());
@@ -192,7 +192,7 @@ TEST_F(
             {});
       });
 
-  consistencyManager_.updateCurrentRevision(0, newRootShadowNode2);
+  consistencyManager_.updateCurrentRevision(0);
 
   EXPECT_EQ(
       consistencyManager_.getCurrentRevision(0).get(),
@@ -265,7 +265,7 @@ TEST_F(
   EXPECT_EQ(
       consistencyManager_.getCurrentRevision(0).get(), newRootShadowNode.get());
 
-  consistencyManager_.updateCurrentRevision(0, newRootShadowNode2);
+  consistencyManager_.updateCurrentRevision(0);
 
   // Updated
   EXPECT_EQ(
@@ -344,7 +344,9 @@ TEST_F(LazyShadowTreeRevisionConsistencyManagerTest, testUpdateToUnmounted) {
   EXPECT_EQ(
       consistencyManager_.getCurrentRevision(0).get(), newRootShadowNode.get());
 
-  consistencyManager_.updateCurrentRevision(0, nullptr);
+  shadowTreeRegistry_.remove(0);
+
+  consistencyManager_.updateCurrentRevision(0);
 
   // Updated
   EXPECT_EQ(consistencyManager_.getCurrentRevision(0).get(), nullptr);