Refactor out AccumulatorMap from groupBy (#3568)

Author: IvanGoncharov

src/execution/collectFields.ts

@@ -1,3 +1,4 @@
+import { AccumulatorMap } from '../jsutils/AccumulatorMap';
 import type { ObjMap } from '../jsutils/ObjMap';
 
 import type {
@@ -37,7 +38,7 @@ export function collectFields(
   runtimeType: GraphQLObjectType,
   selectionSet: SelectionSetNode,
 ): Map<string, ReadonlyArray<FieldNode>> {
-  const fields = new Map();
+  const fields = new AccumulatorMap<string, FieldNode>();
   collectFieldsImpl(
     schema,
     fragments,
@@ -67,7 +68,7 @@ export function collectSubfields(
   returnType: GraphQLObjectType,
   fieldNodes: ReadonlyArray<FieldNode>,
 ): Map<string, ReadonlyArray<FieldNode>> {
-  const subFieldNodes = new Map();
+  const subFieldNodes = new AccumulatorMap<string, FieldNode>();
   const visitedFragmentNames = new Set<string>();
   for (const node of fieldNodes) {
     if (node.selectionSet) {
@@ -91,7 +92,7 @@ function collectFieldsImpl(
   variableValues: { [variable: string]: unknown },
   runtimeType: GraphQLObjectType,
   selectionSet: SelectionSetNode,
-  fields: Map<string, Array<FieldNode>>,
+  fields: AccumulatorMap<string, FieldNode>,
   visitedFragmentNames: Set<string>,
 ): void {
   for (const selection of selectionSet.selections) {
@@ -100,13 +101,7 @@ function collectFieldsImpl(
         if (!shouldIncludeNode(variableValues, selection)) {
           continue;
         }
-        const name = getFieldEntryKey(selection);
-        const fieldList = fields.get(name);
-        if (fieldList !== undefined) {
-          fieldList.push(selection);
-        } else {
-          fields.set(name, [selection]);
-        }
+        fields.add(getFieldEntryKey(selection), selection);
         break;
       }
       case Kind.INLINE_FRAGMENT: {

src/jsutils/AccumulatorMap.ts

@@ -0,0 +1,17 @@
+/**
+ * ES6 Map with additional `add` method to accumulate items.
+ */
+export class AccumulatorMap<K, T> extends Map<K, Array<T>> {
+  get [Symbol.toStringTag]() {
+    return 'AccumulatorMap';
+  }
+
+  add(key: K, item: T): void {
+    const group = this.get(key);
+    if (group === undefined) {
+      this.set(key, [item]);
+    } else {
+      group.push(item);
+    }
+  }
+}

src/jsutils/__tests__/AccumulatorMap-test.ts

@@ -0,0 +1,36 @@
+import { expect } from 'chai';
+import { describe, it } from 'mocha';
+
+import { AccumulatorMap } from '../AccumulatorMap';
+
+function expectMap<K, V>(map: Map<K, V>) {
+  return expect(Object.fromEntries(map.entries()));
+}
+
+describe('AccumulatorMap', () => {
+  it('can be Object.toStringified', () => {
+    const accumulatorMap = new AccumulatorMap();
+
+    expect(Object.prototype.toString.call(accumulatorMap)).to.equal(
+      '[object AccumulatorMap]',
+    );
+  });
+
+  it('accumulate items', () => {
+    const accumulatorMap = new AccumulatorMap<string, number>();
+
+    expectMap(accumulatorMap).to.deep.equal({});
+
+    accumulatorMap.add('a', 1);
+    accumulatorMap.add('b', 2);
+    accumulatorMap.add('c', 3);
+    accumulatorMap.add('b', 4);
+    accumulatorMap.add('c', 5);
+    accumulatorMap.add('c', 6);
+    expectMap(accumulatorMap).to.deep.equal({
+      a: [1],
+      b: [2, 4],
+      c: [3, 5, 6],
+    });
+  });
+});

src/jsutils/groupBy.ts

@@ -1,19 +1,15 @@
+import { AccumulatorMap } from './AccumulatorMap';
+
 /**
  * Groups array items into a Map, given a function to produce grouping key.
  */
 export function groupBy<K, T>(
   list: ReadonlyArray<T>,
   keyFn: (item: T) => K,
 ): Map<K, ReadonlyArray<T>> {
-  const result = new Map<K, Array<T>>();
+  const result = new AccumulatorMap<K, T>();
   for (const item of list) {
-    const key = keyFn(item);
-    const group = result.get(key);
-    if (group === undefined) {
-      result.set(key, [item]);
-    } else {
-      group.push(item);
-    }
+    result.add(keyFn(item), item);
   }
   return result;
 }