State Management in Babel: Patterns and Libraries

State Management in Babel: Patterns and Libraries

When you write a Babel plugin or codemod, you are fundamentally transforming one Abstract Syntax Tree (AST) into another. Along the way, you almost always need to remember things: which nodes you have already visited, which variables are in scope, whether a particular import has been seen, or how many times a helper has been injected. This is the domain of state management in Babel. Despite being a less-discussed topic than state management in UI frameworks, it is just as important — poor state handling is the single most common source of bugs in Babel transforms, from infinite loops to duplicated helper code.

This tutorial walks through what state means in the context of Babel, why it matters, the built-in mechanisms Babel provides, common patterns used by experienced plugin authors, and the libraries that can make your life easier. By the end, you will have a clear mental model for keeping your transforms correct, idempotent, and performant.

What Is State in a Babel Transform?

A Babel plugin is essentially a visitor object whose methods are called as Babel walks the AST. Each visitor method receives a path argument (a wrapper around a node that also tracks its relationship to the rest of the tree) and a state argument. State, in this context, refers to any information that must persist across visitor calls during a single compilation pass — or, in some cases, across multiple files.

There are several distinct kinds of state you may need to track:

• Per-node state: information tied to a specific AST node, such as "has this function already been transformed?"
• Per-file state: information scoped to a single source file, such as the set of imported identifiers or whether a helper has been injected.
• Per-program state: information that spans the entire build, such as a cache of resolved module paths.
• Traversal state: transient information needed while descending into children, such as the nearest enclosing function or class.

Confusing these categories is a frequent source of bugs. For example, storing per-file information in a module-level variable will leak between files when Babel is used in a multi-file build, leading to transforms that work in isolation but fail in production.

Why State Management Matters

Babel transforms are run on real-world codebases that can contain thousands of files and hundreds of thousands of nodes. Naive state handling can cause three classes of problems:

• Correctness bugs: re-transforming an already-transformed node, or failing to track scope, can produce invalid output.
• Performance issues: storing state in plain objects keyed by node references prevents garbage collection of AST nodes after they are no longer needed, leading to memory bloat.
• Re-entrancy bugs: Babel may re-traverse parts of the tree (for example, after inserting new nodes), and shared mutable state can cause subtle ordering issues.

Good state management keeps transforms predictable, idempotent, and cheap. The rest of this tutorial shows you how.

The Built-in state Object

Every visitor method receives a state object as its second argument. This object is an instance of PluginPass and is created fresh for each file. It exposes a small but useful API:

• state.file — the File object, which contains metadata about the source file.
• state.opts — the plugin options object, merged from the user's Babel config.
• state.filename — the absolute path of the file being compiled (may be undefined in some contexts).
• state.cwd — the current working directory.
• state.get(key) / state.set(key, value) — a key/value store scoped to the current file.

You can also attach arbitrary properties directly to state from within a visitor, but using get/set is preferred because it avoids accidental collisions with Babel's own properties.

A First Example: Counting Function Declarations

Here is a minimal plugin that uses the built-in state to count function declarations per file and injects a comment at the top of the file with the total:

module.exports = function countFunctionsPlugin(babel) {
  const { types: t } = babel;

  return {
    name: "count-functions",
    visitor: {
      FunctionDeclaration(path, state) {
        const current = state.get("functionCount") || 0;
        state.set("functionCount", current + 1);
      },
      Program: {
        exit(path, state) {
          const count = state.get("functionCount") || 0;
          if (count > 0) {
            path.unshiftContainer("body", t.addComment(
              t.emptyStatement(),
              "leading",
              ` This file contains ${count} function declaration(s) `
            ));
          }
        }
      }
    }
  };
};

Notice that we read and update the counter through state.get and state.set. Because the state object is recreated for each file, the counter resets automatically — exactly what we want for per-file state.

Using this Inside Visitors

Inside a visitor method, this is bound to the same state object that is passed as the second argument. This means you can write this.set("foo", 1) instead of state.set("foo", 1). Many plugins in the Babel ecosystem use this style for brevity. Just be careful not to use arrow functions for visitor methods, because arrow functions do not bind their own this.

module.exports = function ({ types: t }) {
  return {
    visitor: {
      Identifier(path) {
        // 'this' is the PluginPass instance
        const seen = this.get("seenIdentifiers") || new Set();
        seen.add(path.node.name);
        this.set("seenIdentifiers", seen);
      }
    }
  };
};

