traverser.ts

import { VisitorKeys } from 'estraverse';
import * as ESTree from 'estree';

type Cursor = [number, number];

function init(node: ESTree.Node): Cursor | null {
  const keys: string[] | undefined = (VisitorKeys as any)[node.type];
  if (!keys || keys.length <= 0) {
    return null;
  }
  return [0, 0];
}

function move(node: ESTree.Node, [i, j]: Cursor): [ESTree.Node | undefined, Cursor | null] {
  const keys: string[] = (VisitorKeys as any)[node.type];
  const child: ESTree.Node | ESTree.Node[] = (node as any)[keys[i]];
  if (Array.isArray(child)) {
    return [child[j], j < child.length ? [i, j + 1] : i < keys.length ? [i + 1, 0] : null];
  }
  return [child, i < keys.length ? [i + 1, 0] : null]
}

export interface Frame {
  node: ESTree.Node;
  cursor: Cursor | null;
}

export interface Step {
  type: 'enter' | 'leave',
  node: ESTree.Node,
}

export class Traverser {
  private stack: Frame[] = [];

  private initNode: ESTree.Node | null;

  constructor(node: ESTree.Node) {
    this.initNode = node;
  }

  next(): Step | null {
    if (this.initNode) {
      const node = this.initNode;
      this.initNode = null;

      this.stack.push({
        node,
        cursor: init(node),
      });

      return {
        type: 'enter',
        node,
      };
    }

    const top = this.stack[this.stack.length - 1];

    if (!top) {
      return null
    }

    if (!top.cursor) {
      this.stack.pop();
      return {
        type: 'leave',
        node: top.node,
      };
    }

    const [child, nextCursor] = move(top.node, top.cursor);
    top.cursor = nextCursor;

    if (!child) {
      return this.next();
    } else {
      this.stack.push({
        node: child,
        cursor: init(child),
      });
      return {
        type: 'enter',
        node: child,
      };
    }
  }
}