TypeScript implementation of tree traversing iterator

treeIterator.ts

interface TreeNode<T> {
    value: T
    children?: TreeNode<T>[]
}

function preTraverse<T>(node: TreeNode<T>): IterableIterator<T> {
    const stack = [{ node, index: -1 }]
    
    return {
        [Symbol.iterator]() {
            return this
        },
        next() {
            while (stack.length) {
                const top = stack[stack.length - 1]
                
                if (top.index === -1) {
                    ++top.index
                    return { value: top.node.value }
                }
                
                if (top.node.children && top.index < top.node.children.length) {
                    stack.push({ node: top.node.children[top.index++], index: -1 })
                } else {
                    stack.pop();
                }
            }
            
            return { value: undefined, done: true }
        }
    }
}

function postTraverse<T>(node: TreeNode<T>): IterableIterator<T> {
    const stack = [{ node, index: 0 }]
    
    return {
        [Symbol.iterator]() {
            return this
        },
        next() {
            while (stack.length) {
                const top = stack[stack.length - 1]
                
                if (top.node.children && top.index < top.node.children.length) {
                    stack.push({
                        node: top.node.children[top.index], 
                        index: 0,
                    })
                    top.index++
                } else {
                    stack.pop()
                    return { value: top.node.value }
                }
            }
            
            return { value: undefined, done: true }
        }
    };
}

Usage

const tree = {
    value: 0,
    children: [
        {
            value: 1,
            children: [
                {
                    value: 2,
                },
                {
                    value: 3,
                }
            ]
        },
        {
            value: 4,
            children: [
                {
                    value: 5,
                },
                {
                    value: 6,
                }
            ]
        }
    ]
}

for (const value of preTraverse(tree)) {
    console.log(value)
}
// 0 -> 1 -> 2 -> 3 -> 4 -> 5 -> 6

for (const value of postTraverse(tree)) {
    console.log(value)
}
// 2 -> 3 -> 1 -> 5 -> 6 -> 4 -> 0

Note

Both implementation tries to consume minimum memory as possible as I can. They're implemented by me at first, then refined by Claude 3.5 sonnet.

In both case, they consume memory at most O(d) where d is the maximum depth of given tree. Note that there is no children.reverse to prevent pushing prematurely their children.