Per-Node State with WeakMaps

The built-in state object is great for per-file data, but it is the wrong tool for per-node data. If you need to attach information to specific nodes — for example, marking a function as already transformed so you do not process it twice — you should use a WeakMap declared at the plugin module level.

Why a WeakMap? Because AST nodes are plain objects, and a WeakMap keyed by those objects will automatically drop its entries when the nodes are garbage collected. A plain Map or object would keep the nodes alive forever, leaking memory across a large build.

Example: Idempotent Transform

Suppose we want to rewrite calls to myDebug(...) into console.debug(...), but we want to make sure that if our plugin runs twice (or if the input already contains console.debug calls we generated), we do not transform them again. We can use a WeakSet to mark nodes we have produced:

const generatedNodes = new WeakSet();

module.exports = function ({ types: t }) {
  return {
    name: "rewrite-my-debug",
    visitor: {
      CallExpression(path) {
        const callee = path.node.callee;

        // Skip nodes we generated ourselves.
        if (generatedNodes.has(path.node)) return;

        if (t.isIdentifier(callee, { name: "myDebug" })) {
          const newCall = t.callExpression(
            t.memberExpression(
              t.identifier("console"),
              t.identifier("debug")
            ),
            path.node.arguments
          );
          generatedNodes.add(newCall);
          path.replaceWith(newCall);
        }
      }
    }
  };
};

This pattern is extremely common in production plugins. The WeakSet lives at module scope, so it persists for the lifetime of the process, but because it holds only weak references, it does not prevent garbage collection.

Traversal State: Tracking the Enclosing Context

Sometimes you need to know something about the nearest enclosing construct — for example, whether the current node is inside an async function, or what the current class name is. There are two common approaches.

Approach 1: Walk Up the Tree

The simplest approach is to use path.findParent or path.parentPath to walk up the tree until you find what you need. This is easy to write but can be expensive if done on every node in a large file.

function isInAsyncFunction(path) {
  let p = path.parentPath;
  while (p) {
    if (p.isFunction() && p.node.async) return true;
    p = p.parentPath;
  }
  return false;
}

Approach 2: Maintain a Stack

For better performance, maintain an explicit stack using the enter and exit phases of the visitor. Babel calls enter when it first visits a node and exit after it has finished visiting all of the node's children. This gives you a natural push/pop mechanism.

module.exports = function ({ types: t }) {
  const asyncStack = [];

  return {
    visitor: {
      Function: {
        enter(path) {
          asyncStack.push(path.node.async);
        },
        exit() {
          asyncStack.pop();
        }
      },
      AwaitExpression(path) {
        const currentlyInAsync = asyncStack[asyncStack.length - 1];
        if (!currentlyInAsync) {
          // Report an error or transform accordingly.
          throw path.buildCodeFrameError(
            "await is only allowed inside async functions"
          );
        }
      }
    }
  };
};

This stack-based approach is O(1) per visit and is the recommended pattern when you need enclosing-context information on a hot path.

File-Level State: Imports and Helpers

A very common use case is tracking imports. For example, a plugin that transforms JSX may need to know whether React has been imported and under what name. The idiomatic approach is to record import information in the state object during the ImportDeclaration visitor, and then consult it later.

module.exports = function ({ types: t }) {
  return {
    name: "track-react-import",
    visitor: {
      ImportDeclaration(path, state) {
        if (path.node.source.value === "react") {
          const names = path.node.specifiers.map(s => s.local.name);
          state.set("reactImports", names);
        }
      },
      JSXElement(path, state) {
        const imports = state.get("reactImports") || [];
        if (!imports.includes("React")) {
          throw path.buildCodeFrameError(
            "JSX requires React to be in scope, but it was not imported."
          );
        }
      }
    }
  };
};

For injecting helpers, Babel provides a first-class API: babel.addHelper(name). This automatically deduplicates helpers within a file and across files, so you do not need to track injection state yourself. Prefer it over manually inserting helper functions.

module.exports = function ({ types: t, template, helper }) {
  return {
    visitor: {
      TaggedTemplateExpression(path, state) {
        // addHelper returns an identifier referencing the helper,
        // injecting the helper function exactly once per file.
        const helperId = state.addHelper("typeof");
        path.replaceWith(
          t.callExpression(helperId, [path.node.quasi])
        );
      }
    }
  };
};

Cross-File State

