From 614ffab278c1417240be74ec6fbb0a8b48828149 Mon Sep 17 00:00:00 2001 From: Dan Lynch Date: Thu, 30 Jul 2026 20:33:55 +0000 Subject: [PATCH] =?UTF-8?q?feat(traverse):=20mutation-capable=20traversal?= =?UTF-8?q?=20=E2=80=94=20MutablePath=20with=20replaceWith/remove/insert/s?= =?UTF-8?q?kip/stop=20and=20enter/exit=20visitors?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- packages/traverse/__tests__/mutate.test.ts | 209 +++++++++++++ packages/traverse/src/index.ts | 2 + packages/traverse/src/mutate.ts | 338 +++++++++++++++++++++ 3 files changed, 549 insertions(+) create mode 100644 packages/traverse/__tests__/mutate.test.ts create mode 100644 packages/traverse/src/mutate.ts diff --git a/packages/traverse/__tests__/mutate.test.ts b/packages/traverse/__tests__/mutate.test.ts new file mode 100644 index 000000000..fbc6edcf8 --- /dev/null +++ b/packages/traverse/__tests__/mutate.test.ts @@ -0,0 +1,209 @@ +import { MutablePath, traverse } from '../src'; + +const select = (relname: string) => ({ + SelectStmt: { + targetList: [ + { ResTarget: { val: { ColumnRef: { fields: [{ A_Star: {} }] } } } } + ], + fromClause: [ + { RangeVar: { schemaname: 'app', relname, inh: true, relpersistence: 'p' } } + ], + limitOption: 'LIMIT_OPTION_DEFAULT', + op: 'SETOP_NONE' + } +}); + +const funcCallAst = () => ({ + FuncCall: { + funcname: [ + { String: { sval: 'auth' } }, + { String: { sval: 'uid' } } + ], + funcformat: 'COERCE_EXPLICIT_CALL' + } +}); + +describe('traverse: read semantics', () => { + it('visits nodes with parent paths and key paths like walk', () => { + const seen: Array<[string, (string | number)[]]> = []; + traverse(select('posts'), { + enter(path: MutablePath) { + seen.push([path.tag, path.path]); + } + }); + const tags = seen.map(([t]) => t); + expect(tags).toContain('SelectStmt'); + expect(tags).toContain('RangeVar'); + expect(tags).toContain('ResTarget'); + const rangeVar = seen.find(([t]) => t === 'RangeVar')!; + expect(rangeVar[1]).toEqual(['fromClause', 0]); + }); + + it('return false skips children', () => { + const tags: string[] = []; + traverse(select('posts'), { + enter(path: MutablePath) { + tags.push(path.tag); + if (path.tag === 'ResTarget') return false; + } + }); + expect(tags).toContain('ResTarget'); + expect(tags).not.toContain('ColumnRef'); + }); + + it('skip() skips children; stop() ends the traversal', () => { + const tags: string[] = []; + traverse(select('posts'), { + enter(path: MutablePath) { + tags.push(path.tag); + if (path.tag === 'ResTarget') path.skip(); + } + }); + expect(tags).not.toContain('ColumnRef'); + + const seen: string[] = []; + traverse(select('posts'), { + enter(path: MutablePath) { + seen.push(path.tag); + if (path.tag === 'RangeVar') path.stop(); + } + }); + expect(seen).toContain('RangeVar'); + expect(seen).not.toContain('ResTarget'); + }); + + it('runs exit handlers post-order', () => { + const order: string[] = []; + traverse(select('posts'), { + SelectStmt: { + enter() { + order.push('enter:SelectStmt'); + }, + exit() { + order.push('exit:SelectStmt'); + } + }, + RangeVar(path: MutablePath) { + order.push(`enter:${path.tag}`); + } + }); + expect(order[0]).toBe('enter:SelectStmt'); + expect(order[order.length - 1]).toBe('exit:SelectStmt'); + expect(order).toContain('enter:RangeVar'); + }); +}); + +describe('traverse: mutation', () => { + it('replaceWith swaps a tagged node in an array container', () => { + const ast = select('posts'); + traverse(ast, { + RangeVar(path: MutablePath) { + path.replaceWith({ + RangeVar: { relname: 'users', inh: true, relpersistence: 'p' } + }); + } + }); + expect(ast.SelectStmt.fromClause[0]).toEqual({ + RangeVar: { relname: 'users', inh: true, relpersistence: 'p' } + }); + }); + + it('replaceWith traverses the replacement children but does not re-invoke on the replacement', () => { + const ast = select('posts'); + let rangeVarVisits = 0; + const innerTags: string[] = []; + traverse(ast, { + ResTarget(path: MutablePath) { + path.replaceWith({ + ResTarget: { val: { FuncCall: funcCallAst().FuncCall } } + }); + }, + FuncCall(path: MutablePath) { + innerTags.push(path.tag); + }, + RangeVar() { + rangeVarVisits++; + } + }); + expect(innerTags).toEqual(['FuncCall']); + expect(rangeVarVisits).toBe(1); + }); + + it('remove splices from an array container and keeps sibling iteration aligned', () => { + const ast = { + List: { + items: [ + { String: { sval: 'a' } }, + { String: { sval: 'b' } }, + { String: { sval: 'c' } } + ] + } + }; + const visited: string[] = []; + traverse(ast, { + String(path: MutablePath) { + visited.push(path.node.sval); + if (path.node.sval === 'b') path.remove(); + } + }); + expect(visited).toEqual(['a', 'b', 'c']); + expect(ast.List.items.map((i: any) => i.String.sval)).toEqual(['a', 'c']); + }); + + it('remove deletes an object field container entry', () => { + const ast = select('posts'); + traverse(ast, { + RangeVar(path: MutablePath) { + path.remove(); + } + }); + expect(ast.SelectStmt.fromClause).toEqual([]); + }); + + it('insertBefore and insertAfter add unvisited siblings', () => { + const ast = { + List: { + items: [{ String: { sval: 'mid' } }] + } + }; + const visited: string[] = []; + traverse(ast, { + String(path: MutablePath) { + visited.push(path.node.sval); + if (path.node.sval === 'mid') { + path.insertBefore({ String: { sval: 'pre' } }); + path.insertAfter({ String: { sval: 'post' } }); + } + } + }); + expect(visited).toEqual(['mid']); + expect(ast.List.items.map((i: any) => i.String.sval)).toEqual(['pre', 'mid', 'post']); + }); + + it('mutates concrete typed fields stored as bare objects (CreatePolicyStmt.table)', () => { + const ast = { + CreatePolicyStmt: { + policy_name: 'p', + table: { schemaname: 'app', relname: 'posts', inh: true, relpersistence: 'p' }, + cmd_name: 'select', + permissive: true + } + }; + traverse(ast, { + RangeVar(path: MutablePath) { + path.replaceWith({ schemaname: 'tenant', relname: 'posts', inh: true, relpersistence: 'p' }); + } + }); + expect(ast.CreatePolicyStmt.table.schemaname).toBe('tenant'); + }); + + it('throws when mutating a detached root', () => { + expect(() => + traverse(select('posts'), { + SelectStmt(path: MutablePath) { + path.remove(); + } + }) + ).toThrow(/detached/); + }); +}); diff --git a/packages/traverse/src/index.ts b/packages/traverse/src/index.ts index f2ef5db10..f9c4d4fca 100644 --- a/packages/traverse/src/index.ts +++ b/packages/traverse/src/index.ts @@ -1,2 +1,4 @@ +export type { EnterExit, MutableVisitor, MutableWalker } from './mutate'; +export { MutablePath, traverse } from './mutate'; export type { NodeTag,Visitor, VisitorContext, Walker } from './traverse'; export { NodePath,visit, walk } from './traverse'; diff --git a/packages/traverse/src/mutate.ts b/packages/traverse/src/mutate.ts new file mode 100644 index 000000000..aedd45e41 --- /dev/null +++ b/packages/traverse/src/mutate.ts @@ -0,0 +1,338 @@ +/** + * Mutation-capable traversal for PostgreSQL ASTs. + * + * `traverse(root, visitor)` walks the same tree shapes as `walk` (tagged + * nodes, concrete typed fields, bare ParseResult/ScanResult roots) but hands + * visitors a {@link MutablePath} that supports Babel-style operations: + * + * - `path.replaceWith(value)` — replace this node's stored value in its + * container. The replacement's children are traversed; the visitor is not + * re-invoked on the replacement itself (no self-requeue, so a visitor that + * replaces a node with the same tag cannot loop). + * - `path.remove()` — remove this node (splice from an array container or + * delete the field). Children are not traversed. + * - `path.insertBefore(...values)` / `path.insertAfter(...values)` — insert + * siblings in an array container. Inserted values are not visited. + * - `path.skip()` — do not traverse this node's children. + * - `path.stop()` — end the entire traversal. + * + * Visitors may be plain functions (enter-only, `return false` skips children + * like `walk`) or `{ enter?, exit? }` pairs; the tag-keyed map may also carry + * catch-all `enter` / `exit` handlers invoked for every node. + * + * Traversal order is pre-order (enter), children, post-order (exit) — + * deterministic and schema-driven, matching `walk`. + */ +import type { NodeSpec } from './18/runtime-schema'; +import { runtimeSchema } from './18/runtime-schema'; + +const schemaMap = new Map(runtimeSchema.map((spec: NodeSpec) => [spec.name, spec])); + +function isTaggedNode(value: any): boolean { + if (typeof value !== 'object' || value === null || Array.isArray(value)) { + return false; + } + const keys = Object.keys(value); + return keys.length === 1 && /^[A-Z]/.test(keys[0]); +} + +function detectUntaggedRootTag(root: any): string | null { + if (Array.isArray(root.stmts) && typeof root.version === 'number') { + return 'ParseResult'; + } + if (Array.isArray(root.tokens) && typeof root.version === 'number') { + return 'ScanResult'; + } + return null; +} + +class TraversalState { + stopped = false; +} + +/** + * A path handed to `traverse` visitors: the node, its tag, its parent path, + * the root-relative key path, and — when the node lives inside a container — + * mutation operations on that container slot. + */ +export class MutablePath { + /** Number of array entries the visitor inserted before this node. */ + _insertedBefore = 0; + /** Number of array entries the visitor inserted after this node. */ + _insertedAfter = 0; + _removed = false; + _skipped = false; + _replacedWith: any = undefined; + _didReplace = false; + + constructor( + public tag: TTag, + public node: any, + public parent: MutablePath | null, + public keyPath: readonly (string | number)[], + /** The object or array that physically holds this node's stored value. */ + public container: any | null, + /** The key of this node's stored value within `container`. */ + public containerKey: string | number | null, + private state: TraversalState + ) {} + + get path(): (string | number)[] { + return [...this.keyPath]; + } + + get key(): string | number { + return this.keyPath[this.keyPath.length - 1] ?? ''; + } + + get removed(): boolean { + return this._removed; + } + + /** Do not traverse this node's children. */ + skip(): void { + this._skipped = true; + } + + /** End the entire traversal. */ + stop(): void { + this.state.stopped = true; + } + + private assertAttached(op: string): void { + if (this.container === null || this.containerKey === null) { + throw new Error(`Cannot ${op} a detached path (the traversal root has no container)`); + } + } + + /** + * Replace this node's stored value. Pass the value exactly as it should be + * stored: a tagged wrapper (`{ SelectStmt: {...} }`) where the tree stores + * tagged nodes, or bare node data for concrete typed fields. + */ + replaceWith(value: any): void { + this.assertAttached('replaceWith'); + if (this._removed) { + throw new Error('Cannot replaceWith after remove'); + } + this.container[this.containerKey as any] = value; + this._replacedWith = value; + this._didReplace = true; + } + + /** Remove this node from its container. Children are not traversed. */ + remove(): void { + this.assertAttached('remove'); + if (Array.isArray(this.container)) { + const idx = this.containerKey as number; + this.container.splice(idx, 1); + } else { + delete this.container[this.containerKey as any]; + } + this._removed = true; + } + + /** Insert siblings before this node (array containers only, not visited). */ + insertBefore(...values: any[]): void { + this.assertAttached('insertBefore'); + if (!Array.isArray(this.container)) { + throw new Error('insertBefore requires an array container'); + } + const idx = this.containerKey as number; + this.container.splice(idx, 0, ...values); + this.containerKey = idx + values.length; + this._insertedBefore += values.length; + } + + /** Insert siblings after this node (array containers only, not visited). */ + insertAfter(...values: any[]): void { + this.assertAttached('insertAfter'); + if (!Array.isArray(this.container)) { + throw new Error('insertAfter requires an array container'); + } + const idx = (this.containerKey as number) + 1 + this._insertedAfter; + this.container.splice(idx, 0, ...values); + this._insertedAfter += values.length; + } +} + +export type MutableWalker = (path: MutablePath) => boolean | void; + +export type EnterExit = { + enter?: MutableWalker; + exit?: (path: MutablePath) => void; +}; + +export type MutableVisitor = { + [tag: string]: MutableWalker | EnterExit; +}; + +function handlersFor(visitor: MutableVisitor, tag: string): { enter?: MutableWalker; exit?: (p: MutablePath) => void }[] { + const out: { enter?: MutableWalker; exit?: (p: MutablePath) => void }[] = []; + for (const key of [tag, 'enter', 'exit'] as const) { + const h = visitor[key]; + if (!h) continue; + if (key === 'enter' && typeof h === 'function') { + out.push({ enter: h as MutableWalker }); + } else if (key === 'exit' && typeof h === 'function') { + out.push({ exit: h as (p: MutablePath) => void }); + } else if (key === tag) { + if (typeof h === 'function') out.push({ enter: h as MutableWalker }); + else out.push(h as EnterExit); + } + } + return out; +} + +/** + * Traverse `root` with a mutation-capable visitor. See module docs for the + * mutation semantics. + */ +export function traverse(root: any, visitor: MutableVisitor): void { + const state = new TraversalState(); + visitValue(root, null, [], null, null, visitor, state); +} + +/** + * Visit one stored value (which may be a tagged node, a bare typed node when + * `declaredType` names it, an array, or a plain object to descend through). + * Returns the net change in the parent array's length caused by the visit + * (insertions minus removal), so array iteration can stay aligned. + */ +function visitValue( + value: any, + parent: MutablePath | null, + keyPath: readonly (string | number)[], + container: any | null, + containerKey: string | number | null, + visitor: MutableVisitor, + state: TraversalState, + declaredType?: string +): number { + if (state.stopped || typeof value !== 'object' || value === null) { + return 0; + } + + if (Array.isArray(value)) { + for (let i = 0; i < value.length; i++) { + if (state.stopped) break; + const delta = visitValue(value[i], parent, [...keyPath, i], value, i, visitor, state); + i += delta; + } + return 0; + } + + const keys = Object.keys(value); + if (keys.length === 1 && /^[A-Z]/.test(keys[0])) { + return visitNode(keys[0], value[keys[0]], parent, keyPath, container, containerKey, visitor, state); + } + if (declaredType && declaredType !== 'Node' && schemaMap.has(declaredType) && !isTaggedNode(value)) { + // Concrete typed field stored as a bare untagged object. + return visitNode(declaredType, value, parent, keyPath, container, containerKey, visitor, state); + } + if (parent === null && keyPath.length === 0) { + const rootTag = detectUntaggedRootTag(value); + if (rootTag) { + return visitNode(rootTag, value, parent, keyPath, container, containerKey, visitor, state); + } + } + for (const key of keys) { + if (state.stopped) break; + visitValue(value[key], parent, [...keyPath, key], value, key, visitor, state); + } + return 0; +} + +function visitNode( + tag: string, + nodeData: any, + parent: MutablePath | null, + keyPath: readonly (string | number)[], + container: any | null, + containerKey: string | number | null, + visitor: MutableVisitor, + state: TraversalState +): number { + if (typeof nodeData !== 'object' || nodeData === null) { + return 0; + } + + const path = new MutablePath(tag, nodeData, parent, keyPath, container, containerKey, state); + const handlers = handlersFor(visitor, tag); + + let skipChildren = false; + for (const h of handlers) { + if (!h.enter) continue; + if (h.enter(path) === false) skipChildren = true; + if (path._removed || state.stopped) break; + } + + const arrayDelta = path._insertedBefore + path._insertedAfter - (path._removed ? 1 : 0); + + if (path._removed || state.stopped) { + return Array.isArray(container) ? arrayDelta : 0; + } + + if (path._didReplace) { + // Traverse the replacement's children without re-invoking the visitor on + // the replacement itself (no self-requeue). + const replacement = path._replacedWith; + if (isTaggedNode(replacement)) { + const rTag = Object.keys(replacement)[0]; + traverseChildren(rTag, replacement[rTag], path, keyPath, visitor, state); + } else if (typeof replacement === 'object' && replacement !== null) { + traverseChildren(tag, replacement, path, keyPath, visitor, state); + } + return Array.isArray(container) ? arrayDelta : 0; + } + + if (!path._skipped && !skipChildren) { + traverseChildren(tag, nodeData, path, keyPath, visitor, state); + } + + if (!state.stopped) { + for (const h of handlers) { + if (h.exit) h.exit(path); + if (state.stopped) break; + } + } + + return Array.isArray(container) ? arrayDelta : 0; +} + +function traverseChildren( + tag: string, + nodeData: any, + path: MutablePath, + keyPath: readonly (string | number)[], + visitor: MutableVisitor, + state: TraversalState +): void { + const nodeSpec = schemaMap.get(tag); + if (nodeSpec) { + for (const field of nodeSpec.fields) { + if (state.stopped) break; + const isNodeType = field.type === 'Node' || schemaMap.has(field.type); + if (!isNodeType || nodeData[field.name] == null) continue; + const value = nodeData[field.name]; + if (field.isArray && Array.isArray(value)) { + for (let i = 0; i < value.length; i++) { + if (state.stopped) break; + const delta = visitValue( + value[i], path, [...keyPath, field.name, i], value, i, visitor, state, field.type + ); + i += delta; + } + } else if (!field.isArray) { + visitValue( + value, path, [...keyPath, field.name], nodeData, field.name, visitor, state, field.type + ); + } + } + } else { + for (const key of Object.keys(nodeData)) { + if (state.stopped) break; + visitValue(nodeData[key], path, [...keyPath, key], nodeData, key, visitor, state); + } + } +}