visitor: speed up visitInParallel by dropping support for unknown AST nodes

Author: IvanGoncharov

benchmark/visit-benchmark.js

@@ -0,0 +1,25 @@
+'use strict';
+
+const { parse } = require('graphql/language/parser.js');
+const { visit } = require('graphql/language/visitor.js');
+
+const { bigSchemaSDL } = require('./fixtures.js');
+
+const documentAST = parse(bigSchemaSDL);
+
+const visitor = {
+  enter() {
+    /* do nothing */
+  },
+  leave() {
+    /* do nothing */
+  },
+};
+
+module.exports = {
+  name: 'Visit all AST nodes',
+  count: 10,
+  measure() {
+    visit(documentAST, visitor);
+  },
+};

benchmark/visitInParallel-benchmark.js

@@ -0,0 +1,25 @@
+'use strict';
+
+const { parse } = require('graphql/language/parser.js');
+const { visit, visitInParallel } = require('graphql/language/visitor.js');
+
+const { bigSchemaSDL } = require('./fixtures.js');
+
+const documentAST = parse(bigSchemaSDL);
+
+const visitors = new Array(50).fill({
+  enter() {
+    /* do nothing */
+  },
+  leave() {
+    /* do nothing */
+  },
+});
+
+module.exports = {
+  name: 'Visit all AST nodes in parallel',
+  count: 10,
+  measure() {
+    visit(documentAST, visitInParallel(visitors));
+  },
+};

src/index.ts

@@ -209,6 +209,7 @@ export {
   visit,
   visitInParallel,
   getVisitFn,
+  getEnterLeaveForKind,
   BREAK,
   Kind,
   DirectiveLocation,

src/language/__tests__/visitor-test.ts

@@ -420,49 +420,6 @@ describe('Visitor', () => {
     ]);
   });
 
-  it('visit nodes with unknown kinds but does not traverse deeper', () => {
-    const customAST = parse('{ a }');
-    // @ts-expect-error
-    customAST.definitions[0].selectionSet.selections.push({
-      kind: 'CustomField',
-      name: { kind: 'Name', value: 'NamedNodeToBeSkipped' },
-      selectionSet: {
-        kind: 'SelectionSet',
-        selections: [
-          {
-            kind: 'CustomField',
-            name: { kind: 'Name', value: 'NamedNodeToBeSkipped' },
-          },
-        ],
-      },
-    });
-
-    const visited: Array<any> = [];
-    visit(customAST, {
-      enter(node) {
-        visited.push(['enter', node.kind, getValue(node)]);
-      },
-      leave(node) {
-        visited.push(['leave', node.kind, getValue(node)]);
-      },
-    });
-
-    expect(visited).to.deep.equal([
-      ['enter', 'Document', undefined],
-      ['enter', 'OperationDefinition', undefined],
-      ['enter', 'SelectionSet', undefined],
-      ['enter', 'Field', undefined],
-      ['enter', 'Name', 'a'],
-      ['leave', 'Name', 'a'],
-      ['leave', 'Field', undefined],
-      ['enter', 'CustomField', undefined],
-      ['leave', 'CustomField', undefined],
-      ['leave', 'SelectionSet', undefined],
-      ['leave', 'OperationDefinition', undefined],
-      ['leave', 'Document', undefined],
-    ]);
-  });
-
   it('visits only the specified `Kind` in visitorKeyMap', () => {
     const visited: Array<any> = [];
 

src/language/index.ts

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

src/language/visitor.ts

@@ -1,8 +1,8 @@
 import { inspect } from '../jsutils/inspect';
-import type { Maybe } from '../jsutils/Maybe';
 
 import type { ASTNode, ASTKindToNode } from './ast';
 import { isNode } from './ast';
+import { Kind } from './kinds';
 
 /**
  * A visitor is provided to visit, it contains the collection of
@@ -256,6 +256,15 @@ export function visit(
   visitor: ASTVisitor | ASTReducer<any>,
   visitorKeys: ASTVisitorKeyMap = QueryDocumentKeys,
 ): any {
+  const enterLeaveMap = new Map<
+    keyof ASTKindToNode,
+    EnterLeaveVisitor<ASTNode>
+  >();
+
+  for (const kind of Object.values(Kind)) {
+    enterLeaveMap.set(kind, getEnterLeaveForKind(visitor, kind));
+  }
+
   /* eslint-disable no-undef-init */
   let stack: any = undefined;
   let inArray = Array.isArray(root);
@@ -318,29 +327,30 @@ export function visit(
       if (!isNode(node)) {
         throw new Error(`Invalid AST Node: ${inspect(node)}.`);
       }
-      const visitFn = getVisitFn(visitor, node.kind, isLeaving);
-      if (visitFn) {
-        result = visitFn.call(visitor, node, key, parent, path, ancestors);
+      const visitFn = isLeaving
+        ? enterLeaveMap.get(node.kind)?.leave
+        : enterLeaveMap.get(node.kind)?.enter;
 
-        if (result === BREAK) {
-          break;
-        }
+      result = visitFn?.call(visitor, node, key, parent, path, ancestors);
+
+      if (result === BREAK) {
+        break;
+      }
 
-        if (result === false) {
-          if (!isLeaving) {
+      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;
           }
-        } else if (result !== undefined) {
-          edits.push([key, result]);
-          if (!isLeaving) {
-            if (isNode(result)) {
-              node = result;
-            } else {
-              path.pop();
-              continue;
-            }
-          }
         }
       }
     }
@@ -380,16 +390,31 @@ export function visit(
 export function visitInParallel(
   visitors: ReadonlyArray<ASTVisitor>,
 ): ASTVisitor {
-  const skipping = new Array(visitors.length);
-
-  return {
-    enter(...args) {
-      const node = args[0];
-      for (let i = 0; i < visitors.length; i++) {
-        if (skipping[i] == null) {
-          const fn = getVisitFn(visitors[i], node.kind, /* isLeaving */ false);
-          if (fn) {
-            const result = fn.apply(visitors[i], args);
+  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) {
@@ -399,50 +424,65 @@ export function visitInParallel(
             }
           }
         }
-      }
-    },
-    leave(...args) {
-      const node = args[0];
-      for (let i = 0; i < visitors.length; i++) {
-        if (skipping[i] == null) {
-          const fn = getVisitFn(visitors[i], node.kind, /* isLeaving */ true);
-          if (fn) {
-            const result = fn.apply(visitors[i], args);
+      },
+      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;
           }
-        } else if (skipping[i] === node) {
-          skipping[i] = null;
         }
-      }
-    },
-  };
+      },
+    };
+
+    mergedVisitor[kind] = mergedEnterLeave;
+  }
+
+  return mergedVisitor;
 }
 
 /**
- * Given a visitor instance, if it is leaving or not, and a node kind, return
- * the function the visitor runtime should call.
+ * Given a visitor instance and a node kind, return EnterLeaveVisitor for that kind.
  */