Occasionally you need state that persists across files — for example, a cache of already-resolved module paths to avoid redundant filesystem lookups. The correct place for this is a module-level variable inside your plugin file, because Babel loads each plugin module once per compilation. Be sure to use a data structure that does not grow without bound; a Map keyed by resolved path is usually fine, but consider adding a size cap if your build is very large.

const resolutionCache = new Map();

function resolveCached(filename, resolver) {
  if (resolutionCache.has(filename)) {
    return resolutionCache.get(filename);
  }
  const result = resolver(filename);
  resolutionCache.set(filename, result);
  return result;
}

module.exports = function ({ types: t }) {
  return {
    visitor: {
      ImportDeclaration(path, state) {
        const resolved = resolveCached(
          path.node.source.value,
          (spec) => state.file.resolveModule(spec)
        );
        // ... use resolved ...
      }
    }
  };
};

Avoid putting cross-file state on the state object, because that object is per-file. Avoid putting it on the plugin's return value, because Babel may instantiate the plugin multiple times.

Libraries That Help

While Babel's built-in APIs cover most needs, several libraries and packages make state management easier for specific scenarios.

@babel/traverse

The @babel/traverse package is what Babel uses internally to walk the AST. When you are writing a standalone codemod (not a Babel plugin), you can call traverse(ast, visitor, scope, state) directly and pass a custom state object. This is useful for scripts that need fine-grained control over traversal state.

const { parse } = require("@babel/parser");
const traverse = require("@babel/traverse").default;
const generate = require("@babel/generator").default;

const code = "const x = 1; const y = 2;";
const ast = parse(code);

const state = { constCount: 0 };

traverse(ast, {
  VariableDeclaration(path) {
    if (path.node.kind === "const") {
      this.constCount++;
    }
  }
}, undefined, state);

console.log(state.constCount); // 2

babel-plugin-tester

Although primarily a testing utility, babel-plugin-tester helps you verify that your state handling is correct by letting you run the same plugin multiple times against the same input and assert idempotency. This is invaluable for catching the "runs twice, breaks" class of bugs that plague stateful transforms.

jscodeshift

For large-scale codemods, jscodeshift wraps Babel (or Recast) and provides its own collection-based API. It manages some traversal state for you, such as tracking which nodes have been modified, which can simplify certain transforms. However, the patterns described above still apply when you drop down to the underlying Babel APIs.

ast-types and Scope Helpers

Babel's path.scope object is itself a form of managed state: it tracks bindings within each scope. Methods like scope.getBinding(name), scope.hasBinding(name), and scope.rename(oldName, newName) handle the bookkeeping of variable references for you. Whenever you are tempted to manually track variable names, check whether scope already does what you need.

module.exports = function ({ types: t }) {
  return {
    visitor: {
      Identifier(path) {
        // Only rename if the identifier is a binding we own.
        const binding = path.scope.getBinding(path.node.name);
        if (binding && binding.constant) {
          path.scope.rename(path.node.name, "renamed_" + path.node.name);
        }
      }
    }
  };
};

Best Practices

• Use the right scope for the right state. Per-node data belongs in a WeakMap or WeakSet; per-file data belongs on the state object; cross-file data belongs in a module-level variable.
• Prefer WeakMap/WeakSet over Map/Set for node-keyed data to avoid retaining AST nodes in memory after they are no longer needed.
• Make transforms idempotent. Always check whether a node has already been transformed before transforming it. This protects against double-application and against re-traversal triggered by your own mutations.
• Use enter/exit for stack-based context rather than walking up the tree on every visit, especially on hot paths.
• Prefer state.addHelper over manual helper injection. It handles deduplication and naming collisions for you.
• Never store mutable state on the plugin's returned object. Babel may invoke the plugin function multiple times, and shared state there will leak unpredictably.
• Avoid arrow functions for visitor methods if you rely on this being the state object.
• Test with multiple files. Bugs from leaking per-file state only appear when more than one file is compiled in the same process. Run your test suite with at least two files.
• Use path.buildCodeFrameError instead of throwing plain errors, so users get helpful source-location diagnostics.

Conclusion

State management in Babel is less about choosing a single library and more about matching the right mechanism to the right scope of data. The built-in state object handles per-file concerns, WeakMap and WeakSet handle per-node concerns without leaking memory, explicit enter/exit stacks handle traversal context efficiently, and module-level variables handle cross-file caches. By following the patterns and best practices outlined here — and leaning on Babel's own scope and helper APIs where they exist — you can write transforms that are correct, idempotent, and fast, even on the largest codebases.