Fix YogaLayoutableShadowNode handling of non-layoutable children (#36325)

Summary:
Pull Request resolved: https://github.com/facebook/react-native/pull/36325

A lot of the code in YogaLayoutableShadowNode expects a 1:1 mapping from ShadowNode children to Yoga Node children, which is not always the case, e.g. for `RawTextShadowNode` (if text, or a number coerced to text, is rendered outside a Text component).

In S323291 we saw this cause memory corruption due to later invalid static_cast. I changed the casting mechanism to terminte instead of corrupting memory, but there is logic here that we need to make permissive (I think D29894182 (https://github.com/facebook/react-native/commit/d3e836245b5fab2e0299dd06901de74dff767b63) also tried to previously do this?).

Normally Text will be rendered inside of a ParagraphShadowNode, where "Text" from React is mapped to ParagraphShadowNode (which is layoutable), "VirtualText" is mapped to TextShadowNode (which is also not layoutable), then finally the text fragment is mapped to "RawTextShadowNode". Arguably React renderer behavior should be not to send anything to native, but we can provide a generalized solution in YogaLayoutableShadowNode to handle any other cases of this issue.

This solution works by filtering ShadowNode children to those which are YogaLayoutable, then only ever operating on that list of children. This means a guaranteed invariant of the nodes we operate on being layoutable (vs the adhoc error handling right now for when they are not), and means we maintain the index based mapping of ShadowNode children to Yoga Node children.

Note, there is another similar API, `getLayoutableChildNodes()` which is protected and returns a filtered list of LayoutableShadowNode. This is public, to allow querying/setting layout results, whearas the similar operations in YogaLayoutableShadowNode are instead an implementation detail.

Changelog:
[General][Fixed] - Fix YogaLayoutableShadowNode handling of non-layoutable children

Reviewed By: sammy-SC

Differential Revision: D43657405

fbshipit-source-id: 8ed136b03b4da15a5e7dfbdd5539b04a2952420d

ReactCommon/react/renderer/components/view/YogaLayoutableShadowNode.cpp

@@ -115,6 +115,12 @@ YogaLayoutableShadowNode::YogaLayoutableShadowNode(
               .yogaNode_.isDirty() == yogaNode_.isDirty() &&
       "Yoga node must inherit dirty flag.");
 
+  for (auto &child : getChildren()) {
+    if (auto layoutableChild = traitCast<YogaLayoutableShadowNode>(child)) {
+      yogaLayoutableChildren_.push_back(layoutableChild);
+    }
+  }
+
   yogaNode_.setContext(this);
   yogaNode_.setOwner(nullptr);
   updateYogaChildrenOwnersIfNeeded();
@@ -159,17 +165,17 @@ void YogaLayoutableShadowNode::enableMeasurement() {
       YogaLayoutableShadowNode::yogaNodeMeasureCallbackConnector);
 }
 
-void YogaLayoutableShadowNode::appendYogaChild(ShadowNode const &childNode) {
+void YogaLayoutableShadowNode::appendYogaChild(
+    YogaLayoutableShadowNode::Shared const &childNode) {
   // The caller must check this before calling this method.
   react_native_assert(
       !getTraits().check(ShadowNodeTraits::Trait::LeafYogaNode));
 
   ensureYogaChildrenLookFine();
 
-  auto &layoutableChildNode =
-      traitCast<YogaLayoutableShadowNode const &>(childNode);
+  yogaLayoutableChildren_.push_back(childNode);
   yogaNode_.insertChild(
-      &layoutableChildNode.yogaNode_,
+      &childNode->yogaNode_,
       static_cast<uint32_t>(yogaNode_.getChildren().size()));
 
   ensureYogaChildrenLookFine();
@@ -183,48 +189,21 @@ void YogaLayoutableShadowNode::adoptYogaChild(size_t index) {
   react_native_assert(
       !getTraits().check(ShadowNodeTraits::Trait::LeafYogaNode));
 
-  auto &children = getChildren();
+  auto &childNode =
+      traitCast<YogaLayoutableShadowNode const &>(*getChildren().at(index));
 
-  // Overflow checks.
-  react_native_assert(children.size() > index);
-  react_native_assert(children.size() >= yogaNode_.getChildren().size());
-
-  auto &childNode = *children.at(index);
-
-  auto &layoutableChildNode =
-      traitCast<YogaLayoutableShadowNode const &>(childNode);
-
-  // Note, the following (commented out) assert is conceptually valid but still
-  // might produce false-positive signals because of the ABA problem (different
-  // objects with non-interleaving life-times being allocated on the same
-  // address). react_native_assert(layoutableChildNode.yogaNode_.getOwner() !=
-  // &yogaNode_);
-
-  if (layoutableChildNode.yogaNode_.getOwner() == nullptr) {
+  if (childNode.yogaNode_.getOwner() == nullptr) {
     // The child node is not owned.
-    layoutableChildNode.yogaNode_.setOwner(&yogaNode_);
+    childNode.yogaNode_.setOwner(&yogaNode_);
     // At this point the child yoga node must be already inserted by the caller.
     // react_native_assert(layoutableChildNode.yogaNode_.isDirty());
   } else {
     // The child is owned by some other node, we need to clone that.
     // TODO: At this point, React has wrong reference to the node. (T138668036)
     auto clonedChildNode = childNode.clone({});
-    auto &layoutableClonedChildNode =
-        traitCast<YogaLayoutableShadowNode &>(*clonedChildNode);
-
-    // The owner must be nullptr for a newly cloned node.
-    react_native_assert(
-        layoutableClonedChildNode.yogaNode_.getOwner() == nullptr);
-
-    // Establishing ownership.
-    layoutableClonedChildNode.yogaNode_.setOwner(&yogaNode_);
 
     // Replace the child node with a newly cloned one in the children list.
     replaceChild(childNode, clonedChildNode, static_cast<int>(index));
-
-    // Replace the Yoga node inside the Yoga node children list.
-    yogaNode_.replaceChild(
-        &layoutableClonedChildNode.yogaNode_, static_cast<int>(index));
   }
 
   ensureYogaChildrenLookFine();
@@ -243,35 +222,81 @@ void YogaLayoutableShadowNode::appendChild(
     return;
   }
 
-  // Here we don't have information about the previous structure of the node (if
-  // it that existed before), so we don't have anything to compare the Yoga node
-  // with (like a previous version of this node). Therefore we must dirty the
-  // node.
-  yogaNode_.setDirty(true);
+  if (auto yogaLayoutableChild =
+          traitCast<YogaLayoutableShadowNode>(childNode)) {
+    // Here we don't have information about the previous structure of the node
+    // (if it that existed before), so we don't have anything to compare the
+    // Yoga node with (like a previous version of this node). Therefore we must
+    // dirty the node.
+    yogaNode_.setDirty(true);
 
-  // All children of a non-leaf `YogaLayoutableShadowNode` must be a
-  // `YogaLayoutableShadowNode`s to be appended. This happens when invalid
-  // string/numeric child is passed which is not YogaLayoutableShadowNode
-  // (e.g. RCTRawText). This used to throw an error, but we are ignoring it
-  // because we want core library components to be fault-tolerant and degrade
-  // gracefully. A soft error will be emitted from JavaScript.
-  if (traitCast<YogaLayoutableShadowNode const *>(childNode.get()) != nullptr) {
     // Appending the Yoga node.
-    appendYogaChild(*childNode);
+    appendYogaChild(yogaLayoutableChild);
 
     ensureYogaChildrenLookFine();
-    ensureYogaChildrenAlighment();
+    ensureYogaChildrenAlignment();
 
     // Adopting the Yoga node.
     adoptYogaChild(getChildren().size() - 1);
 
     ensureConsistency();
-  } else {
-    react_native_log_error(
-        "Text strings must be rendered within a <Text> component.");
   }
 }
 
+void YogaLayoutableShadowNode::replaceChild(
+    ShadowNode const &oldChild,
+    ShadowNode::Shared const &newChild,
+    size_t suggestedIndex) {
+  LayoutableShadowNode::replaceChild(oldChild, newChild, suggestedIndex);
+
+  ensureUnsealed();
+  ensureYogaChildrenLookFine();
+
+  auto layoutableOldChild =
+      traitCast<YogaLayoutableShadowNode const *>(&oldChild);
+  auto layoutableNewChild = traitCast<YogaLayoutableShadowNode>(newChild);
+
+  if (layoutableOldChild == nullptr && layoutableNewChild == nullptr) {
+    // No need to mutate yogaLayoutableChildren_
+    return;
+  }
+
+  bool suggestedIndexAccurate = suggestedIndex >= 0 &&
+      suggestedIndex < yogaLayoutableChildren_.size() &&
+      yogaLayoutableChildren_[suggestedIndex].get() == layoutableOldChild;
+
+  auto oldChildIter = suggestedIndexAccurate
+      ? yogaLayoutableChildren_.begin() + suggestedIndex
+      : std::find_if(
+            yogaLayoutableChildren_.begin(),
+            yogaLayoutableChildren_.end(),
+            [&](YogaLayoutableShadowNode::Shared const &layoutableChild) {
+              return layoutableChild.get() == layoutableOldChild;
+            });
+  auto oldChildIndex =
+      static_cast<int32_t>(oldChildIter - yogaLayoutableChildren_.begin());
+
+  if (oldChildIter == yogaLayoutableChildren_.end()) {
+    // oldChild does not exist as part of our node
+    return;
+  }
+
+  if (layoutableNewChild) {
+    // Both children are layoutable, replace the old one with the new one
+    react_native_assert(layoutableNewChild->yogaNode_.getOwner() == nullptr);
+    layoutableNewChild->yogaNode_.setOwner(&yogaNode_);
+    *oldChildIter = layoutableNewChild;
+    yogaNode_.replaceChild(&layoutableNewChild->yogaNode_, oldChildIndex);
+  } else {
+    // Layoutable child replaced with non layoutable child. Remove the previous
+    // child from the layoutable children list.
+    yogaLayoutableChildren_.erase(oldChildIter);
+    yogaNode_.removeChild(oldChildIndex);
+  }
+
+  ensureYogaChildrenLookFine();
+}
+
 bool YogaLayoutableShadowNode::doesOwn(
     YogaLayoutableShadowNode const &child) const {
   return child.yogaNode_.getOwner() == &yogaNode_;
@@ -297,23 +322,28 @@ void YogaLayoutableShadowNode::updateYogaChildren() {
 
   auto oldYogaChildren = isClean ? yogaNode_.getChildren() : YGVector{};
   yogaNode_.setChildren({});
+  yogaLayoutableChildren_.clear();
 
   for (size_t i = 0; i < getChildren().size(); i++) {
-    appendYogaChild(*getChildren().at(i));
-    adoptYogaChild(i);
+    if (auto yogaLayoutableChild =
+            traitCast<YogaLayoutableShadowNode>(getChildren()[i])) {
+      appendYogaChild(yogaLayoutableChild);
+      adoptYogaChild(i);
 
-    if (isClean) {
-      auto &oldYogaChildNode = *oldYogaChildren.at(i);
-      auto &newYogaChildNode =
-          traitCast<YogaLayoutableShadowNode const &>(*getChildren().at(i))
-              .yogaNode_;
+      if (isClean) {
+        auto yogaChildIndex = yogaLayoutableChildren_.size() - 1;
+        auto &oldYogaChildNode = *oldYogaChildren.at(yogaChildIndex);
+        auto &newYogaChildNode =
+            yogaLayoutableChildren_.at(yogaChildIndex)->yogaNode_;
 
-      isClean = isClean && !newYogaChildNode.isDirty() &&
-          (newYogaChildNode.getStyle() == oldYogaChildNode.getStyle());
+        isClean = isClean && !newYogaChildNode.isDirty() &&
+            (newYogaChildNode.getStyle() == oldYogaChildNode.getStyle());
+      }
     }
   }
 
-  react_native_assert(getChildren().size() == yogaNode_.getChildren().size());
+  react_native_assert(
+      yogaLayoutableChildren_.size() == yogaNode_.getChildren().size());
 
   yogaNode_.setDirty(!isClean);
 }
@@ -720,12 +750,9 @@ void YogaLayoutableShadowNode::swapLeftAndRightInTree(
   swapLeftAndRightInYogaStyleProps(shadowNode);
   swapLeftAndRightInViewProps(shadowNode);
 
-  for (auto &child : shadowNode.getChildren()) {
-    auto const yogaLayoutableChild =
-        traitCast<YogaLayoutableShadowNode const *>(child.get());
-    if ((yogaLayoutableChild != nullptr) &&
-        !yogaLayoutableChild->doesOwn(shadowNode)) {
-      swapLeftAndRightInTree(*yogaLayoutableChild);
+  for (auto &child : shadowNode.yogaLayoutableChildren_) {
+    if (!child->doesOwn(shadowNode)) {
+      swapLeftAndRightInTree(*child);
     }
   }
 }
@@ -837,7 +864,7 @@ void YogaLayoutableShadowNode::swapLeftAndRightInViewProps(
 
 void YogaLayoutableShadowNode::ensureConsistency() const {
   ensureYogaChildrenLookFine();
-  ensureYogaChildrenAlighment();
+  ensureYogaChildrenAlignment();
   ensureYogaChildrenOwnersConsistency();
 }
 
@@ -874,7 +901,7 @@ void YogaLayoutableShadowNode::ensureYogaChildrenLookFine() const {
 #endif
 }
 
-void YogaLayoutableShadowNode::ensureYogaChildrenAlighment() const {
+void YogaLayoutableShadowNode::ensureYogaChildrenAlignment() const {
 #ifdef REACT_NATIVE_DEBUG
   // If the node is not a leaf node, checking that:
   // - All children are `YogaLayoutableShadowNode` subclasses.
@@ -882,7 +909,7 @@ void YogaLayoutableShadowNode::ensureYogaChildrenAlighment() const {
   //   this node.
 
   auto &yogaChildren = yogaNode_.getChildren();
-  auto &children = getChildren();
+  auto &children = yogaLayoutableChildren_;
 
   if (getTraits().check(ShadowNodeTraits::Trait::LeafYogaNode)) {
     react_native_assert(yogaChildren.empty());

ReactCommon/react/renderer/components/view/YogaLayoutableShadowNode.h

@@ -26,9 +26,9 @@ class YogaLayoutableShadowNode : public LayoutableShadowNode {
   using CompactValue = facebook::yoga::detail::CompactValue;
 
  public:
-  using UnsharedList = butter::small_vector<
-      YogaLayoutableShadowNode *,
-      kShadowNodeChildrenSmallVectorSize>;
+  using Shared = std::shared_ptr<YogaLayoutableShadowNode const>;
+  using ListOfShared =
+      butter::small_vector<Shared, kShadowNodeChildrenSmallVectorSize>;
 
   static ShadowNodeTraits BaseTraits();
   static ShadowNodeTraits::Trait IdentifierTrait();
@@ -52,7 +52,11 @@ class YogaLayoutableShadowNode : public LayoutableShadowNode {
    */
   void enableMeasurement();
 
-  void appendChild(ShadowNode::Shared const &child);
+  void appendChild(ShadowNode::Shared const &child) override;
+  void replaceChild(
+      ShadowNode const &oldChild,
+      ShadowNode::Shared const &newChild,
+      size_t suggestedIndex = -1) override;
 
   void updateYogaChildren();
 
@@ -121,7 +125,7 @@ class YogaLayoutableShadowNode : public LayoutableShadowNode {
    * The method does *not* do anything besides that (no cloning or `owner` field
    * adjustment).
    */
-  void appendYogaChild(ShadowNode const &childNode);
+  void appendYogaChild(YogaLayoutableShadowNode::Shared const &childNode);
 
   /*
    * Makes the child node with a given `index` (and Yoga node associated with) a
@@ -187,9 +191,15 @@ class YogaLayoutableShadowNode : public LayoutableShadowNode {
 #pragma mark - Consistency Ensuring Helpers
 
   void ensureConsistency() const;
-  void ensureYogaChildrenAlighment() const;
+  void ensureYogaChildrenAlignment() const;
   void ensureYogaChildrenOwnersConsistency() const;
   void ensureYogaChildrenLookFine() const;
+
+#pragma mark - Private member variables
+  /*
+   * List of children which derive from YogaLayoutableShadowNode
+   */
+  ListOfShared yogaLayoutableChildren_;
 };
 
 } // namespace react

ReactCommon/react/renderer/core/ShadowNode.cpp

@@ -228,7 +228,7 @@ void ShadowNode::appendChild(const ShadowNode::Shared &child) {
 void ShadowNode::replaceChild(
     ShadowNode const &oldChild,
     ShadowNode::Shared const &newChild,
-    int suggestedIndex) {
+    size_t suggestedIndex) {
   ensureUnsealed();
 
   cloneChildrenIfShared();
@@ -239,7 +239,7 @@ void ShadowNode::replaceChild(
       *std::const_pointer_cast<ShadowNode::ListOfShared>(children_);
   auto size = children.size();
 
-  if (suggestedIndex != -1 && static_cast<size_t>(suggestedIndex) < size) {
+  if (suggestedIndex != -1 && suggestedIndex < size) {
     // If provided `suggestedIndex` is accurate,
     // replacing in place using the index.
     if (children.at(suggestedIndex).get() == &oldChild) {

ReactCommon/react/renderer/core/ShadowNode.h

@@ -161,11 +161,11 @@ class ShadowNode : public Sealable, public DebugStringConvertible {
 
 #pragma mark - Mutating Methods
 
-  void appendChild(Shared const &child);
-  void replaceChild(
+  virtual void appendChild(Shared const &child);
+  virtual void replaceChild(
       ShadowNode const &oldChild,
       Shared const &newChild,
-      int suggestedIndex = -1);
+      size_t suggestedIndex = -1);
 
   /*
    * Performs all side effects associated with mounting/unmounting in one place.