Abstract Syntax Trees (AST) in JavaScript Guide

// Source code:
const greeting = "hello";
 
// AST representation (ESTree format):
const ast = {
  type: "Program",
  body: [
    {
      type: "VariableDeclaration",
      kind: "const",
      declarations: [
        {
          type: "VariableDeclarator",
          id: {
            type: "Identifier",
            name: "greeting",
          },
          init: {
            type: "Literal",
            value: "hello",
            raw: '"hello"',
          },
        },
      ],
    },
  ],
  sourceType: "module",
};
 
// More complex: function declaration
// function add(a, b) { return a + b; }
const functionAST = {
  type: "FunctionDeclaration",
  id: { type: "Identifier", name: "add" },
  params: [
    { type: "Identifier", name: "a" },
    { type: "Identifier", name: "b" },
  ],
  body: {
    type: "BlockStatement",
    body: [
      {
        type: "ReturnStatement",
        argument: {
          type: "BinaryExpression",
          operator: "+",
          left: { type: "Identifier", name: "a" },
          right: { type: "Identifier", name: "b" },
        },
      },
    ],
  },
};
 
// Arrow function: (x) => x * 2
const arrowAST = {
  type: "ArrowFunctionExpression",
  params: [{ type: "Identifier", name: "x" }],
  body: {
    type: "BinaryExpression",
    operator: "*",
    left: { type: "Identifier", name: "x" },
    right: { type: "Literal", value: 2 },
  },
  expression: true, // concise body (no braces)
};

// Acorn is a small, fast JavaScript parser used by many tools
import * as acorn from "acorn";
 
// Parse source code into AST
const source = `
function fibonacci(n) {
  if (n <= 1) return n;
  return fibonacci(n - 1) + fibonacci(n - 2);
}
`;
 
const ast = acorn.parse(source, {
  ecmaVersion: 2025,
  sourceType: "module",
  locations: true, // Include line/column info
});
 
// The AST preserves all structural information
console.log(ast.type); // "Program"
console.log(ast.body[0].type); // "FunctionDeclaration"
console.log(ast.body[0].id.name); // "fibonacci"
console.log(ast.body[0].params[0].name); // "n"
 
// Accessing the if statement
const ifStmt = ast.body[0].body.body[0];
console.log(ifStmt.type); // "IfStatement"
console.log(ifStmt.test.type); // "BinaryExpression"
console.log(ifStmt.test.operator); // "<="
 
// Location information
console.log(ifStmt.loc.start); // { line: 3, column: 2 }
console.log(ifStmt.loc.end);   // { line: 3, column: 24 }
 
// Parse with comments
const astWithComments = acorn.parse(source, {
  ecmaVersion: 2025,
  onComment: (isBlock, text, start, end) => {
    console.log(`${isBlock ? "Block" : "Line"} comment: ${text}`);
  },
});
 
// Handling parse errors
try {
  acorn.parse("function { broken", { ecmaVersion: 2025 });
} catch (err) {
  console.log(err.message); // "Unexpected token (1:9)"
  console.log(err.pos);     // 9
  console.log(err.loc);     // { line: 1, column: 9 }
}

// Walking an AST to visit every node
 
// Simple recursive traversal
function traverse(node, visitor) {
  if (!node || typeof node !== "object") return;
 
  // Call visitor for this node type
  const handler = visitor[node.type];
  if (handler) handler(node);
 
  // Visit all child nodes
  for (const key of Object.keys(node)) {
    const child = node[key];
    if (Array.isArray(child)) {
      child.forEach((item) => traverse(item, visitor));
    } else if (child && typeof child === "object" && child.type) {
      traverse(child, visitor);
    }
  }
}
 
// Usage: Find all function names in a program
const functionNames = [];
traverse(ast, {
  FunctionDeclaration(node) {
    functionNames.push(node.id.name);
  },
  ArrowFunctionExpression(node) {
    functionNames.push("(arrow)");
  },
});
 
// Advanced traversal with enter/leave phases
function walk(node, visitors) {
  if (!node || typeof node !== "object") return;
 
  const visitor = visitors[node.type];
  if (visitor?.enter) visitor.enter(node);
 
  for (const key of Object.keys(node)) {
    const child = node[key];
    if (Array.isArray(child)) {
      child.forEach((item) => walk(item, visitors));
    } else if (child && typeof child === "object" && child.type) {
      walk(child, visitors);
    }
  }
 
  if (visitor?.leave) visitor.leave(node);
}
 
// Track scope depth
let scopeDepth = 0;
walk(ast, {
  FunctionDeclaration: {
    enter(node) {
      scopeDepth++;
      console.log(`${" ".repeat(scopeDepth * 2)}Enter function: ${node.id.name}`);
    },
    leave(node) {
      console.log(`${" ".repeat(scopeDepth * 2)}Leave function: ${node.id.name}`);
      scopeDepth--;
    },
  },
  BlockStatement: {
    enter() { scopeDepth++; },
    leave() { scopeDepth--; },
  },
});

