refactor(incremental): simplify incremental graph by allowing mutations

src/execution/IncrementalGraph.ts

@@ -15,32 +15,16 @@ import type {
   StreamRecord,
   SubsequentResultRecord,
 } from './types.js';
-import { isDeferredGroupedFieldSetRecord } from './types.js';
-
-interface DeferredFragmentNode {
-  deferredFragmentRecord: DeferredFragmentRecord;
-  deferredGroupedFieldSetRecords: Set<DeferredGroupedFieldSetRecord>;
-  reconcilableResults: Set<ReconcilableDeferredGroupedFieldSetResult>;
-  children: Set<SubsequentResultNode>;
-}
-
-function isDeferredFragmentNode(
-  node: SubsequentResultNode | undefined,
-): node is DeferredFragmentNode {
-  return node !== undefined && 'deferredFragmentRecord' in node;
-}
-
-type SubsequentResultNode = DeferredFragmentNode | StreamRecord;
+import {
+  isDeferredFragmentRecord,
+  isDeferredGroupedFieldSetRecord,
+} from './types.js';
 
 /**
  * @internal
  */
 export class IncrementalGraph {
-  private _rootNodes: Set<SubsequentResultNode>;
-  private _deferredFragmentNodes: Map<
-    DeferredFragmentRecord,
-    DeferredFragmentNode
-  >;
+  private _rootNodes: Set<SubsequentResultRecord>;
 
   private _completedQueue: Array<IncrementalDataRecordResult>;
   private _nextQueue: Array<
@@ -49,15 +33,14 @@ export class IncrementalGraph {
 
   constructor() {
     this._rootNodes = new Set();
-    this._deferredFragmentNodes = new Map();
     this._completedQueue = [];
     this._nextQueue = [];
   }
 
   getNewRootNodes(
     incrementalDataRecords: ReadonlyArray<IncrementalDataRecord>,
   ): ReadonlyArray<SubsequentResultRecord> {
-    const initialResultChildren = new Set<SubsequentResultNode>();
+    const initialResultChildren = new Set<SubsequentResultRecord>();
     this._addIncrementalDataRecords(
       incrementalDataRecords,
       undefined,
@@ -69,13 +52,12 @@ export class IncrementalGraph {
   addCompletedReconcilableDeferredGroupedFieldSet(
     reconcilableResult: ReconcilableDeferredGroupedFieldSetResult,
   ): void {
-    for (const deferredFragmentNode of this._fragmentsToNodes(
-      reconcilableResult.deferredGroupedFieldSetRecord.deferredFragmentRecords,
-    )) {
-      deferredFragmentNode.deferredGroupedFieldSetRecords.delete(
+    for (const deferredFragmentRecord of reconcilableResult
+      .deferredGroupedFieldSetRecord.deferredFragmentRecords) {
+      deferredFragmentRecord.deferredGroupedFieldSetRecords.delete(
         reconcilableResult.deferredGroupedFieldSetRecord,
       );
-      deferredFragmentNode.reconcilableResults.add(reconcilableResult);
+      deferredFragmentRecord.reconcilableResults.add(reconcilableResult);
     }
 
     const incrementalDataRecords = reconcilableResult.incrementalDataRecords;
@@ -131,64 +113,50 @@ export class IncrementalGraph {
         reconcilableResults: ReadonlyArray<ReconcilableDeferredGroupedFieldSetResult>;
       }
     | undefined {
-    const deferredFragmentNode = this._deferredFragmentNodes.get(
-      deferredFragmentRecord,
-    );
     // TODO: add test case?
     /* c8 ignore next 3 */
-    if (deferredFragmentNode === undefined) {
+    if (!this._rootNodes.has(deferredFragmentRecord)) {
       return;
     }
-    if (deferredFragmentNode.deferredGroupedFieldSetRecords.size > 0) {
+    if (deferredFragmentRecord.deferredGroupedFieldSetRecords.size > 0) {
       return;
     }
     const reconcilableResults = Array.from(
-      deferredFragmentNode.reconcilableResults,
+      deferredFragmentRecord.reconcilableResults,
     );
-    this._removeRootNode(deferredFragmentNode);
+    this._removeRootNode(deferredFragmentRecord);
     for (const reconcilableResult of reconcilableResults) {
-      for (const otherDeferredFragmentNode of this._fragmentsToNodes(
-        reconcilableResult.deferredGroupedFieldSetRecord
-          .deferredFragmentRecords,
-      )) {
-        otherDeferredFragmentNode.reconcilableResults.delete(
+      for (const otherDeferredFragmentRecord of reconcilableResult
+        .deferredGroupedFieldSetRecord.deferredFragmentRecords) {
+        otherDeferredFragmentRecord.reconcilableResults.delete(
           reconcilableResult,
         );
       }
     }
     const newRootNodes = this._promoteNonEmptyToRoot(
-      deferredFragmentNode.children,
+      deferredFragmentRecord.children,
     );
     return { newRootNodes, reconcilableResults };
   }
 
   removeDeferredFragment(
     deferredFragmentRecord: DeferredFragmentRecord,
   ): boolean {
-    const deferredFragmentNode = this._deferredFragmentNodes.get(
-      deferredFragmentRecord,
-    );
-    if (deferredFragmentNode === undefined) {
+    if (!this._rootNodes.has(deferredFragmentRecord)) {
       return false;
     }
-    this._removeRootNode(deferredFragmentNode);
-    this._deferredFragmentNodes.delete(deferredFragmentRecord);
-    // TODO: add test case for an erroring deferred fragment with child defers
-    /* c8 ignore next 5 */
-    for (const child of deferredFragmentNode.children) {
-      if (isDeferredFragmentNode(child)) {
-        this.removeDeferredFragment(child.deferredFragmentRecord);
-      }
-    }
+    this._removeRootNode(deferredFragmentRecord);
     return true;
   }
 
   removeStream(streamRecord: StreamRecord): void {
     this._removeRootNode(streamRecord);
   }
 
-  private _removeRootNode(subsequentResultNode: SubsequentResultNode): void {
-    this._rootNodes.delete(subsequentResultNode);
+  private _removeRootNode(
+    subsequentResultRecord: SubsequentResultRecord,
+  ): void {
+    this._rootNodes.delete(subsequentResultRecord);
     if (this._rootNodes.size === 0) {
       for (const resolve of this._nextQueue) {
         resolve({ value: undefined, done: true });
@@ -199,16 +167,16 @@ export class IncrementalGraph {
   private _addIncrementalDataRecords(
     incrementalDataRecords: ReadonlyArray<IncrementalDataRecord>,
     parents: ReadonlyArray<DeferredFragmentRecord> | undefined,
-    initialResultChildren?: Set<SubsequentResultNode> | undefined,
+    initialResultChildren?: Set<SubsequentResultRecord> | undefined,
   ): void {
     for (const incrementalDataRecord of incrementalDataRecords) {
       if (isDeferredGroupedFieldSetRecord(incrementalDataRecord)) {
         for (const deferredFragmentRecord of incrementalDataRecord.deferredFragmentRecords) {
-          const deferredFragmentNode = this._addDeferredFragmentNode(
+          this._addDeferredFragment(
             deferredFragmentRecord,
             initialResultChildren,
           );
-          deferredFragmentNode.deferredGroupedFieldSetRecords.add(
+          deferredFragmentRecord.deferredGroupedFieldSetRecords.add(
             incrementalDataRecord,
           );
         }
@@ -220,33 +188,29 @@ export class IncrementalGraph {
         initialResultChildren.add(incrementalDataRecord);
       } else {
         for (const parent of parents) {
-          const deferredFragmentNode = this._addDeferredFragmentNode(
-            parent,
-            initialResultChildren,
-          );
-          deferredFragmentNode.children.add(incrementalDataRecord);
+          this._addDeferredFragment(parent, initialResultChildren);
+          parent.children.add(incrementalDataRecord);
         }
       }
     }
   }
 
   private _promoteNonEmptyToRoot(
-    maybeEmptyNewRootNodes: Set<SubsequentResultNode>,
+    maybeEmptyNewRootNodes: Set<SubsequentResultRecord>,
   ): ReadonlyArray<SubsequentResultRecord> {
     const newRootNodes: Array<SubsequentResultRecord> = [];
     for (const node of maybeEmptyNewRootNodes) {
-      if (isDeferredFragmentNode(node)) {
+      if (isDeferredFragmentRecord(node)) {
         if (node.deferredGroupedFieldSetRecords.size > 0) {
           for (const deferredGroupedFieldSetRecord of node.deferredGroupedFieldSetRecords) {
             if (!this._completesRootNode(deferredGroupedFieldSetRecord)) {
               this._onDeferredGroupedFieldSet(deferredGroupedFieldSetRecord);
             }
           }
           this._rootNodes.add(node);
-          newRootNodes.push(node.deferredFragmentRecord);
+          newRootNodes.push(node);
           continue;
         }
-        this._deferredFragmentNodes.delete(node.deferredFragmentRecord);
         for (const child of node.children) {
           maybeEmptyNewRootNodes.add(child);
         }
@@ -264,53 +228,26 @@ export class IncrementalGraph {
   private _completesRootNode(
     deferredGroupedFieldSetRecord: DeferredGroupedFieldSetRecord,
   ): boolean {
-    return this._fragmentsToNodes(
-      deferredGroupedFieldSetRecord.deferredFragmentRecords,
-    ).some((node) => this._rootNodes.has(node));
-  }
-
-  private _fragmentsToNodes(
-    deferredFragmentRecords: ReadonlyArray<DeferredFragmentRecord>,
-  ): Array<DeferredFragmentNode> {
-    return deferredFragmentRecords
-      .map((deferredFragmentRecord) =>
-        this._deferredFragmentNodes.get(deferredFragmentRecord),
-      )
-      .filter<DeferredFragmentNode>(isDeferredFragmentNode);
+    return deferredGroupedFieldSetRecord.deferredFragmentRecords.some(
+      (deferredFragmentRecord) => this._rootNodes.has(deferredFragmentRecord),
+    );
   }
 
-  private _addDeferredFragmentNode(
+  private _addDeferredFragment(
     deferredFragmentRecord: DeferredFragmentRecord,
-    initialResultChildren: Set<SubsequentResultNode> | undefined,
-  ): DeferredFragmentNode {
-    let deferredFragmentNode = this._deferredFragmentNodes.get(
-      deferredFragmentRecord,
-    );
-    if (deferredFragmentNode !== undefined) {
-      return deferredFragmentNode;
+    initialResultChildren: Set<SubsequentResultRecord> | undefined,
+  ): void {
+    if (this._rootNodes.has(deferredFragmentRecord)) {
+      return;
     }
-    deferredFragmentNode = {
-      deferredFragmentRecord,
-      deferredGroupedFieldSetRecords: new Set(),
-      reconcilableResults: new Set(),
-      children: new Set(),
-    };
-    this._deferredFragmentNodes.set(
-      deferredFragmentRecord,
-      deferredFragmentNode,
-    );
     const parent = deferredFragmentRecord.parent;
     if (parent === undefined) {
       invariant(initialResultChildren !== undefined);
-      initialResultChildren.add(deferredFragmentNode);
-      return deferredFragmentNode;
+      initialResultChildren.add(deferredFragmentRecord);
+      return;
     }
-    const parentNode = this._addDeferredFragmentNode(
-      parent,
-      initialResultChildren,
-    );
-    parentNode.children.add(deferredFragmentNode);
-    return deferredFragmentNode;
+    parent.children.add(deferredFragmentRecord);
+    this._addDeferredFragment(parent, initialResultChildren);
   }
 
   private _onDeferredGroupedFieldSet(

src/execution/execute.ts

@@ -63,7 +63,6 @@ import { buildIncrementalResponse } from './IncrementalPublisher.js';
 import { mapAsyncIterable } from './mapAsyncIterable.js';
 import type {
   CancellableStreamRecord,
-  DeferredFragmentRecord,
   DeferredGroupedFieldSetRecord,
   DeferredGroupedFieldSetResult,
   ExecutionResult,
@@ -73,6 +72,7 @@ import type {
   StreamItemResult,
   StreamRecord,
 } from './types.js';
+import { DeferredFragmentRecord } from './types.js';
 import {
   getArgumentValues,
   getDirectiveValues,
@@ -1676,11 +1676,11 @@ function addNewDeferredFragments(
         : deferredFragmentRecordFromDeferUsage(parentDeferUsage, newDeferMap);
 
     // Instantiate the new record.
-    const deferredFragmentRecord: DeferredFragmentRecord = {
+    const deferredFragmentRecord = new DeferredFragmentRecord(
       path,
-      label: newDeferUsage.label,
+      newDeferUsage.label,
       parent,
-    };
+    );
 
     // Update the map.
     newDeferMap.set(newDeferUsage, deferredFragmentRecord);

src/execution/types.ts

@@ -214,11 +214,34 @@ export interface DeferredGroupedFieldSetRecord {
 
 export type SubsequentResultRecord = DeferredFragmentRecord | StreamRecord;
 
-export interface DeferredFragmentRecord {
+/** @internal */
+export class DeferredFragmentRecord {
   path: Path | undefined;
   label: string | undefined;
   id?: string | undefined;
   parent: DeferredFragmentRecord | undefined;
+  deferredGroupedFieldSetRecords: Set<DeferredGroupedFieldSetRecord>;
+  reconcilableResults: Set<ReconcilableDeferredGroupedFieldSetResult>;
+  children: Set<SubsequentResultRecord>;
+
+  constructor(
+    path: Path | undefined,
+    label: string | undefined,
+    parent: DeferredFragmentRecord | undefined,
+  ) {
+    this.path = path;
+    this.label = label;
+    this.parent = parent;
+    this.deferredGroupedFieldSetRecords = new Set();
+    this.reconcilableResults = new Set();
+    this.children = new Set();
+  }
+}
+
+export function isDeferredFragmentRecord(
+  subsequentResultRecord: SubsequentResultRecord,
+): subsequentResultRecord is DeferredFragmentRecord {
+  return subsequentResultRecord instanceof DeferredFragmentRecord;
 }
 
 export interface StreamItemResult {