Fix display: contents nodes not being cloned with the wrong owner (#1826)

Summary:
X-link: https://github.com/facebook/react-native/pull/52530

This PR fixes two issues with `display: contents` implementation:
1. When a node with `display: contents` set is a leaf, it won't be cloned after the initial tree is built. The added test case covers this scenario.
2. It was possible for the subtree of `display: contents` nodes not to be cloned during layout. I don't have a minimal reproduction for this one, unfortunately. It was discovered in the Expensify app: https://github.com/Expensify/App/issues/65268, along with a consistent reproduction. In that specific case, it seems to be heavily tied to `react-native-onyx`, which is a state management library.

Changelog: [GENERAL][FIXED] - Fixed nodes with `display: contents` set being cloned with the wrong owner

Pull Request resolved: https://github.com/facebook/yoga/pull/1826

Reviewed By: adityasharat, NickGerleman

Differential Revision: D78084270

Pulled By: j-piasecki

fbshipit-source-id: eb81f6d7dcd1665974d07261ba693e2abea239bb

tests/YGPersistentNodeCloningTest.cpp

@@ -181,4 +181,34 @@ TEST_F(
   EXPECT_EQ(nodesCloned[0], a.get());
 }
 
+TEST_F(YGPersistentNodeCloningTest, clone_leaf_display_contents_node) {
+  // <View id="A">
+  //   <View id="B" style={{ display: 'contents' }} />
+  // </View>
+
+  auto b = std::make_shared<NodeWrapper>(config);
+  auto a = std::make_shared<NodeWrapper>(config, std::vector{b});
+  YGNodeStyleSetDisplay(b->node, YGDisplayContents);
+
+  // We don't expect any cloning during the first layout
+  onClone = [](...) { FAIL(); };
+
+  YGNodeCalculateLayout(a->node, YGUndefined, YGUndefined, YGDirectionLTR);
+
+  auto aPrime = std::make_shared<NodeWrapper>(config, std::vector{b});
+
+  std::vector<NodeWrapper*> nodesCloned;
+  // We should clone "C"
+  onClone = [&](YGNodeConstRef oldNode,
+                YGNodeConstRef /*owner*/,
+                size_t /*childIndex*/) {
+    nodesCloned.push_back(static_cast<NodeWrapper*>(YGNodeGetContext(oldNode)));
+  };
+
+  YGNodeCalculateLayout(aPrime->node, 100, 100, YGDirectionLTR);
+
+  EXPECT_EQ(nodesCloned.size(), 1);
+  EXPECT_EQ(nodesCloned[0], b.get());
+}
+
 } // namespace facebook::yoga

yoga/algorithm/CalculateLayout.cpp

@@ -478,6 +478,8 @@ static void zeroOutLayoutRecursively(yoga::Node* const node) {
 }
 
 static void cleanupContentsNodesRecursively(yoga::Node* const node) {
+  if (node->hasContentsChildren()) [[unlikely]] {
+    node->cloneContentsChildrenIfNeeded();
     for (auto child : node->getChildren()) {
       if (child->style().display() == Display::Contents) {
         child->getLayout() = {};
@@ -490,6 +492,7 @@ static void cleanupContentsNodesRecursively(yoga::Node* const node) {
         cleanupContentsNodesRecursively(child);
       }
     }
+  }
 }
 
 static float calculateAvailableInnerDimension(

yoga/node/Node.cpp

@@ -181,6 +181,17 @@ void Node::setDirty(bool isDirty) {
   }
 }
 
+void Node::setChildren(const std::vector<Node*>& children) {
+  children_ = children;
+
+  contentsChildrenCount_ = 0;
+  for (const auto& child : children) {
+    if (child->style().display() == Display::Contents) {
+      contentsChildrenCount_++;
+    }
+  }
+}
+
 bool Node::removeChild(Node* child) {
   auto p = std::find(children_.begin(), children_.end(), child);
   if (p != children_.end()) {
@@ -380,6 +391,23 @@ void Node::cloneChildrenIfNeeded() {
     if (child->getOwner() != this) {
       child = resolveRef(config_->cloneNode(child, this, i));
       child->setOwner(this);
+
+      if (child->hasContentsChildren()) [[unlikely]] {
+        child->cloneContentsChildrenIfNeeded();
+      }
+    }
+    i += 1;
+  }
+}
+
+void Node::cloneContentsChildrenIfNeeded() {
+  size_t i = 0;
+  for (Node*& child : children_) {
+    if (child->style().display() == Display::Contents &&
+        child->getOwner() != this) {
+      child = resolveRef(config_->cloneNode(child, this, i));
+      child->setOwner(this);
+      child->cloneChildrenIfNeeded();
     }
     i += 1;
   }

yoga/node/Node.h

@@ -96,6 +96,10 @@ class YG_EXPORT Node : public ::YGNode {
     return config_->hasErrata(errata);
   }
 
+  bool hasContentsChildren() const {
+    return contentsChildrenCount_ != 0;
+  }
+
   YGDirtiedFunc getDirtiedFunc() const {
     return dirtiedFunc_;
   }
@@ -244,15 +248,12 @@ class YG_EXPORT Node : public ::YGNode {
     owner_ = owner;
   }
 
-  void setChildren(const std::vector<Node*>& children) {
-    children_ = children;
-  }
-
   // TODO: rvalue override for setChildren
 
   void setConfig(Config* config);
 
   void setDirty(bool isDirty);
+  void setChildren(const std::vector<Node*>& children);
   void setLayoutLastOwnerDirection(Direction direction);
   void setLayoutComputedFlexBasis(FloatOptional computedFlexBasis);
   void setLayoutComputedFlexBasisGeneration(
@@ -286,6 +287,7 @@ class YG_EXPORT Node : public ::YGNode {
   void removeChild(size_t index);
 
   void cloneChildrenIfNeeded();
+  void cloneContentsChildrenIfNeeded();
   void markDirtyAndPropagate();
   float resolveFlexGrow() const;
   float resolveFlexShrink() const;