A general depth-first traversal utility function in JavaScript

depth_first_traverse.js

// This is a depth-first traversal function with a few nice features.
//
// * Call this function as depthFirstTraverse(root, tree, fn)
//   `tree` is an object whose properties are nodes in a tree; the
//          values are arrays of that node's child nodes.
//   `root` is the starting point for the traversal; a property in `tree`.
//   `fn` is called as in fn(node, depth, childNum) for each node.
//          childNum is the index of node as a child of its parent;
//          as a special case, the childNum of `root` is undefined.
//          `depth` is 0 for the root, and in general indicates how
//          many edges `node` is away from the root.
//
// Features:
//
// 1. You can optionally send in both fn1() and fn2().
//    They receive the same arguments. The difference is that fn1() is
//    called before the node's children are traversed; fn2() is called
//    afterwards. (So each node would receive two calls, one from fn1()
//    then later one from fn2().)
//
// 2. You can provide a fifth argument called `opts`, an object with
//    property 'nodeSet' whose value is an object specifying which nodes
//    to include. fn1() and fn2() are only called on nodes where
//    `node in opts.nodeSet` is true. If `opts` is missing, all ndoes
//    are included.
//
// 3. Your fn1() can return the string 'skip', in which case all the
//    children of the current node are skipped.
//
// 4. Your fn1() can return the string 'break', in which case traversal
//    stops; neither fn1() nor fn2() will be called again.
//

function depthFirstTraverse(root, tree, fn1, fn2, opts) {
    let {depth, childNum, nodeSet} = opts || {};
    if (depth === undefined) depth = 0;
    let reply = fn1(root, depth, childNum);
    if (reply === 'break') return reply;
    if (reply !== 'skip' && tree[root] !== undefined) {
        for (const [i, node] of tree[root].entries()) {
            if (nodeSet && !(node in nodeSet)) continue;
            reply = depthFirstTraverse(node, tree, fn1, fn2,
                {depth: depth + 1, childNum: i, nodeSet});
            if (reply === 'break') return reply;
        }
    }
    if (fn2) fn2(root, depth, childNum);
}