Add SurfaceHandlerCommitHook responsible for preserving layout constraints (#54836)

Summary:
Pull Request resolved: #54836

Changelog: [General][Added] Added a `SurfaceHandlerCommitHook` to preserve layout constraints during commits from different threads

Since layout constraints and layout context are being kept in the props of the `RootShadowNode`, it was possible for the JS revision merge to lose an update if it was being commited concurrently.

`SurfaceHandlerCommitHook` fixes this by ensuring that the root node always uses the latest context and constraints.

Differential Revision: D88839901

Author: j-piasecki

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

@@ -224,10 +224,18 @@ void Scheduler::registerSurface(
     const SurfaceHandler& surfaceHandler) const noexcept {
   surfaceHandler.setContextContainer(getContextContainer());
   surfaceHandler.setUIManager(uiManager_.get());
+
+  if (ReactNativeFeatureFlags::enableFabricCommitBranching()) {
+    surfaceHandler.registerCommitHook(uiManager_.get());
+  }
 }
 
 void Scheduler::unregisterSurface(
     const SurfaceHandler& surfaceHandler) const noexcept {
+  if (ReactNativeFeatureFlags::enableFabricCommitBranching()) {
+    surfaceHandler.unregisterCommitHook(uiManager_.get());
+  }
+
   surfaceHandler.setUIManager(nullptr);
 }
 

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

@@ -12,6 +12,8 @@
 #include <react/featureflags/ReactNativeFeatureFlags.h>
 #include <react/renderer/uimanager/UIManager.h>
 
+#include "SurfaceHandlerCommitHook.h"
+
 namespace facebook::react {
 
 using Status = SurfaceHandler::Status;
@@ -21,13 +23,15 @@ SurfaceHandler::SurfaceHandler(
     SurfaceId surfaceId) noexcept {
   parameters_.moduleName = moduleName;
   parameters_.surfaceId = surfaceId;
+  commitHook_ = std::make_shared<SurfaceHandlerCommitHook>(surfaceId);
 }
 
 #pragma mark - Surface Life-Cycle Management
 
 void SurfaceHandler::setContextContainer(
     std::shared_ptr<const ContextContainer> contextContainer) const noexcept {
-  parameters_.contextContainer = std::move(contextContainer);
+  parameters_.contextContainer = contextContainer;
+  commitHook_->setContextContainer(contextContainer);
 }
 
 Status SurfaceHandler::getStatus() const noexcept {
@@ -142,6 +146,7 @@ SurfaceId SurfaceHandler::getSurfaceId() const noexcept {
 void SurfaceHandler::setSurfaceId(SurfaceId surfaceId) const noexcept {
   std::unique_lock lock(parametersMutex_);
   parameters_.surfaceId = surfaceId;
+  commitHook_->setSurfaceId(surfaceId);
 }
 
 std::string SurfaceHandler::getModuleName() const noexcept {
@@ -213,78 +218,6 @@ Size SurfaceHandler::measure(
   return rootShadowNode->getLayoutMetrics().frame.size;
 }
 
-std::shared_ptr<const ShadowNode> SurfaceHandler::dirtyMeasurableNodesRecursive(
-    std::shared_ptr<const ShadowNode> node) const {
-  const auto nodeHasChildren = !node->getChildren().empty();
-  const auto isMeasurableYogaNode =
-      node->getTraits().check(ShadowNodeTraits::Trait::MeasurableYogaNode);
-
-  // Node is not measurable and has no children, its layout will not be affected
-  if (!nodeHasChildren && !isMeasurableYogaNode) {
-    return nullptr;
-  }
-
-  ShadowNode::SharedListOfShared newChildren =
-      ShadowNodeFragment::childrenPlaceholder();
-
-  if (nodeHasChildren) {
-    std::shared_ptr<std::vector<std::shared_ptr<const ShadowNode>>>
-        newChildrenMutable = nullptr;
-    for (size_t i = 0; i < node->getChildren().size(); i++) {
-      const auto& child = node->getChildren()[i];
-
-      if (const auto& layoutableNode =
-              std::dynamic_pointer_cast<const YogaLayoutableShadowNode>(
-                  child)) {
-        auto newChild = dirtyMeasurableNodesRecursive(layoutableNode);
-
-        if (newChild != nullptr) {
-          if (newChildrenMutable == nullptr) {
-            newChildrenMutable = std::make_shared<
-                std::vector<std::shared_ptr<const ShadowNode>>>(
-                node->getChildren());
-            newChildren = newChildrenMutable;
-          }
-
-          (*newChildrenMutable)[i] = newChild;
-        }
-      }
-    }
-
-    // Node is not measurable and its children were not dirtied, its layout will
-    // not be affected
-    if (!isMeasurableYogaNode && newChildrenMutable == nullptr) {
-      return nullptr;
-    }
-  }
-
-  const auto newNode = node->getComponentDescriptor().cloneShadowNode(
-      *node,
-      {
-          .children = newChildren,
-          // Preserve the original state of the node
-          .state = node->getState(),
-      });
-
-  if (isMeasurableYogaNode) {
-    std::static_pointer_cast<YogaLayoutableShadowNode>(newNode)->dirtyLayout();
-  }
-
-  return newNode;
-}
-
-void SurfaceHandler::dirtyMeasurableNodes(ShadowNode& root) const {
-  for (const auto& child : root.getChildren()) {
-    if (const auto& layoutableNode =
-            std::dynamic_pointer_cast<const YogaLayoutableShadowNode>(child)) {
-      const auto newChild = dirtyMeasurableNodesRecursive(layoutableNode);
-      if (newChild != nullptr) {
-        root.replaceChild(*child, newChild);
-      }
-    }
-  }
-}
-
 void SurfaceHandler::constraintLayout(
     const LayoutConstraints& layoutConstraints,
     const LayoutContext& layoutContext) const {
@@ -308,26 +241,16 @@ void SurfaceHandler::constraintLayout(
       return;
     }
 
-    PropsParserContext propsParserContext{
-        parameters_.surfaceId, *parameters_.contextContainer};
+    commitHook_->setLayoutConstraints(layoutConstraints, layoutContext);
 
     react_native_assert(
         link_.shadowTree && "`link_.shadowTree` must not be null.");
     link_.shadowTree->commit(
         [&](const RootShadowNode& oldRootShadowNode) {
-          auto newRoot = oldRootShadowNode.clone(
-              propsParserContext, layoutConstraints, layoutContext);
-
-          // Dirty all measurable nodes when the fontSizeMultiplier changes to
-          // trigger re-measurement.
-          if (ReactNativeFeatureFlags::enableFontScaleChangesUpdatingLayout() &&
-              layoutContext.fontSizeMultiplier !=
-                  oldRootShadowNode.getConcreteProps()
-                      .layoutContext.fontSizeMultiplier) {
-            dirtyMeasurableNodes(*newRoot);
-          }
-
-          return newRoot;
+          // Commit hook will update the layout constraints and context on the
+          // root node
+          return std::make_shared<RootShadowNode>(
+              oldRootShadowNode, ShadowNodeFragment{});
         },
         {/* default commit options */});
   }
@@ -394,6 +317,14 @@ void SurfaceHandler::setUIManager(const UIManager* uiManager) const noexcept {
       uiManager != nullptr ? Status::Registered : Status::Unregistered;
 }
 
+void SurfaceHandler::registerCommitHook(UIManager* uiManager) const noexcept {
+  uiManager->registerCommitHook(*commitHook_);
+}
+
+void SurfaceHandler::unregisterCommitHook(UIManager* uiManager) const noexcept {
+  uiManager->unregisterCommitHook(*commitHook_);
+}
+
 SurfaceHandler::~SurfaceHandler() noexcept {
   react_native_assert(
       link_.status == Status::Unregistered &&

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

@@ -22,6 +22,7 @@ class Scheduler;
 class ShadowTree;
 class MountingCoordinator;
 class UIManager;
+class SurfaceHandlerCommitHook;
 
 /*
  * Represents a running React Native surface and provides control over it.
@@ -148,14 +149,12 @@ class SurfaceHandler {
    * Must be called by `Scheduler` during registration process.
    */
   void setUIManager(const UIManager *uiManager) const noexcept;
+  void registerCommitHook(UIManager *uiManager) const noexcept;
+  void unregisterCommitHook(UIManager *uiManager) const noexcept;
 
   void applyDisplayMode(DisplayMode displayMode) const;
 
-  /*
-   * An utility for dirtying all measurable shadow nodes present in the tree.
-   */
-  void dirtyMeasurableNodes(ShadowNode &root) const;
-  std::shared_ptr<const ShadowNode> dirtyMeasurableNodesRecursive(std::shared_ptr<const ShadowNode> node) const;
+  std::shared_ptr<SurfaceHandlerCommitHook> commitHook_;
 
 #pragma mark - Link & Parameters
 

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

@@ -0,0 +1,155 @@
+/*
+ * 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 "SurfaceHandlerCommitHook.h"
+
+namespace facebook::react {
+
+SurfaceHandlerCommitHook::SurfaceHandlerCommitHook(SurfaceId surfaceId)
+    : surfaceId_(surfaceId) {}
+
+RootShadowNode::Unshared SurfaceHandlerCommitHook::shadowTreeWillCommit(
+    const ShadowTree& shadowTree,
+    const RootShadowNode::Shared& /*oldRootShadowNode*/,
+    const RootShadowNode::Unshared& newRootShadowNode,
+    const ShadowTreeCommitOptions& /*commitOptions*/) noexcept {
+  // Commit for another surface
+  if (surfaceId_ != shadowTree.getSurfaceId()) {
+    return newRootShadowNode;
+  }
+
+  LayoutConstraints layoutConstraints;
+  LayoutContext layoutContext;
+
+  {
+    std::shared_lock lock(mutex_);
+
+    // Layout constraints were not set
+    if (!layoutConstraints_.has_value() || !layoutContext_.has_value()) {
+      return newRootShadowNode;
+    }
+
+    layoutConstraints = layoutConstraints_.value();
+    layoutContext = layoutContext_.value();
+  }
+
+  // Don't clone the root node if the layout constraints and layout context
+  // are not changed
+  if (newRootShadowNode->getConcreteProps().layoutConstraints ==
+          layoutConstraints &&
+      newRootShadowNode->getConcreteProps().layoutContext == layoutContext) {
+    return newRootShadowNode;
+  }
+
+  PropsParserContext propsParserContext{surfaceId_, *contextContainer_};
+
+  auto clonedRoot = newRootShadowNode->clone(
+      propsParserContext, layoutConstraints, layoutContext);
+
+  // Dirty all measurable nodes when the fontSizeMultiplier changes to
+  // trigger re-measurement.
+  if (ReactNativeFeatureFlags::enableFontScaleChangesUpdatingLayout() &&
+      layoutContext.fontSizeMultiplier !=
+          newRootShadowNode->getConcreteProps()
+              .layoutContext.fontSizeMultiplier) {
+    dirtyMeasurableNodes(*clonedRoot);
+  }
+
+  return clonedRoot;
+}
+
+std::shared_ptr<const ShadowNode>
+SurfaceHandlerCommitHook::dirtyMeasurableNodesRecursive(
+    std::shared_ptr<const ShadowNode> node) const {
+  const auto nodeHasChildren = !node->getChildren().empty();
+  const auto isMeasurableYogaNode =
+      node->getTraits().check(ShadowNodeTraits::Trait::MeasurableYogaNode);
+
+  // Node is not measurable and has no children, its layout will not be affected
+  if (!nodeHasChildren && !isMeasurableYogaNode) {
+    return nullptr;
+  }
+
+  ShadowNode::SharedListOfShared newChildren =
+      ShadowNodeFragment::childrenPlaceholder();
+
+  if (nodeHasChildren) {
+    std::shared_ptr<std::vector<std::shared_ptr<const ShadowNode>>>
+        newChildrenMutable = nullptr;
+    for (size_t i = 0; i < node->getChildren().size(); i++) {
+      const auto& child = node->getChildren()[i];
+
+      if (const auto& layoutableNode =
+              std::dynamic_pointer_cast<const YogaLayoutableShadowNode>(
+                  child)) {
+        auto newChild = dirtyMeasurableNodesRecursive(layoutableNode);
+
+        if (newChild != nullptr) {
+          if (newChildrenMutable == nullptr) {
+            newChildrenMutable = std::make_shared<
+                std::vector<std::shared_ptr<const ShadowNode>>>(
+                node->getChildren());
+            newChildren = newChildrenMutable;
+          }
+
+          (*newChildrenMutable)[i] = newChild;
+        }
+      }
+    }
+
+    // Node is not measurable and its children were not dirtied, its layout will
+    // not be affected
+    if (!isMeasurableYogaNode && newChildrenMutable == nullptr) {
+      return nullptr;
+    }
+  }
+
+  const auto newNode = node->getComponentDescriptor().cloneShadowNode(
+      *node,
+      {
+          .children = newChildren,
+          // Preserve the original state of the node
+          .state = node->getState(),
+      });
+
+  if (isMeasurableYogaNode) {
+    std::static_pointer_cast<YogaLayoutableShadowNode>(newNode)->dirtyLayout();
+  }
+
+  return newNode;
+}
+
+void SurfaceHandlerCommitHook::dirtyMeasurableNodes(ShadowNode& root) const {
+  for (const auto& child : root.getChildren()) {
+    if (const auto& layoutableNode =
+            std::dynamic_pointer_cast<const YogaLayoutableShadowNode>(child)) {
+      const auto newChild = dirtyMeasurableNodesRecursive(layoutableNode);
+      if (newChild != nullptr) {
+        root.replaceChild(*child, newChild);
+      }
+    }
+  }
+}
+
+void SurfaceHandlerCommitHook::setContextContainer(
+    std::shared_ptr<const ContextContainer> contextContainer) {
+  contextContainer_ = std::move(contextContainer);
+}
+
+void SurfaceHandlerCommitHook::setLayoutConstraints(
+    const LayoutConstraints& layoutConstraints,
+    const LayoutContext& layoutContext) {
+  std::unique_lock lock(mutex_);
+  layoutConstraints_ = layoutConstraints;
+  layoutContext_ = layoutContext;
+}
+
+void SurfaceHandlerCommitHook::setSurfaceId(SurfaceId surfaceId) {
+  surfaceId_ = surfaceId;
+}
+
+} // namespace facebook::react