Bring visitorKeys back

src/index.ts

@@ -234,6 +234,7 @@ export type {
   /** Visitor utilities */
   ASTVisitor,
   ASTVisitFn,
+  VisitorKeyMap,
   /** AST nodes */
   ASTNode,
   ASTKindToNode,

src/language/__tests__/visitor-test.ts

@@ -3,10 +3,11 @@ import { describe, it } from 'mocha';
 
 import { kitchenSinkQuery } from '../../__testUtils__/kitchenSinkQuery';
 
-import type { ASTNode, SelectionSetNode } from '../ast';
+import type { ASTKindToNode, ASTNode, SelectionSetNode } from '../ast';
 import { isNode } from '../ast';
 import { Kind } from '../kinds';
 import { parse } from '../parser';
+import type { VisitorKeyMap, ASTVisitor } from '../visitor';
 import { visit, visitInParallel, BREAK } from '../visitor';
 
 function checkVisitorFnArgs(ast: any, args: any, isEdited: boolean = false) {
@@ -893,6 +894,41 @@ describe('Visitor', () => {
     ]);
   });
 
+  it('visits only the specified `Kind` in visitorKeyMap', () => {
+    const visited: Array<[string, string, ReturnType<typeof getValue>]> = [];
+
+    const visitorKeyMap: VisitorKeyMap<ASTKindToNode> = {
+      Document: ['definitions'],
+      OperationDefinition: ['name'],
+    };
+
+    const visitor: ASTVisitor = {
+      enter(node) {
+        visited.push(['enter', node.kind, getValue(node)]);
+      },
+      leave(node) {
+        visited.push(['leave', node.kind, getValue(node)]);
+      },
+    };
+
+    const exampleDocumentAST = parse(/* GraphQL */ `
+      query ExampleOperation {
+        someField
+      }
+    `);
+
+    visit(exampleDocumentAST, visitor, visitorKeyMap);
+
+    expect(visited).to.deep.equal([
+      ['enter', Kind.DOCUMENT, undefined],
+      ['enter', Kind.OPERATION_DEFINITION, undefined],
+      ['enter', Kind.NAME, 'ExampleOperation'],
+      ['leave', Kind.NAME, 'ExampleOperation'],
+      ['leave', Kind.OPERATION_DEFINITION, undefined],
+      ['leave', Kind.DOCUMENT, undefined],
+    ]);
+  });
+
   describe('visitInParallel', () => {
     // Note: nearly identical to the above test of the same test but
     // using visitInParallel.

src/language/index.ts

@@ -19,7 +19,7 @@ export type { ParseOptions } from './parser';
 export { print } from './printer';
 
 export { visit, visitInParallel, getVisitFn, BREAK } from './visitor';
-export type { ASTVisitor, ASTVisitFn } from './visitor';
+export type { ASTVisitor, ASTVisitFn, VisitorKeyMap } from './visitor';
 
 export { Location, Token } from './ast';
 export type {

src/language/visitor.ts

@@ -76,9 +76,14 @@ type ReducedField<T, R> = T extends null | undefined
   ? ReadonlyArray<R>
   : R;
 
-const QueryDocumentKeys = {
-  Name: [],
+/**
+ * A KeyMap describes each the traversable properties of each kind of node.
+ */
+export type VisitorKeyMap<KindToNode> = {
+  [P in keyof KindToNode]?: ReadonlyArray<keyof KindToNode[P]>;
+};
 
+export const QueryDocumentKeys: VisitorKeyMap<ASTKindToNode> = {
   Document: ['definitions'],
   OperationDefinition: [
     'name',
@@ -103,12 +108,6 @@ const QueryDocumentKeys = {
     'selectionSet',
   ],
 
-  IntValue: [],
-  FloatValue: [],
-  StringValue: [],
-  BooleanValue: [],
-  NullValue: [],
-  EnumValue: [],
   ListValue: ['values'],
   ObjectValue: ['fields'],
   ObjectField: ['name', 'value'],
@@ -242,11 +241,20 @@ export const BREAK: unknown = Object.freeze({});
  * })
  * ```
  */
-export function visit<N extends ASTNode>(root: N, visitor: ASTVisitor): N;
-export function visit<R>(root: ASTNode, visitor: ASTReducer<R>): R;
+export function visit<N extends ASTNode>(
+  root: N,
+  visitor: ASTVisitor,
+  visitorKeys?: VisitorKeyMap<ASTKindToNode>,
+): N;
+export function visit<R>(
+  root: ASTNode,
+  visitor: ASTReducer<R>,
+  visitorKeys?: VisitorKeyMap<ASTKindToNode>,
+): R;
 export function visit(
   root: ASTNode,
   visitor: ASTVisitor | ASTReducer<any>,
+  visitorKeys: VisitorKeyMap<ASTKindToNode> = QueryDocumentKeys,
 ): any {
   /* eslint-disable no-undef-init */
   let stack: any = undefined;
@@ -346,7 +354,7 @@ export function visit(
     } else {
       stack = { inArray, index, keys, edits, prev: stack };
       inArray = Array.isArray(node);
-      keys = inArray ? node : (QueryDocumentKeys as any)[node.kind] ?? [];
+      keys = inArray ? node : (visitorKeys as any)[node.kind] ?? [];
       index = -1;
       edits = [];
       if (parent) {