import { devAssert } from '../jsutils/devAssert.js';
import { inspect } from '../jsutils/inspect.js';
import type { ASTNode } from './ast.js';
import { isNode, QueryDocumentKeys } from './ast.js';
import { Kind } from './kinds.js';
export type ASTVisitor = EnterLeaveVisitor<ASTNode> | KindVisitor;
type KindVisitor = {
readonly [NodeT in ASTNode as NodeT['kind']]?:
| ASTVisitFn<NodeT>
| EnterLeaveVisitor<NodeT>;
};
interface EnterLeaveVisitor<TVisitedNode extends ASTNode> {
readonly enter?: ASTVisitFn<TVisitedNode> | undefined;
readonly leave?: ASTVisitFn<TVisitedNode> | undefined;
}
export type ASTVisitFn<TVisitedNode extends ASTNode> = (
node: TVisitedNode,
key: string | number | undefined,
parent: ASTNode | ReadonlyArray<ASTNode> | undefined,
path: ReadonlyArray<string | number>,
ancestors: ReadonlyArray<ASTNode | ReadonlyArray<ASTNode>>,
) => any;
export type ASTReducer<R> = {
readonly [NodeT in ASTNode as NodeT['kind']]?: {
readonly enter?: ASTVisitFn<NodeT>;
readonly leave: ASTReducerFn<NodeT, R>;
};
};
type ASTReducerFn<TReducedNode extends ASTNode, R> = (
node: { [K in keyof TReducedNode]: ReducedField<TReducedNode[K], R> },
key: string | number | undefined,
parent: ASTNode | ReadonlyArray<ASTNode> | undefined,
path: ReadonlyArray<string | number>,
ancestors: ReadonlyArray<ASTNode | ReadonlyArray<ASTNode>>,
) => R;
type ReducedField<T, R> = T extends null | undefined
? T
: T extends ReadonlyArray<any>
? ReadonlyArray<R>
: R;
export type ASTVisitorKeyMap = {
[NodeT in ASTNode as NodeT['kind']]?: ReadonlyArray<keyof NodeT>;
};
export const BREAK: unknown = Object.freeze({});
export function visit<N extends ASTNode>(
root: N,
visitor: ASTVisitor,
visitorKeys?: ASTVisitorKeyMap,
): N;
export function visit<R>(
root: ASTNode,
visitor: ASTReducer<R>,
visitorKeys?: ASTVisitorKeyMap,
): R;
export function visit(
root: ASTNode,
visitor: ASTVisitor | ASTReducer<any>,
visitorKeys: ASTVisitorKeyMap = QueryDocumentKeys,
): any {
const enterLeaveMap = new Map<Kind, EnterLeaveVisitor<ASTNode>>();
for (const kind of Object.values(Kind)) {
enterLeaveMap.set(kind, getEnterLeaveForKind(visitor, kind));
}
let stack: any = undefined;
let inArray = Array.isArray(root);
let keys: any = [root];
let index = -1;
let edits = [];
let node: any = root;
let key: any = undefined;
let parent: any = undefined;
const path: any = [];
const ancestors = [];
do {
index++;
const isLeaving = index === keys.length;
const isEdited = isLeaving && edits.length !== 0;
if (isLeaving) {
key = ancestors.length === 0 ? undefined : path[path.length - 1];
node = parent;
parent = ancestors.pop();
if (isEdited) {
if (inArray) {
node = node.slice();
let editOffset = 0;
for (const [editKey, editValue] of edits) {
const arrayKey = editKey - editOffset;
if (editValue === null) {
node.splice(arrayKey, 1);
editOffset++;
} else {
node[arrayKey] = editValue;
}
}
} else {
node = Object.defineProperties(
{},
Object.getOwnPropertyDescriptors(node),
);
for (const [editKey, editValue] of edits) {
node[editKey] = editValue;
}
}
}
index = stack.index;
keys = stack.keys;
edits = stack.edits;
inArray = stack.inArray;
stack = stack.prev;
} else if (parent) {
key = inArray ? index : keys[index];
node = parent[key];
if (node === null || node === undefined) {
continue;
}
path.push(key);
}
let result;
if (!Array.isArray(node)) {
devAssert(isNode(node), `Invalid AST Node: ${inspect(node)}.`);
const visitFn = isLeaving
? enterLeaveMap.get(node.kind)?.leave
: enterLeaveMap.get(node.kind)?.enter;
result = visitFn?.call(visitor, node, key, parent, path, ancestors);
if (result === BREAK) {
break;
}
if (result === false) {
if (!isLeaving) {
path.pop();
continue;
}
} else if (result !== undefined) {
edits.push([key, result]);
if (!isLeaving) {
if (isNode(result)) {
node = result;
} else {
path.pop();
continue;
}
}
}
}
if (result === undefined && isEdited) {
edits.push([key, node]);
}
if (isLeaving) {
path.pop();
} else {
stack = { inArray, index, keys, edits, prev: stack };
inArray = Array.isArray(node);
keys = inArray ? node : (visitorKeys as any)[node.kind] ?? [];
index = -1;
edits = [];
if (parent) {
ancestors.push(parent);
}
parent = node;
}
} while (stack !== undefined);
if (edits.length !== 0) {
return edits[edits.length - 1][1];
}
return root;
}
export function visitInParallel(
visitors: ReadonlyArray<ASTVisitor>,
): ASTVisitor {
const skipping = new Array(visitors.length).fill(null);
const mergedVisitor = Object.create(null);
for (const kind of Object.values(Kind)) {
let hasVisitor = false;
const enterList = new Array(visitors.length).fill(undefined);
const leaveList = new Array(visitors.length).fill(undefined);
for (let i = 0; i < visitors.length; ++i) {
const { enter, leave } = getEnterLeaveForKind(visitors[i], kind);
hasVisitor ||= enter != null || leave != null;
enterList[i] = enter;
leaveList[i] = leave;
}
if (!hasVisitor) {
continue;
}
const mergedEnterLeave: EnterLeaveVisitor<ASTNode> = {
enter(...args) {
const node = args[0];
for (let i = 0; i < visitors.length; i++) {
if (skipping[i] === null) {
const result = enterList[i]?.apply(visitors[i], args);
if (result === false) {
skipping[i] = node;
} else if (result === BREAK) {
skipping[i] = BREAK;
} else if (result !== undefined) {
return result;
}
}
}
},
leave(...args) {
const node = args[0];
for (let i = 0; i < visitors.length; i++) {
if (skipping[i] === null) {
const result = leaveList[i]?.apply(visitors[i], args);
if (result === BREAK) {
skipping[i] = BREAK;
} else if (result !== undefined && result !== false) {
return result;
}
} else if (skipping[i] === node) {
skipping[i] = null;
}
}
},
};
mergedVisitor[kind] = mergedEnterLeave;
}
return mergedVisitor;
}
export function getEnterLeaveForKind(
visitor: ASTVisitor,
kind: Kind,
): EnterLeaveVisitor<ASTNode> {
const kindVisitor:
| ASTVisitFn<ASTNode>
| EnterLeaveVisitor<ASTNode>
| undefined = (visitor as any)[kind];
if (typeof kindVisitor === 'object') {
return kindVisitor;
} else if (typeof kindVisitor === 'function') {
return { enter: kindVisitor, leave: undefined };
}
return { enter: (visitor as any).enter, leave: (visitor as any).leave };
}