lazytree.nim

# If `nextChildIndex` is `rootIndex`, this means the next yielded node will be
# the root node. The actual value is quite intentional because `inc`ing it will
# give the first child index 0.
const rootIndex = -1

type
  Iterator[N] = ref object
    root: N
    nextChildIndex: int

proc newIterator[N](root: N): Iterator[N] =
  Iterator[N](root: root, nextChildIndex: rootIndex)

func hasNextNode[N](iter: Iterator[N]): bool =
  ## Return `true` if the iterator has any more child nodes to yield.
  ## Otherwise return `false`.
  if iter.nextChildIndex == rootIndex:
    return true
  return iter.nextChildIndex < len(iter.root)

proc nextNode[N](iter: Iterator[N]): N =
  ## Return next node from `iter`.
  ## Precondition: `hasNextNode` must have returned `true`.
  if iter.nextChildIndex == rootIndex:
    result = iter.root
  else:
    result = iter.root[iter.nextChildIndex]
  inc(iter.nextChildIndex)

iterator treeIter*[N](root: N): N =
  var rootIterator = newIterator[N](root)
  while rootIterator.hasNextNode():
    yield rootIterator.nextNode()