TreeTraversalState.h

#pragma once
 
#include <CesiumUtility/Assert.h>
 
#include <cstddef>
#include <unordered_map>
#include <vector>
 
namespace CesiumUtility {

template <typename TNodePointer, typename TState> class TreeTraversalState {
public:
  size_t getNodeCountInPreviousTraversal() const {
    return this->_previousTraversal.size();
  }

  size_t getNodeCountInCurrentTraversal() const {
    return this->_currentTraversal.size();
  }

  void beginTraversal() {
    // If this assertion fails, it indicates a traversal is already in progress.
    CESIUM_ASSERT(this->_parentIndices.empty());
 
    this->_previousTraversal.swap(this->_currentTraversal);
    this->_currentTraversal.clear();
    this->_previousTraversalNextNodeIndex = 0;
  }

  void beginNode(const TNodePointer& pNode) {
    int64_t currentTraversalIndex = int64_t(this->_currentTraversal.size());
    int64_t previousTraversalIndex = this->_previousTraversalNextNodeIndex;
 
    if (previousTraversalIndex >= 0 &&
        size_t(previousTraversalIndex) < this->_previousTraversal.size()) {
      const TraversalData& previousData =
          this->_previousTraversal[size_t(previousTraversalIndex)];
      if (previousData.pNode == pNode) {
        ++this->_previousTraversalNextNodeIndex;
      } else {
        previousTraversalIndex = -1;
      }
    } else {
      previousTraversalIndex = -1;
    }
 
    this->_parentIndices.emplace_back(TraversalIndices{
        .previous = previousTraversalIndex,
        .current = currentTraversalIndex});
 
    this->_currentTraversal.emplace_back(TraversalData{pNode, -1, TState()});
  }

  TNodePointer getCurrentNode() const noexcept {
    if (this->_parentIndices.empty())
      return nullptr;
 
    return this->currentData().pNode;
  }

  bool wasCurrentNodePreviouslyTraversed() const noexcept {
    return this->previousState() != nullptr;
  }

  const TState* previousState() const {
    const TraversalData* pData = this->previousData();
    return pData ? &pData->state : nullptr;
  }

  TState& currentState() {
    TraversalData& data = this->currentData();
    return data.state;
  }

  const TState& currentState() const {
    const TraversalData& data = this->currentData();
    return data.state;
  }

  void finishNode([[maybe_unused]] const TNodePointer& pNode) {
    // An assertion failure here indicates mismatched calls to beginNode /
    // finishNode.
    CESIUM_ASSERT(!this->_currentTraversal.empty());
    CESIUM_ASSERT(!this->_parentIndices.empty());
    CESIUM_ASSERT(this->currentData().pNode == pNode);
 
    this->currentData().nextSiblingIndex =
        int64_t(this->_currentTraversal.size());
 
    // Now that this node is done, skip its subtree, if any, in the previous
    // traversal. If this finished node doesn't exist in the previous traversal,
    // look for the next node in the previous traversal at the current position.
    const TraversalData* pPreviousData = this->previousData();
    if (pPreviousData) {
      CESIUM_ASSERT(pPreviousData->nextSiblingIndex >= 0);
      this->_previousTraversalNextNodeIndex = pPreviousData->nextSiblingIndex;
    }
 
    this->_parentIndices.pop_back();
  }

  template <typename Func> void forEachPreviousChild(Func&& callback) const {
    int64_t parentPreviousIndex = this->previousDataIndex();
    if (parentPreviousIndex < 0) {
      // Current node was not previously traversed.
      return;
    }
 
    const TraversalData& parentPreviousData =
        this->_previousTraversal[size_t(parentPreviousIndex)];
 
    for (size_t i = size_t(parentPreviousIndex + 1);
         i < size_t(parentPreviousData.nextSiblingIndex);) {
      CESIUM_ASSERT(i < this->_previousTraversal.size());
      const TraversalData& data = this->_previousTraversal[i];
      callback(data.pNode, data.state);
      CESIUM_ASSERT(
          data.nextSiblingIndex >= 0 && size_t(data.nextSiblingIndex) > i);
      i = size_t(data.nextSiblingIndex);
    }
  }

  template <typename Func>
  void forEachPreviousDescendant(Func&& callback) const {
    int64_t parentPreviousIndex = this->previousDataIndex();
    if (parentPreviousIndex < 0) {
      // Current node was not previously traversed.
      return;
    }
 
    const TraversalData& parentPreviousData =
        this->_previousTraversal[size_t(parentPreviousIndex)];
 
    for (size_t i = size_t(parentPreviousIndex + 1);
         i < size_t(parentPreviousData.nextSiblingIndex);
         ++i) {
      CESIUM_ASSERT(i < this->_previousTraversal.size());
      const TraversalData& data = this->_previousTraversal[i];
      callback(data.pNode, data.state);
    }
  }

  template <typename Func> void forEachCurrentDescendant(Func&& callback) {
    int64_t parentCurrentIndex = this->currentDataIndex();
    CESIUM_ASSERT(parentCurrentIndex >= 0);
 
    const TraversalData& parentCurrentData =
        this->_currentTraversal[size_t(parentCurrentIndex)];
 
    size_t endIndex = parentCurrentData.nextSiblingIndex >= 0
                          ? size_t(parentCurrentData.nextSiblingIndex)
                          : this->_currentTraversal.size();
 
    for (size_t i = size_t(parentCurrentIndex + 1); i < endIndex; ++i) {
      TraversalData& data = this->_currentTraversal[i];
      callback(data.pNode, data.state);
    }
  }

  std::unordered_map<TNodePointer, TState> slowlyGetCurrentStates() const {
    return this->slowlyGetStates(this->_currentTraversal);
  }

  std::unordered_map<TNodePointer, TState> slowlyGetPreviousStates() const {
    return this->slowlyGetStates(this->_previousTraversal);
  }
 
private:
  struct TraversalData {
    TNodePointer pNode;
    int64_t nextSiblingIndex;
    TState state;
  };
 
  struct TraversalIndices {
    int64_t previous;
    int64_t current;
  };
 
  int64_t previousDataIndex() const {
    CESIUM_ASSERT(!this->_parentIndices.empty());
 
    int64_t result = this->_parentIndices.back().previous;
 
    CESIUM_ASSERT(
        result == -1 ||
        (result >= 0 && size_t(result) < this->_previousTraversal.size()));
 
    return result;
  }
 
  int64_t currentDataIndex() const {
    CESIUM_ASSERT(!this->_parentIndices.empty());
 
    // An assertion failure here may indicate beginTraversal wasn't called.
    CESIUM_ASSERT(
        this->_parentIndices.back().current >= 0 &&
        this->_parentIndices.back().current <
            static_cast<int64_t>(this->_currentTraversal.size()));
 
    return this->_parentIndices.back().current;
  }
 
  const TraversalData* previousData() const {
    int64_t previousIndex = this->previousDataIndex();
    if (previousIndex < 0)
      return nullptr;
 
    const TraversalData& previousData =
        this->_previousTraversal[size_t(previousIndex)];
 
    CESIUM_ASSERT(previousData.pNode == this->currentData().pNode);
 
    return &previousData;
  }
 
  TraversalData& currentData() {
    int64_t currentIndex = this->currentDataIndex();
    CESIUM_ASSERT(currentIndex >= 0);
    return this->_currentTraversal[size_t(currentIndex)];
  }
 
  const TraversalData& currentData() const {
    int64_t currentIndex = this->currentDataIndex();
    CESIUM_ASSERT(currentIndex >= 0);
    return this->_currentTraversal[size_t(currentIndex)];
  }
 
  std::unordered_map<TNodePointer, TState>
  slowlyGetStates(const std::vector<TraversalData>& traversalData) const {
    std::unordered_map<TNodePointer, TState> result;
 
    for (const TraversalData& data : traversalData) {
      result[data.pNode] = data.state;
    }
 
    return result;
  }
 
  std::vector<TraversalData> _previousTraversal;
  std::vector<TraversalData> _currentTraversal;
 
  // A stack of indices into the previous and current traversals. When
  // `beginNode` is called, the index of that node within each traversal is
  // pushed onto the end of this vector. This is always the index of the _last_
  // node in the `_currentTraversal`, because `beginNode` always adds a new
  // entry to the end of the `currentTraversal`. The previous traversal index
  // may be -1 if the current node was not visited at all in the previous
  // traversal. `finishNode` pops the last entry off the end of this array.
  std::vector<TraversalIndices> _parentIndices;
 
  // The index of the next node in the previous traversal. In `beginNode`, this
  // new node is added to the end of `_currentTraversal`. In
  // `_previousTraversal`, if it exists at all, it will be found at this index.
  int64_t _previousTraversalNextNodeIndex = 0;
};
 
} // namespace CesiumUtility