From 349096036ec1625f8f4b8b8a1e328ddb073d4d5a Mon Sep 17 00:00:00 2001
From: Ivan Goncharov <ivan.goncharov.ua@gmail.com>
Date: Wed, 22 Aug 2018 15:54:10 +0300
Subject: [PATCH] Add 'getDefinitionsMap' to 'ASTValidationContext'

---
 src/validation/ValidationContext.js           | 40 +++++++++++------
 src/validation/rules/KnownDirectives.js       |  8 ++--
 .../rules/LoneAnonymousOperation.js           |  9 +---
 src/validation/rules/NoUnusedFragments.js     | 43 +++++++------------
 4 files changed, 47 insertions(+), 53 deletions(-)

diff --git a/src/validation/ValidationContext.js b/src/validation/ValidationContext.js
index 170cb36787..7e462e4551 100644
--- a/src/validation/ValidationContext.js
+++ b/src/validation/ValidationContext.js
@@ -7,11 +7,13 @@
  * @flow strict
  */
 
+import keyMap from '../jsutils/keyMap';
 import type { ObjMap } from '../jsutils/ObjMap';
 import type { GraphQLError } from '../error';
 import { visit, visitWithTypeInfo } from '../language/visitor';
 import { Kind } from '../language/kinds';
 import type {
+  ASTKindToNode,
   DocumentNode,
   OperationDefinitionNode,
   VariableNode,
@@ -38,6 +40,10 @@ type VariableUsage = {|
   +defaultValue: ?mixed,
 |};
 
+type KindToNodesMap = $Shape<
+  $ObjMap<ASTKindToNode, <Node>(Node) => ?Array<Node>>,
+>;
+
 /**
  * An instance of this class is passed as the "this" context to all validators,
  * allowing access to commonly useful contextual information from within a
@@ -46,19 +52,29 @@ type VariableUsage = {|
 export class ASTValidationContext {
   _ast: DocumentNode;
   _errors: Array<GraphQLError>;
-  _fragments: ?ObjMap<FragmentDefinitionNode>;
+  _fragments: ObjMap<FragmentDefinitionNode>;
   _fragmentSpreads: Map<SelectionSetNode, $ReadOnlyArray<FragmentSpreadNode>>;
   _recursivelyReferencedFragments: Map<
     OperationDefinitionNode,
     $ReadOnlyArray<FragmentDefinitionNode>,
   >;
+  _defsByKind: KindToNodesMap;
 
   constructor(ast: DocumentNode): void {
     this._ast = ast;
     this._errors = [];
-    this._fragments = undefined;
     this._fragmentSpreads = new Map();
     this._recursivelyReferencedFragments = new Map();
+
+    const defsByKind = Object.create(null);
+    for (const node of ast.definitions) {
+      if (defsByKind[node.kind]) {
+        defsByKind[node.kind].push(node);
+      } else {
+        defsByKind[node.kind] = [node];
+      }
+    }
+    this._defsByKind = defsByKind;
   }
 
   reportError(error: GraphQLError): void {
@@ -73,20 +89,18 @@ export class ASTValidationContext {
     return this._ast;
   }
 
+  getDefinitionsMap(): KindToNodesMap {
+    return this._defsByKind;
+  }
+
   getFragment(name: string): ?FragmentDefinitionNode {
-    let fragments = this._fragments;
-    if (!fragments) {
-      this._fragments = fragments = this.getDocument().definitions.reduce(
-        (frags, statement) => {
-          if (statement.kind === Kind.FRAGMENT_DEFINITION) {
-            frags[statement.name.value] = statement;
-          }
-          return frags;
-        },
-        Object.create(null),
+    if (!this._fragments) {
+      this._fragments = keyMap(
+        this.getDefinitionsMap().FragmentDefinition || [],
+        def => def.name.value,
       );
     }
-    return fragments[name];
+    return this._fragments[name];
   }
 
   getFragmentSpreads(
diff --git a/src/validation/rules/KnownDirectives.js b/src/validation/rules/KnownDirectives.js
index 52c9a7b005..3ed4574315 100644
--- a/src/validation/rules/KnownDirectives.js
+++ b/src/validation/rules/KnownDirectives.js
@@ -47,11 +47,9 @@ export function KnownDirectives(
     locationsMap[directive.name] = directive.locations;
   }
 
-  const astDefinitions = context.getDocument().definitions;
-  for (const def of astDefinitions) {
-    if (def.kind === Kind.DIRECTIVE_DEFINITION) {
-      locationsMap[def.name.value] = def.locations.map(name => name.value);
-    }
+  const defNodes = context.getDefinitionsMap().DirectiveDefinition || [];
+  for (const node of defNodes) {
+    locationsMap[node.name.value] = node.locations.map(name => name.value);
   }
 
   return {
diff --git a/src/validation/rules/LoneAnonymousOperation.js b/src/validation/rules/LoneAnonymousOperation.js
index 2faad31468..398d425902 100644
--- a/src/validation/rules/LoneAnonymousOperation.js
+++ b/src/validation/rules/LoneAnonymousOperation.js
@@ -9,7 +9,6 @@
 
 import type { ASTValidationContext } from '../ValidationContext';
 import { GraphQLError } from '../../error/GraphQLError';
-import { Kind } from '../../language/kinds';
 import type { ASTVisitor } from '../../language/visitor';
 
 export function anonOperationNotAloneMessage(): string {
@@ -25,13 +24,9 @@ export function anonOperationNotAloneMessage(): string {
 export function LoneAnonymousOperation(
   context: ASTValidationContext,
 ): ASTVisitor {
-  let operationCount = 0;
+  const operationDefs = context.getDefinitionsMap().OperationDefinition;
+  const operationCount = operationDefs ? operationDefs.length : 0;
   return {
-    Document(node) {
-      operationCount = node.definitions.filter(
-        definition => definition.kind === Kind.OPERATION_DEFINITION,
-      ).length;
-    },
     OperationDefinition(node) {
       if (!node.name && operationCount > 1) {
         context.reportError(
diff --git a/src/validation/rules/NoUnusedFragments.js b/src/validation/rules/NoUnusedFragments.js
index 486c7b6dbf..1254b82b34 100644
--- a/src/validation/rules/NoUnusedFragments.js
+++ b/src/validation/rules/NoUnusedFragments.js
@@ -22,38 +22,25 @@ export function unusedFragMessage(fragName: string): string {
  * within operations, or spread within other fragments spread within operations.
  */
 export function NoUnusedFragments(context: ASTValidationContext): ASTVisitor {
-  const operationDefs = [];
-  const fragmentDefs = [];
+  const fragmentNameUsed = Object.create(null);
+  const operationDefs = context.getDefinitionsMap().OperationDefinition || [];
+  for (const operation of operationDefs) {
+    for (const fragment of context.getRecursivelyReferencedFragments(
+      operation,
+    )) {
+      fragmentNameUsed[fragment.name.value] = true;
+    }
+  }
 
   return {
-    OperationDefinition(node) {
-      operationDefs.push(node);
-      return false;
-    },
     FragmentDefinition(node) {
-      fragmentDefs.push(node);
+      const fragName = node.name.value;
+      if (fragmentNameUsed[fragName] !== true) {
+        context.reportError(
+          new GraphQLError(unusedFragMessage(fragName), [node]),
+        );
+      }
       return false;
     },
-    Document: {
-      leave() {
-        const fragmentNameUsed = Object.create(null);
-        for (const operation of operationDefs) {
-          for (const fragment of context.getRecursivelyReferencedFragments(
-            operation,
-          )) {
-            fragmentNameUsed[fragment.name.value] = true;
-          }
-        }
-
-        for (const fragmentDef of fragmentDefs) {
-          const fragName = fragmentDef.name.value;
-          if (fragmentNameUsed[fragName] !== true) {
-            context.reportError(
-              new GraphQLError(unusedFragMessage(fragName), [fragmentDef]),
-            );
-          }
-        }
-      },
-    },
   };
 }