yoga/yoga/node/LayoutableChildren.h

/*
 * 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.
 */

#pragma once

#include <cstdint>
#include <forward_list>
#include <utility>

#include <yoga/enums/Display.h>

namespace facebook::yoga {

class Node;

template <typename T>
class LayoutableChildren {
 public:
  struct Iterator {
    using iterator_category = std::input_iterator_tag;
    using difference_type = std::ptrdiff_t;
    using value_type = T*;
    using pointer = T*;
    using reference = T*;

    Iterator() = default;

    Iterator(const T* node, size_t childIndex)
        : node_(node), childIndex_(childIndex) {}

    T* operator*() const {
      return node_->getChild(childIndex_);
    }

    Iterator& operator++() {
      next();
      return *this;
    }

    Iterator operator++(int) {
      Iterator tmp = *this;
      ++(*this);
      return tmp;
    }

    friend bool operator==(const Iterator& a, const Iterator& b) {
      return a.node_ == b.node_ && a.childIndex_ == b.childIndex_;
    }

    friend bool operator!=(const Iterator& a, const Iterator& b) {
      return a.node_ != b.node_ || a.childIndex_ != b.childIndex_;
    }

   private:
    void next() {
      if (childIndex_ + 1 >= node_->getChildCount()) {
        // if the current node has no more children, try to backtrack and
        // visit its successor
        if (backtrack_.empty()) [[likely]] {
          // if there are no nodes to backtrack to, the last node has been
          // visited
          *this = Iterator{};
        } else {
          // pop and restore the latest backtrack entry
          const auto& back = backtrack_.front();
          node_ = back.first;
          childIndex_ = back.second;
          backtrack_.pop_front();

          // go to the next node
          next();
        }
      } else {
        // current node has more children to visit, go to next
        ++childIndex_;
        // skip all display: contents nodes, possibly going deeper into the
        // tree
        if (node_->getChild(childIndex_)->style().display() ==
            Display::Contents) [[unlikely]] {
          skipContentsNodes();
        }
      }
    }

    void skipContentsNodes() {
      // get the node that would be returned from the iterator
      auto currentNode = node_->getChild(childIndex_);
      while (currentNode->style().display() == Display::Contents &&
             currentNode->getChildCount() > 0) {
        // if it has display: contents set, it shouldn't be returned but its
        // children should in its place push the current node and child index
        // so that the current state can be restored when backtracking
        backtrack_.push_front({node_, childIndex_});
        // traverse the child
        node_ = currentNode;
        childIndex_ = 0;

        // repeat until a node without display: contents is found in the
        // subtree or a leaf is reached
        currentNode = currentNode->getChild(childIndex_);
      }

      // if no node without display: contents was found, try to backtrack
      if (currentNode->style().display() == Display::Contents) {
        next();
      }
    }

    const T* node_{nullptr};
    size_t childIndex_{0};
    std::forward_list<std::pair<const T*, size_t>> backtrack_;

    friend LayoutableChildren;
  };

  explicit LayoutableChildren(const T* node) : node_(node) {
    static_assert(std::input_iterator<LayoutableChildren<T>::Iterator>);
    static_assert(
        std::is_base_of<Node, T>::value,
        "Type parameter of LayoutableChildren must derive from yoga::Node");
  }

  Iterator begin() const {
    if (node_->getChildCount() > 0) {
      auto result = Iterator(node_, 0);
      if (node_->getChild(0)->style().display() == Display::Contents)
          [[unlikely]] {
        result.skipContentsNodes();
      }
      return result;
    } else {
      return Iterator{};
    }
  }

  Iterator end() const {
    return Iterator{};
  }

 private:
  const T* node_;
};

} // namespace facebook::yoga