diff --git a/packages/babel-core/src/transformation/index.ts b/packages/babel-core/src/transformation/index.ts index ae780fc2ab02..68f1489c1cfe 100644 --- a/packages/babel-core/src/transformation/index.ts +++ b/packages/babel-core/src/transformation/index.ts @@ -117,7 +117,7 @@ function* transformFile(file: File, pluginPasses: PluginPasses): Handler { passes, file.opts.wrapPluginVisitorMethod, ); - traverse(file.ast, visitor, file.scope); + traverse(file.ast.program, visitor, file.scope, null, file.path, true); for (const [plugin, pass] of passPairs) { const fn = plugin.post; diff --git a/packages/babel-traverse/package.json b/packages/babel-traverse/package.json index 31288112593d..f7e6fe8e0ad4 100644 --- a/packages/babel-traverse/package.json +++ b/packages/babel-traverse/package.json @@ -28,6 +28,7 @@ "globals": "condition:BABEL_8_BREAKING ? ^13.5.0 : ^11.1.0" }, "devDependencies": { + "@babel/core": "workspace:^", "@babel/helper-plugin-test-runner": "workspace:^" }, "engines": { diff --git a/packages/babel-traverse/src/cache.ts b/packages/babel-traverse/src/cache.ts index e17493c93843..33a3ccf3bc83 100644 --- a/packages/babel-traverse/src/cache.ts +++ b/packages/babel-traverse/src/cache.ts @@ -1,8 +1,13 @@ import type { Node } from "@babel/types"; import type NodePath from "./path"; import type Scope from "./scope"; +import type { HubInterface } from "./hub"; -export let path: WeakMap> = new WeakMap(); +let pathsCache: WeakMap< + HubInterface | typeof nullHub, + WeakMap> +> = new WeakMap(); +export { pathsCache as path }; export let scope: WeakMap = new WeakMap(); export function clear() { @@ -11,9 +16,29 @@ export function clear() { } export function clearPath() { - path = new WeakMap(); + pathsCache = new WeakMap(); } export function clearScope() { scope = new WeakMap(); } + +// NodePath#hub can be null, but it's not a valid weakmap key because it +// cannot be collected by GC. Use an object, knowing tht it will not be +// collected anyway. It's not a memory leak because pathsCache.get(nullHub) +// is itself a weakmap, so its entries can still be collected. +const nullHub = Object.freeze({} as const); + +export function getCachedPaths(hub: HubInterface | null, parent: Node) { + return pathsCache.get(hub ?? nullHub)?.get(parent); +} + +export function getOrCreateCachedPaths(hub: HubInterface | null, parent: Node) { + let parents = pathsCache.get(hub ?? nullHub); + if (!parents) pathsCache.set(hub ?? nullHub, (parents = new WeakMap())); + + let paths = parents.get(parent); + if (!paths) parents.set(parent, (paths = new Map())); + + return paths; +} diff --git a/packages/babel-traverse/src/index.ts b/packages/babel-traverse/src/index.ts index 1c6264fb5430..331d0dbf1a83 100644 --- a/packages/babel-traverse/src/index.ts +++ b/packages/babel-traverse/src/index.ts @@ -36,6 +36,7 @@ function traverse( scope: Scope | undefined, state: S, parentPath?: NodePath, + visitSelf?: boolean, ): void; function traverse( @@ -44,6 +45,7 @@ function traverse( scope?: Scope, state?: any, parentPath?: NodePath, + visitSelf?: boolean, ): void; function traverse( @@ -53,6 +55,7 @@ function traverse( scope?: Scope, state?: any, parentPath?: NodePath, + visitSelf?: boolean, ) { if (!parent) return; @@ -66,13 +69,25 @@ function traverse( } } + if (!parentPath && visitSelf) { + throw new Error("visitSelf can only be used when providing a NodePath."); + } + if (!VISITOR_KEYS[parent.type]) { return; } visitors.explode(opts as Visitor); - traverseNode(parent, opts as ExplodedVisitor, scope, state, parentPath); + traverseNode( + parent, + opts as ExplodedVisitor, + scope, + state, + parentPath, + /* skipKeys */ null, + visitSelf, + ); } export default traverse; @@ -100,8 +115,6 @@ traverse.node = function ( traverse.clearNode = function (node: t.Node, opts?: RemovePropertiesOptions) { removeProperties(node, opts); - - cache.path.delete(node); }; traverse.removeProperties = function ( diff --git a/packages/babel-traverse/src/path/index.ts b/packages/babel-traverse/src/path/index.ts index 8f9e3f401ed9..f713ab4d75d0 100644 --- a/packages/babel-traverse/src/path/index.ts +++ b/packages/babel-traverse/src/path/index.ts @@ -8,7 +8,7 @@ import type { Visitor } from "../types"; import Scope from "../scope"; import { validate } from "@babel/types"; import * as t from "@babel/types"; -import { path as pathCache } from "../cache"; +import { getOrCreateCachedPaths } from "../cache"; import generator from "@babel/generator"; // NodePath is split across many files. @@ -92,11 +92,7 @@ class NodePath { // @ts-expect-error key must present in container container[key]; - let paths = pathCache.get(parent); - if (!paths) { - paths = new Map(); - pathCache.set(parent, paths); - } + const paths = getOrCreateCachedPaths(hub, parent); let path = paths.get(targetNode); if (!path) { diff --git a/packages/babel-traverse/src/path/modification.ts b/packages/babel-traverse/src/path/modification.ts index f9e7e3f9257e..015932e86cce 100644 --- a/packages/babel-traverse/src/path/modification.ts +++ b/packages/babel-traverse/src/path/modification.ts @@ -1,6 +1,6 @@ // This file contains methods that modify the path/node in some ways. -import { path as pathCache } from "../cache"; +import { getCachedPaths } from "../cache"; import PathHoister from "./lib/hoister"; import NodePath from "./index"; import { @@ -279,7 +279,7 @@ export function updateSiblingKeys( ) { if (!this.parent) return; - const paths = pathCache.get(this.parent); + const paths = getCachedPaths(this.hub, this.parent) || ([] as never[]); for (const [, path] of paths) { if (typeof path.key === "number" && path.key >= fromIndex) { diff --git a/packages/babel-traverse/src/path/removal.ts b/packages/babel-traverse/src/path/removal.ts index e1a6abef9b04..85d86c756a9f 100644 --- a/packages/babel-traverse/src/path/removal.ts +++ b/packages/babel-traverse/src/path/removal.ts @@ -1,7 +1,7 @@ // This file contains methods responsible for removing a node. import { hooks } from "./lib/removal-hooks"; -import { path as pathCache } from "../cache"; +import { getCachedPaths } from "../cache"; import type NodePath from "./index"; import { REMOVED, SHOULD_SKIP } from "./index"; @@ -46,7 +46,9 @@ export function _remove(this: NodePath) { export function _markRemoved(this: NodePath) { // this.shouldSkip = true; this.removed = true; this._traverseFlags |= SHOULD_SKIP | REMOVED; - if (this.parent) pathCache.get(this.parent).delete(this.node); + if (this.parent) { + getCachedPaths(this.hub, this.parent).delete(this.node); + } this.node = null; } diff --git a/packages/babel-traverse/src/path/replacement.ts b/packages/babel-traverse/src/path/replacement.ts index 71f23f13ddfd..4fc4fa3fe6a7 100644 --- a/packages/babel-traverse/src/path/replacement.ts +++ b/packages/babel-traverse/src/path/replacement.ts @@ -3,7 +3,7 @@ import { codeFrameColumns } from "@babel/code-frame"; import traverse from "../index"; import NodePath from "./index"; -import { path as pathCache } from "../cache"; +import { getCachedPaths } from "../cache"; import { parse } from "@babel/parser"; import { FUNCTION_TYPES, @@ -47,7 +47,7 @@ export function replaceWithMultiple( nodes = this._verifyNodeList(nodes); inheritLeadingComments(nodes[0], this.node); inheritTrailingComments(nodes[nodes.length - 1], this.node); - pathCache.get(this.parent)?.delete(this.node); + getCachedPaths(this.hub, this.parent)?.delete(this.node); this.node = // @ts-expect-error this.key must present in this.container this.container[this.key] = null; @@ -210,7 +210,7 @@ export function _replaceWith(this: NodePath, node: t.Node) { } this.debug(`Replace with ${node?.type}`); - pathCache.get(this.parent)?.set(node, this).delete(this.node); + getCachedPaths(this.hub, this.parent)?.set(node, this).delete(this.node); this.node = // @ts-expect-error this.key must present in this.container diff --git a/packages/babel-traverse/src/traverse-node.ts b/packages/babel-traverse/src/traverse-node.ts index b9d866f254ac..37f571953baa 100644 --- a/packages/babel-traverse/src/traverse-node.ts +++ b/packages/babel-traverse/src/traverse-node.ts @@ -24,13 +24,19 @@ export function traverseNode( state?: any, path?: NodePath, skipKeys?: Record, + visitSelf?: boolean, ): boolean { const keys = VISITOR_KEYS[node.type]; if (!keys) return false; const context = new TraversalContext(scope, opts, state, path); + if (visitSelf) { + if (skipKeys?.[path.parentKey]) return false; + return context.visitQueue([path]); + } + for (const key of keys) { - if (skipKeys && skipKeys[key]) continue; + if (skipKeys?.[key]) continue; if (context.visit(node, key)) { return true; } diff --git a/packages/babel-traverse/test/hub.js b/packages/babel-traverse/test/hub.js index 6be4fe71e7d3..6724472755ea 100644 --- a/packages/babel-traverse/test/hub.js +++ b/packages/babel-traverse/test/hub.js @@ -1,10 +1,46 @@ -import assert from "assert"; +import { transformSync } from "@babel/core"; import { Hub } from "../lib/index.js"; describe("hub", function () { it("default buildError should return TypeError", function () { const hub = new Hub(); const msg = "test_msg"; - assert.deepEqual(hub.buildError(null, msg), new TypeError(msg)); + expect(hub.buildError(null, msg)).toEqual(new TypeError(msg)); + }); + + it("should be preserved across nested traversals", function () { + let origHub; + let innerHub = {}; + let exprHub; + function plugin({ types: t, traverse }) { + return { + visitor: { + Identifier(path) { + if (path.node.name !== "foo") return; + origHub = path.hub; + + const mem = t.memberExpression( + t.identifier("property"), + t.identifier("access"), + ); + traverse(mem, { + noScope: true, + Identifier(path) { + if (path.node.name === "property") innerHub = path.hub; + }, + }); + const [p2] = path.insertAfter(mem); + + exprHub = p2.get("expression").hub; + }, + }, + }; + } + + transformSync("foo;", { configFile: false, plugins: [plugin] }); + + expect(origHub).toBeInstanceOf(Object); + expect(exprHub).toBe(origHub); + expect(innerHub).toBeUndefined(); }); }); diff --git a/yarn.lock b/yarn.lock index 3e0ad25927c5..56e5fa621b6e 100644 --- a/yarn.lock +++ b/yarn.lock @@ -3971,6 +3971,7 @@ __metadata: resolution: "@babel/traverse@workspace:packages/babel-traverse" dependencies: "@babel/code-frame": "workspace:^" + "@babel/core": "workspace:^" "@babel/generator": "workspace:^" "@babel/helper-environment-visitor": "workspace:^" "@babel/helper-function-name": "workspace:^"