-export function getVisitFn(
+export function getEnterLeaveForKind(
   visitor: ASTVisitor,
   kind: keyof ASTKindToNode,
-  isLeaving: boolean,
-): Maybe<ASTVisitFn<ASTNode>> {
+): EnterLeaveVisitor<ASTNode> {
   const kindVisitor:
     | ASTVisitFn<ASTNode>
     | EnterLeaveVisitor<ASTNode>
     | undefined = (visitor as any)[kind];
-  if (kindVisitor) {
-    if (typeof kindVisitor === 'function') {
-      // { Kind() {} }
-      return isLeaving ? undefined : kindVisitor;
-    }
+
+  if (typeof kindVisitor === 'object') {
     // { Kind: { enter() {}, leave() {} } }
-    return isLeaving ? kindVisitor.leave : kindVisitor.enter;
+    return kindVisitor;
+  } else if (typeof kindVisitor === 'function') {
+    // { Kind() {} }
+    return { enter: kindVisitor, leave: undefined };
   }
+
   // { enter() {}, leave() {} }
-  return isLeaving ? (visitor as any).leave : (visitor as any).enter;
+  return { enter: (visitor as any).enter, leave: (visitor as any).leave };
+}
+
+/**
+ * Given a visitor instance, if it is leaving or not, and a node kind, return
+ * the function the visitor runtime should call.
+ *
+ * @deprecated Please use `getEnterLeaveForKind` instead. Will be removed in v17
+ */
+export function getVisitFn(
+  visitor: ASTVisitor,
+  kind: keyof ASTKindToNode,
+  isLeaving: boolean,
+): ASTVisitFn<ASTNode> | undefined {
+  const { enter, leave } = getEnterLeaveForKind(visitor, kind);
+  return isLeaving ? leave : enter;
 }

src/utilities/TypeInfo.ts

@@ -2,7 +2,7 @@ import type { ASTVisitor } from '../language/visitor';
 import type { ASTNode, FieldNode } from '../language/ast';
 import { Kind } from '../language/kinds';
 import { isNode } from '../language/ast';
-import { getVisitFn } from '../language/visitor';
+import { getEnterLeaveForKind } from '../language/visitor';
 
 import type { Maybe } from '../jsutils/Maybe';
 
@@ -344,7 +344,7 @@ export function visitWithTypeInfo(
     enter(...args) {
       const node = args[0];
       typeInfo.enter(node);
-      const fn = getVisitFn(visitor, node.kind, /* isLeaving */ false);
+      const fn = getEnterLeaveForKind(visitor, node.kind).enter;
       if (fn) {
         const result = fn.apply(visitor, args);
         if (result !== undefined) {
@@ -358,7 +358,7 @@ export function visitWithTypeInfo(
     },
     leave(...args) {
       const node = args[0];
-      const fn = getVisitFn(visitor, node.kind, /* isLeaving */ true);
+      const fn = getEnterLeaveForKind(visitor, node.kind).leave;
       let result;
       if (fn) {
         result = fn.apply(visitor, args);