fix: traverse.visitors.merge unable to handle enter,exit

packages/babel-core/src/config/validation/plugins.ts

@@ -38,7 +38,7 @@ const VALIDATORS: ValidatorSet = {
 function assertVisitorMap(loc: OptionPath, value: unknown): Visitor {
   const obj = assertObject(loc, value);
   if (obj) {
-    Object.keys(obj).forEach(prop => assertVisitorHandler(prop, obj[prop]));
+    if (obj._exploded) return;
 
     if (obj.enter || obj.exit) {
       throw new Error(
@@ -47,14 +47,13 @@ function assertVisitorMap(loc: OptionPath, value: unknown): Visitor {
         )} cannot contain catch-all "enter" or "exit" handlers. Please target individual nodes.`,
       );
     }
+
+    Object.keys(obj).forEach(prop => assertVisitorHandler(prop, obj[prop]));
   }
   return obj as Visitor;
 }
 
-function assertVisitorHandler(
-  key: string,
-  value: unknown,
-): VisitorHandler | void {
+function assertVisitorHandler(key: string, value: unknown) {
   if (value && typeof value === "object") {
     Object.keys(value).forEach((handler: string) => {
       if (handler !== "enter" && handler !== "exit") {
@@ -66,17 +65,8 @@ function assertVisitorHandler(
   } else if (typeof value !== "function") {
     throw new Error(`.visitor["${key}"] must be a function`);
   }
-
-  return value as any;
 }
 
-type VisitorHandler =
-  | Function
-  | {
-      enter?: Function;
-      exit?: Function;
-    };
-
 export type PluginObject<S extends PluginPass = PluginPass> = {
   name?: string;
   manipulateOptions?: (

packages/babel-traverse/src/visitors.ts

@@ -5,8 +5,10 @@ import {
   FLIPPED_ALIAS_KEYS,
   TYPES,
   __internal__deprecationWarning as deprecationWarning,
+  type Node,
 } from "@babel/types";
-import type { ExplodedVisitor, NodePath, Visitor } from "./index";
+import type { NodePath, Visitor } from "./index";
+import type { VisitNodeFunction, ExplodedVisitor } from "./types";
 
 type VIRTUAL_TYPES = keyof typeof virtualTypes;
 function isVirtualType(type: string): type is VIRTUAL_TYPES {
@@ -224,71 +226,88 @@ export function merge(
   states: any[] = [],
   wrapper?: Function | null,
 ) {
-  const rootVisitor: Visitor = {};
+  const rootVisitor: ExplodedVisitor = { _exploded: true, _verified: true };
 
   for (let i = 0; i < visitors.length; i++) {
-    const visitor = visitors[i];
+    const visitor = explode(visitors[i]);
     const state = states[i];
 
-    explode(visitor);
+    for (const type of Object.keys(visitor) as (keyof ExplodedVisitor)[]) {
+      const isCatchAll = type === "enter" || type === "exit";
 
-    for (const type of Object.keys(visitor) as (keyof Visitor)[]) {
-      let visitorType = visitor[type];
+      if (type === "_exploded" || type === "_verified") continue;
+      if (isCatchAll) rootVisitor[type] ||= [];
+
+      let newVisitor = {};
 
       // if we have state or wrapper then overload the callbacks to take it
       if (state || wrapper) {
-        visitorType = wrapWithStateOrWrapper(visitorType, state, wrapper);
+        if (isCatchAll) {
+          rootVisitor[type].push(
+            ...wrapWithStateOrWrapper(type, visitor[type], state, wrapper),
+          );
+          continue;
+        }
+
+        const visitorType = visitor[type];
+
+        for (const key of Object.keys(visitorType)) {
+          // @ts-expect-error TODO: improve types
+          const fns = visitorType[key];
+
+          // not an enter/exit array of callbacks
+          if (!Array.isArray(fns)) continue;
+
+          // @ts-expect-error TODO: improve types
+          newVisitor[key] = wrapWithStateOrWrapper(
+            key,
+            fns as VisitNodeFunction<unknown, Node>[],
+            state,
+            wrapper,
+          );
+        }
+      } else if (isCatchAll) {
+        rootVisitor[type].push(...visitor[type]);
+        continue;
+      } else {
+        newVisitor = visitor[type];
       }
 
-      // @ts-expect-error: Expression produces a union type that is too complex to represent.
       const nodeVisitor = (rootVisitor[type] ||= {});
-      mergePair(nodeVisitor, visitorType);
+      mergePair(nodeVisitor, newVisitor);
     }
   }
 
   return rootVisitor;
 }
 
 function wrapWithStateOrWrapper<State>(
-  oldVisitor: Visitor<State>,
+  key: string,
+  fns: VisitNodeFunction<unknown, Node>[],
   state: State,
   wrapper?: Function | null,
 ) {
-  const newVisitor: Visitor = {};
+  return fns.map(fn => {
+    let newFn = fn;
 
-  for (const key of Object.keys(oldVisitor) as (keyof Visitor<State>)[]) {
-    let fns = oldVisitor[key];
-
-    // not an enter/exit array of callbacks
-    if (!Array.isArray(fns)) continue;
-
-    fns = fns.map(function (fn) {
-      let newFn = fn;
-
-      if (state) {
-        newFn = function (path: NodePath) {
-          fn.call(state, path, state);
-        };
-      }
-
-      if (wrapper) {
-        // @ts-expect-error Fixme: document state.key
-        newFn = wrapper(state.key, key, newFn);
-      }
-
-      // Override toString in case this function is printed, we want to print the wrapped function, same as we do in `wrapCheck`
-      if (newFn !== fn) {
-        newFn.toString = () => fn.toString();
-      }
+    if (state) {
+      newFn = function (path: NodePath) {
+        fn.call(state, path, state);
+      };
+    }
 
-      return newFn;
-    });
+    if (wrapper) {
+      // @ts-expect-error Fixme: document state.key
+      newFn = wrapper(state.key, key, newFn);
+    }
 
-    // @ts-expect-error: Expression produces a union type that is too complex to represent.
-    newVisitor[key] = fns;
-  }
+    // Override toString in case this function is printed, we want to print the wrapped function, same as we do in `wrapCheck`
+    if (newFn !== fn) {
+      newFn.toString = () => fn.toString();
+    }
 
-  return newVisitor;
+    return newFn;
+  });
 }
 
 function ensureEntranceObjects(obj: Visitor) {
@@ -329,7 +348,7 @@ function shouldIgnoreKey(
   | "skipKeys"
   | "blacklist" {
   // internal/hidden key
-  if (key[0] === "_") return true;
+  if (key.startsWith("_")) return true;
 
   // ignore function keys
   if (key === "enter" || key === "exit" || key === "shouldSkip") return true;

packages/babel-traverse/test/visitors.js

@@ -0,0 +1,85 @@
+import { parse } from "@babel/parser";
+
+import _traverse, { visitors } from "../lib/index.js";
+const traverse = _traverse.default || _traverse;
+
+describe("visitors", () => {
+  it("should set `_verified` and `_exploded` to `true` if merging catch-all visitors", () => {
+    const visitor = visitors.merge([{ enter() {} }, { enter() {} }]);
+    expect(visitor._verified).toBe(true);
+    expect(visitor._exploded).toBe(true);
+  });
+
+  it("should work when merging node type visitors", () => {
+    const ast = parse("1");
+    const visitor = visitors.merge([
+      { ArrayExpression() {} },
+      { ArrayExpression() {} },
+    ]);
+    traverse(ast, visitor);
+    expect(visitor).toMatchInlineSnapshot(`
+      Object {
+        "ArrayExpression": Object {
+          "enter": Array [
+            [Function],
+            [Function],
+          ],
+        },
+        "_exploded": true,
+        "_verified": true,
+      }
+    `);
+  });
+
+  it("enter", () => {
+    const ast = parse("1");
+    const visitor = visitors.merge([{ enter() {} }, { enter() {} }]);
+    traverse(ast, visitor);
+    expect(visitor).toMatchInlineSnapshot(`
+      Object {
+        "_exploded": true,
+        "_verified": true,
+        "enter": Array [
+          [Function],
+          [Function],
+        ],
+      }
+    `);
+  });
+
+  it("enter with states", () => {
+    const ast = parse("1");
+    const visitor = visitors.merge([{ enter() {} }, { enter() {} }], [{}, {}]);
+    traverse(ast, visitor);
+    expect(visitor).toMatchInlineSnapshot(`
+      Object {
+        "_exploded": true,
+        "_verified": true,
+        "enter": Array [
+          [Function],
+          [Function],
+        ],
+      }
+    `);
+  });
+
+  it("enter with wrapper", () => {
+    const ast = parse("1");
+    const visitor = visitors.merge(
+      [{ enter() {} }, { enter() {} }],
+      [{}, {}],
+      (stateKey, key, fn) => fn,
+    );
+    traverse(ast, visitor);
+    expect(visitor).toMatchInlineSnapshot(`
+      Object {
+        "_exploded": true,
+        "_verified": true,
+        "enter": Array [
+          [Function],
+          [Function],
+        ],
+      }
+    `);
+  });
+});