// Transforming AST nodes to modify or generate code
 
// Example: Convert var to const/let
function transformVarDeclarations(ast) {
  traverse(ast, {
    VariableDeclaration(node) {
      if (node.kind === "var") {
        // Check if any declarator is reassigned
        const names = node.declarations.map((d) => d.id.name);
        // Simple heuristic: use 'let' (a full implementation would
        // check all references in scope for reassignment)
        node.kind = "let";
      }
    },
  });
  return ast;
}
 
// Example: Add console.log to every function entry
function instrumentFunctions(ast) {
  traverse(ast, {
    FunctionDeclaration(node) {
      const logStatement = {
        type: "ExpressionStatement",
        expression: {
          type: "CallExpression",
          callee: {
            type: "MemberExpression",
            object: { type: "Identifier", name: "console" },
            property: { type: "Identifier", name: "log" },
          },
          arguments: [
            {
              type: "Literal",
              value: `Entering ${node.id.name}`,
            },
          ],
        },
      };
 
      // Insert at the beginning of the function body
      node.body.body.unshift(logStatement);
    },
  });
  return ast;
}
 
// Example: Dead code elimination
function removeDeadCode(ast) {
  traverse(ast, {
    IfStatement(node) {
      // Remove if(false) blocks
      if (node.test.type === "Literal" && node.test.value === false) {
        // Replace with the else block or empty statement
        if (node.alternate) {
          Object.assign(node, node.alternate);
        } else {
          node.type = "EmptyStatement";
          delete node.test;
          delete node.consequent;
          delete node.alternate;
        }
      }
    },
  });
  return ast;
}
 
// Example: Rename variables (simple scope-unaware version)
function renameIdentifier(ast, oldName, newName) {
  traverse(ast, {
    Identifier(node) {
      if (node.name === oldName) {
        node.name = newName;
      }
    },
  });
  return ast;
}

// Generating source code from an AST
 
class CodeGenerator {
  #indent = 0;
 
  generate(node) {
    const method = this[`gen${node.type}`];
    if (!method) throw new Error(`Unknown node type: ${node.type}`);
    return method.call(this, node);
  }
 
  genProgram(node) {
    return node.body.map((stmt) => this.generate(stmt)).join("\n");
  }
 
  genVariableDeclaration(node) {
    const decls = node.declarations
      .map((d) => this.generate(d))
      .join(", ");
    return `${node.kind} ${decls};`;
  }
 
  genVariableDeclarator(node) {
    const id = this.generate(node.id);
    if (node.init) {
      return `${id} = ${this.generate(node.init)}`;
    }
    return id;
  }
 
  genFunctionDeclaration(node) {
    const params = node.params.map((p) => this.generate(p)).join(", ");
    const body = this.generate(node.body);
    return `function ${node.id.name}(${params}) ${body}`;
  }
 
  genBlockStatement(node) {
    this.#indent++;
    const pad = "  ".repeat(this.#indent);
    const body = node.body
      .map((stmt) => `${pad}${this.generate(stmt)}`)
      .join("\n");
    this.#indent--;
    const outerPad = "  ".repeat(this.#indent);
    return `{\n${body}\n${outerPad}}`;
  }
 
  genReturnStatement(node) {
    if (node.argument) {
      return `return ${this.generate(node.argument)};`;
    }
    return "return;";
  }
 
  genBinaryExpression(node) {
    const left = this.generate(node.left);
    const right = this.generate(node.right);
    return `${left} ${node.operator} ${right}`;
  }
 
  genIdentifier(node) { return node.name; }
 
  genLiteral(node) {
    if (typeof node.value === "string") return `"${node.value}"`;
    return String(node.value);
  }
 
  genIfStatement(node) {
    let code = `if (${this.generate(node.test)}) ${this.generate(node.consequent)}`;
    if (node.alternate) {
      code += ` else ${this.generate(node.alternate)}`;
    }
    return code;
  }
 
  genExpressionStatement(node) {
    return `${this.generate(node.expression)};`;
  }
 
  genCallExpression(node) {
    const callee = this.generate(node.callee);
    const args = node.arguments.map((a) => this.generate(a)).join(", ");
    return `${callee}(${args})`;
  }
 
  genMemberExpression(node) {
    return `${this.generate(node.object)}.${this.generate(node.property)}`;
  }
}
 
// Usage
const generator = new CodeGenerator();
const output = generator.generate(ast);
console.log(output);
// function fibonacci(n) {
//   if (n <= 1) return n;
//   return fibonacci(n - 1) + fibonacci(n - 2);
// }