Include context in memoization, Factor out memoize3

Author: leebyron

src/execution/execute.js

@@ -12,6 +12,7 @@ import { GraphQLError, locatedError } from '../error';
 import invariant from '../jsutils/invariant';
 import isInvalid from '../jsutils/isInvalid';
 import isNullish from '../jsutils/isNullish';
+import memoize3 from '../jsutils/memoize3';
 import type { ObjMap } from '../jsutils/ObjMap';
 import type { MaybePromise } from '../jsutils/MaybePromise';
 
@@ -1245,25 +1246,12 @@ function collectAndExecuteSubfields(
  * type. Memoizing ensures the subfields are not repeatedly calculated, which
  * saves overhead when resolving lists of values.
  */
-const subfieldCache: WeakMap<
-  $ReadOnlyArray<FieldNode>,
-  WeakMap<GraphQLObjectType, ObjMap<Array<FieldNode>>>,
-> = new WeakMap();
-function collectSubfields(
+const collectSubfields = memoize3(_collectSubfields);
+function _collectSubfields(
   exeContext: ExecutionContext,
   returnType: GraphQLObjectType,
   fieldNodes: $ReadOnlyArray<FieldNode>,
 ): ObjMap<Array<FieldNode>> {
-  let cacheByReturnType = subfieldCache.get(fieldNodes);
-  if (cacheByReturnType) {
-    const cachedSubFieldNodes = cacheByReturnType.get(returnType);
-    if (cachedSubFieldNodes) {
-      return cachedSubFieldNodes;
-    }
-  } else {
-    cacheByReturnType = new WeakMap();
-    subfieldCache.set(fieldNodes, cacheByReturnType);
-  }
   let subFieldNodes = Object.create(null);
   const visitedFragmentNames = Object.create(null);
   for (let i = 0; i < fieldNodes.length; i++) {
@@ -1278,7 +1266,6 @@ function collectSubfields(
       );
     }
   }
-  cacheByReturnType.set(returnType, subFieldNodes);
   return subFieldNodes;
 }
 

src/jsutils/memoize3.js

@@ -0,0 +1,44 @@
+/**
+ * Copyright (c) 2017-present, Facebook, Inc.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * @flow
+ */
+
+/**
+ * Memoizes the provided three-argument function.
+ */
+export default function memoize3<T: (a1: any, a2: any, a3: any) => any>(
+  fn: T,
+): T {
+  let cache0;
+  function memoized(a1, a2, a3) {
+    if (!cache0) {
+      cache0 = new WeakMap();
+    }
+    let cache1 = cache0.get(a1);
+    let cache2;
+    if (cache1) {
+      cache2 = cache1.get(a2);
+      if (cache2) {
+        const cachedValue = cache2.get(a3);
+        if (cachedValue) {
+          return cachedValue;
+        }
+      }
+    } else {
+      cache1 = new WeakMap();
+      cache0.set(a1, cache1);
+    }
+    if (!cache2) {
+      cache2 = new WeakMap();
+      cache1.set(a2, cache2);
+    }
+    const newValue = fn.apply(this, arguments);
+    cache2.set(a3, newValue);
+    return newValue;
+  }
+  return (memoized: any);
+}