Use the iterator for traversing children

yoga/algorithm/Baseline.cpp

@@ -32,9 +32,7 @@ float calculateBaseline(const yoga::Node* node) {
   }
 
   yoga::Node* baselineChild = nullptr;
-  const size_t childCount = node->getLayoutChildCount();
-  for (size_t i = 0; i < childCount; i++) {
-    auto child = node->getLayoutChild(i);
+  for (auto child : node->getLayoutChildren()) {
     if (child->getLineIndex() > 0) {
       break;
     }
@@ -67,9 +65,7 @@ bool isBaselineLayout(const yoga::Node* node) {
   if (node->style().alignItems() == Align::Baseline) {
     return true;
   }
-  const auto childCount = node->getLayoutChildCount();
-  for (size_t i = 0; i < childCount; i++) {
-    auto child = node->getLayoutChild(i);
+  for (auto child : node->getLayoutChildren()) {
     if (child->style().positionType() != PositionType::Absolute &&
         child->style().alignSelf() == Align::Baseline) {
       return true;

yoga/algorithm/CalculateLayout.cpp

@@ -1454,9 +1454,9 @@ static void calculateLayoutImpl(
   }
   // STEP 4: COLLECT FLEX ITEMS INTO FLEX LINES
 
-  // Indexes of children that represent the first and last items in the line.
-  size_t startOfLineIndex = 0;
-  size_t endOfLineIndex = 0;
+  // Iterator representing the beginning of the current line
+  Node::LayoutableChildren::Iterator startOfLineIterator =
+      node->getLayoutChildren().begin();
 
   // Number of lines.
   size_t lineCount = 0;
@@ -1469,20 +1469,17 @@ static void calculateLayoutImpl(
 
   // Max main dimension of all the lines.
   float maxLineMainDim = 0;
-  for (; endOfLineIndex < childCount;
-       lineCount++, startOfLineIndex = endOfLineIndex) {
+  for (; startOfLineIterator != node->getLayoutChildren().end(); lineCount++) {
     auto flexLine = calculateFlexLine(
         node,
         ownerDirection,
         ownerWidth,
         mainAxisOwnerSize,
         availableInnerWidth,
         availableInnerMainDim,
-        startOfLineIndex,
+        startOfLineIterator,
         lineCount);
 
-    endOfLineIndex = flexLine.endOfLineIndex;
-
     // If we don't need to measure the cross axis, we can skip the entire flex
     // step.
     const bool canSkipFlex =
@@ -1834,17 +1831,18 @@ static void calculateLayoutImpl(
       case Align::Baseline:
         break;
     }
-    size_t endIndex = 0;
+    Node::LayoutableChildren::Iterator endIterator =
+        node->getLayoutChildren().begin();
     for (size_t i = 0; i < lineCount; i++) {
-      const size_t startIndex = endIndex;
-      size_t ii = startIndex;
+      const Node::LayoutableChildren::Iterator startIterator = endIterator;
+      auto iterator = startIterator;
 
       // compute the line's height and find the endIndex
       float lineHeight = 0;
       float maxAscentForCurrentLine = 0;
       float maxDescentForCurrentLine = 0;
-      for (; ii < childCount; ii++) {
-        const auto child = node->getLayoutChild(ii);
+      for (; iterator != node->getLayoutChildren().end(); iterator++) {
+        const auto child = *iterator;
         if (child->style().display() == Display::None) {
           continue;
         }
@@ -1877,11 +1875,11 @@ static void calculateLayoutImpl(
           }
         }
       }
-      endIndex = ii;
+      endIterator = iterator;
       currentLead += i != 0 ? crossAxisGap : 0;
 
-      for (ii = startIndex; ii < endIndex; ii++) {
-        const auto child = node->getLayoutChild(ii);
+      for (iterator = startIterator; iterator != endIterator; iterator++) {
+        const auto child = *iterator;
         if (child->style().display() == Display::None) {
           continue;
         }
@@ -2084,8 +2082,7 @@ static void calculateLayoutImpl(
   // As we only wrapped in normal direction yet, we need to reverse the
   // positions on wrap-reverse.
   if (performLayout && node->style().flexWrap() == Wrap::WrapReverse) {
-    for (size_t i = 0; i < childCount; i++) {
-      const auto child = node->getLayoutChild(i);
+    for (auto child : node->getLayoutChildren()) {
       if (child->style().positionType() != PositionType::Absolute) {
         child->setLayoutPosition(
             node->getLayout().measuredDimension(dimension(crossAxis)) -
@@ -2102,8 +2099,7 @@ static void calculateLayoutImpl(
     const bool needsCrossTrailingPos = needsTrailingPosition(crossAxis);
 
     if (needsMainTrailingPos || needsCrossTrailingPos) {
-      for (size_t i = 0; i < childCount; i++) {
-        const auto child = node->getLayoutChild(i);
+      for (auto child : node->getLayoutChildren()) {
         // Absolute children will be handled by their containing block since we
         // cannot guarantee that their positions are set when their parents are
         // done with layout.

yoga/algorithm/FlexLine.cpp

@@ -20,17 +20,17 @@ FlexLine calculateFlexLine(
     const float mainAxisownerSize,
     const float availableInnerWidth,
     const float availableInnerMainDim,
-    const size_t startOfLineIndex,
+    Node::LayoutableChildren::Iterator& iterator,
     const size_t lineCount) {
   std::vector<yoga::Node*> itemsInFlow;
-  itemsInFlow.reserve(node->getLayoutChildren().size());
+  itemsInFlow.reserve(node->getChildCount());
 
   float sizeConsumed = 0.0f;
   float totalFlexGrowFactors = 0.0f;
   float totalFlexShrinkScaledFactors = 0.0f;
   size_t numberOfAutoMargins = 0;
-  size_t endOfLineIndex = startOfLineIndex;
-  size_t firstElementInLineIndex = startOfLineIndex;
+  size_t endOfLineIndex = iterator.index();
+  size_t firstElementInLineIndex = iterator.index();
 
   float sizeConsumedIncludingMinConstraint = 0;
   const Direction direction = node->resolveDirection(ownerDirection);
@@ -41,8 +41,9 @@ FlexLine calculateFlexLine(
       node->style().computeGapForAxis(mainAxis, availableInnerMainDim);
 
   // Add items to the current line until it's full or we run out of items.
-  for (; endOfLineIndex < node->getLayoutChildren().size(); endOfLineIndex++) {
-    auto child = node->getLayoutChild(endOfLineIndex);
+  for (; iterator != node->getLayoutChildren().end();
+       iterator++, endOfLineIndex = iterator.index()) {
+    auto child = *iterator;
     if (child->style().display() == Display::None ||
         child->style().positionType() == PositionType::Absolute) {
       if (firstElementInLineIndex == endOfLineIndex) {

yoga/algorithm/FlexLine.h

@@ -72,7 +72,7 @@ FlexLine calculateFlexLine(
     float mainAxisownerSize,
     float availableInnerWidth,
     float availableInnerMainDim,
-    size_t startOfLineIndex,
+    Node::LayoutableChildren::Iterator& iterator,
     size_t lineCount);
 
 } // namespace facebook::yoga

yoga/node/Node.cpp

@@ -41,6 +41,7 @@ Node::Node(Node&& node) noexcept
       style_(std::move(node.style_)),
       layout_(node.layout_),
       lineIndex_(node.lineIndex_),
+      contentsChildren_(node.contentsChildren_),
       owner_(node.owner_),
       children_(std::move(node.children_)),
       config_(node.config_),
@@ -116,14 +117,37 @@ void Node::setMeasureFunc(YGMeasureFunc measureFunc) {
 }
 
 void Node::replaceChild(Node* child, size_t index) {
+  auto previousChild = children_[index];
+  if (previousChild->style().display() == Display::Contents &&
+      child->style().display() != Display::Contents) {
+    contentsChildren_--;
+  } else if (
+      previousChild->style().display() != Display::Contents &&
+      child->style().display() == Display::Contents) {
+    contentsChildren_++;
+  }
+
   children_[index] = child;
 }
 
 void Node::replaceChild(Node* oldChild, Node* newChild) {
+  if (oldChild->style().display() == Display::Contents &&
+      newChild->style().display() != Display::Contents) {
+    contentsChildren_--;
+  } else if (
+      oldChild->style().display() != Display::Contents &&
+      newChild->style().display() == Display::Contents) {
+    contentsChildren_++;
+  }
+
   std::replace(children_.begin(), children_.end(), oldChild, newChild);
 }
 
 void Node::insertChild(Node* child, size_t index) {
+  if (child->style().display() == Display::Contents) {
+    contentsChildren_++;
+  }
+
   children_.insert(children_.begin() + static_cast<ptrdiff_t>(index), child);
 }
 
@@ -160,13 +184,21 @@ void Node::setDirty(bool isDirty) {
 bool Node::removeChild(Node* child) {
   auto p = std::find(children_.begin(), children_.end(), child);
   if (p != children_.end()) {
+    if (child->style().display() == Display::Contents) {
+      contentsChildren_--;
+    }
+
     children_.erase(p);
     return true;
   }
   return false;
 }
 
 void Node::removeChild(size_t index) {
+  if (children_[index]->style().display() == Display::Contents) {
+    contentsChildren_--;
+  }
+
   children_.erase(children_.begin() + static_cast<ptrdiff_t>(index));
 }
 

yoga/node/Node.h

@@ -266,21 +266,22 @@ class YG_EXPORT Node : public ::YGNode {
     return children_.size();
   }
 
-  // This needs to be optimized, not to create a new vector every time
-  const std::vector<Node*> getLayoutChildren() const {
-    std::vector<Node*> result;
-    for (auto child : LayoutableChildren(this)) {
-      result.push_back(child);
-    }
-    return result;
-  }
-
-  Node* getLayoutChild(size_t index) const {
-    return getLayoutChildren()[index];
+  const LayoutableChildren getLayoutChildren() const {
+    return LayoutableChildren(this);
   }
 
   size_t getLayoutChildCount() const {
-    return getLayoutChildren().size();
+    if (contentsChildren_ == 0) {
+      return children_.size();
+    } else {
+      size_t count = 0;
+      for (auto iter = getLayoutChildren().begin();
+           iter != getLayoutChildren().end();
+           iter++) {
+        count++;
+      }
+      return count;
+    }
   }
 
   const Config* getConfig() const {
@@ -437,6 +438,7 @@ class YG_EXPORT Node : public ::YGNode {
   Style style_;
   LayoutResults layout_;
   size_t lineIndex_ = 0;
+  size_t contentsChildren_ = 0;
   Node* owner_ = nullptr;
   std::vector<Node*> children_;
   const Config* config_;