From 009142154a0471c442e41e2914dd94d8b6690ac5 Mon Sep 17 00:00:00 2001
From: Arda TANRIKULU <ardatanrikulu@gmail.com>
Date: Mon, 20 Sep 2021 12:00:14 -0400
Subject: [PATCH] Bring `visitorKeys` back (#3250)

Co-authored-by: Ivan Goncharov <ivan.goncharov.ua@gmail.com>
---
 src/index.ts                           |  1 +
 src/language/__tests__/visitor-test.ts | 36 ++++++++++++++++++++++++++
 src/language/index.ts                  |  2 +-
 src/language/visitor.ts                | 30 +++++++++++++--------
 4 files changed, 57 insertions(+), 12 deletions(-)

diff --git a/src/index.ts b/src/index.ts
index be0eda9b0f..836ca22d0d 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -234,6 +234,7 @@ export type {
   /** Visitor utilities */
   ASTVisitor,
   ASTVisitFn,
+  ASTVisitorKeyMap,
   /** AST nodes */
   ASTNode,
   ASTKindToNode,
diff --git a/src/language/__tests__/visitor-test.ts b/src/language/__tests__/visitor-test.ts
index 93b8fe072e..e63f03ed4a 100644
--- a/src/language/__tests__/visitor-test.ts
+++ b/src/language/__tests__/visitor-test.ts
@@ -4,6 +4,7 @@ import { describe, it } from 'mocha';
 import { kitchenSinkQuery } from '../../__testUtils__/kitchenSinkQuery';
 
 import type { ASTNode, SelectionSetNode } from '../ast';
+import type { ASTVisitorKeyMap, ASTVisitor } from '../visitor';
 import { isNode } from '../ast';
 import { Kind } from '../kinds';
 import { parse } from '../parser';
@@ -462,6 +463,41 @@ describe('Visitor', () => {
     ]);
   });
 
+  it('visits only the specified `Kind` in visitorKeyMap', () => {
+    const visited: Array<any> = [];
+
+    const visitorKeyMap: ASTVisitorKeyMap = {
+      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(`
+      query ExampleOperation {
+        someField
+      }
+    `);
+
+    visit(exampleDocumentAST, visitor, visitorKeyMap);
+
+    expect(visited).to.deep.equal([
+      ['enter', 'Document', undefined],
+      ['enter', 'OperationDefinition', undefined],
+      ['enter', 'Name', 'ExampleOperation'],
+      ['leave', 'Name', 'ExampleOperation'],
+      ['leave', 'OperationDefinition', undefined],
+      ['leave', 'Document', undefined],
+    ]);
+  });
+
   it('Legacy: visits variables defined in fragments', () => {
     const ast = parse('fragment a($v: Boolean = false) on t { f }', {
       noLocation: true,
diff --git a/src/language/index.ts b/src/language/index.ts
index dfe4e53584..070ca45d23 100644
--- a/src/language/index.ts
+++ b/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, ASTVisitorKeyMap } from './visitor';
 
 export { Location, Token } from './ast';
 export type {
diff --git a/src/language/visitor.ts b/src/language/visitor.ts
index 458a7af2f3..df3081bedd 100644
--- a/src/language/visitor.ts
+++ b/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 ASTVisitorKeyMap = {
+  [P in keyof ASTKindToNode]?: ReadonlyArray<keyof ASTKindToNode[P]>;
+};
 
+export const QueryDocumentKeys: ASTVisitorKeyMap = {
   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?: 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 {
   /* 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) {