AST Traversal

/**
 * Traverse an AST with visitor functions and ancestry tracking
 * @param node - Root node to start traversal from
 * @param handlers - Visitor functions (enter/exit) or single handler function
 * @param state - Optional state object passed to all handlers
 */
function traverse<T>(
  node: t.Node,
  handlers: TraversalHandler<T> | TraversalHandlers<T>,
  state?: T,
): void;

/**
 * Single handler function signature
 */
type TraversalHandler<T> = (
  node: t.Node,
  parent: TraversalAncestors,
  state: T,
) => void;

/**
 * Handler configuration object with enter/exit functions
 */
interface TraversalHandlers<T> {
  enter?: TraversalHandler<T>;
  exit?: TraversalHandler<T>;
}

/**
 * Array of ancestor information providing context during traversal
 */
type TraversalAncestors = Array<{
  node: t.Node;
  key: string;
  index?: number;
}>;

/**
 * Fast traversal with simple enter function
 * @param node - Root node to traverse (can be null/undefined)
 * @param enter - Function called for each node during traversal
 * @param opts - Optional options object passed to enter function
 * @returns true if traversal was stopped, false otherwise
 */
function traverseFast<Options = object>(
  node: t.Node | null | undefined,
  enter: (
    node: t.Node,
    opts?: Options,
  ) => void | typeof traverseFast.skip | typeof traverseFast.stop,
  opts?: Options,
): boolean;

// Control symbols for traversal flow
traverseFast.skip: symbol; // Skip current node's children
traverseFast.stop: symbol; // Stop entire traversal

import * as t from "@babel/types";
import { traverse } from "@babel/types";
import { parse } from "@babel/parser";

// Parse some code
const code = `
  function add(a, b) {
    return a + b;
  }
  add(1, 2);
`;
const ast = parse(code);

// Simple traversal with single handler
traverse(ast, (node, ancestors, state) => {
  console.log(`Visiting ${node.type}`);
  
  // Access parent information
  if (ancestors.length > 0) {
    const parent = ancestors[ancestors.length - 1];
    console.log(`  Parent: ${parent.node.type} (key: ${parent.key})`);
  }
});

// Traversal with both enter and exit handlers
traverse(ast, {
  enter(node, ancestors, state) {
    console.log(`Entering ${node.type}`);
  },
  exit(node, ancestors, state) {
    console.log(`Exiting ${node.type}`);
  }
});

interface CollectorState {
  identifiers: string[];
  functions: string[];
}

const state: CollectorState = {
  identifiers: [],
  functions: []
};

traverse(ast, (node, ancestors, state) => {
  if (t.isIdentifier(node)) {
    state.identifiers.push(node.name);
  }
  
  if (t.isFunctionDeclaration(node) && node.id) {
    state.functions.push(node.id.name);
  }
}, state);

console.log("Identifiers:", state.identifiers);
console.log("Functions:", state.functions);

traverse(ast, (node, ancestors, state) => {
  if (t.isIdentifier(node)) {
    // Find if this identifier is inside a function
    const functionAncestor = ancestors.find(ancestor => 
      t.isFunctionDeclaration(ancestor.node) || 
      t.isFunctionExpression(ancestor.node)
    );
    
    if (functionAncestor) {
      console.log(`Identifier "${node.name}" is inside a function`);
    } else {
      console.log(`Identifier "${node.name}" is at module level`);
    }
  }
});

import { traverseFast } from "@babel/types";

// Count all nodes quickly
let nodeCount = 0;
traverseFast(ast, (node) => {
  nodeCount++;
});
console.log(`Total nodes: ${nodeCount}`);

// Collect all string literals quickly
const strings: string[] = [];
traverseFast(ast, (node) => {
  if (t.isStringLiteral(node)) {
    strings.push(node.value);
  }
});
console.log("String literals:", strings);

// Skip function bodies during traversal
traverseFast(ast, (node, opts) => {
  console.log(`Visiting ${node.type}`);
  
  // Skip traversing inside function declarations
  if (t.isFunctionDeclaration(node)) {
    console.log("Skipping function body");
    return traverseFast.skip;
  }
  
  // Stop traversal entirely if we find an error
  if (t.isThrowStatement(node)) {
    console.log("Found throw statement, stopping traversal");
    return traverseFast.stop;
  }
});

// Find all binary expressions inside functions
traverse(ast, (node, ancestors, state) => {
  if (t.isBinaryExpression(node)) {
    // Check if we're inside a function
    const inFunction = ancestors.some(ancestor => 
      t.isFunctionDeclaration(ancestor.node) || 
      t.isFunctionExpression(ancestor.node) ||
      t.isArrowFunctionExpression(ancestor.node)
    );
    
    if (inFunction) {
      console.log(`Found binary expression: ${node.operator}`);
    }
  }
});

traverse(ast, (node, ancestors, state) => {
  if (t.isVariableDeclarator(node)) {
    // Find which declarator in the declaration this is
    const parent = ancestors[ancestors.length - 1];
    if (parent && parent.index !== undefined) {
      console.log(`Variable declarator at index ${parent.index}`);
    }
  }
});

const startTime = Date.now();

// Fast traversal - better for simple operations
let fastCount = 0;
traverseFast(ast, (node) => {
  fastCount++;
});

const fastTime = Date.now() - startTime;

const start2 = Date.now();

// Regular traversal - better when you need ancestry info
let regularCount = 0;
traverse(ast, (node, ancestors, state) => {
  regularCount++;
});

const regularTime = Date.now() - start2;

console.log(`Fast traversal: ${fastCount} nodes in ${fastTime}ms`);
console.log(`Regular traversal: ${regularCount} nodes in ${regularTime}ms`);