fix(apollo-forest-run): index covered operations (#661)

* fix(apollo-forest-run): index covered operations

* add benchmark

* Change files

---------

Co-authored-by: vrazuvaev <vrazuvaev@microsoft.com_msteamsmdb>

Author: vladar

change/@graphitation-apollo-forest-run-bd962896-4b98-4c62-853f-2cf1f3b74811.json

@@ -0,0 +1,7 @@
+{
+  "type": "patch",
+  "comment": "fix(apollo-forest-run): index covered operations",
+  "packageName": "@graphitation/apollo-forest-run",
+  "email": "vrazuvaev@microsoft.com_msteamsmdb",
+  "dependentChangeType": "patch"
+}

packages/apollo-forest-run-benchmark/src/scenarios.ts

@@ -1,5 +1,15 @@
 import type { ForestRun } from "@graphitation/apollo-forest-run";
 import type { Scenario, ScenarioContext } from "./types";
+import { parse } from "graphql";
+
+// Pre-parse 50 unique named queries to simulate a real app with many active operations.
+// Each has its own operation name and data, creating separate trees in the forest.
+// This is the setup that exposes O(n) scans in getCoveringOperationIds.
+const EXTRA_OPS_COUNT = 50;
+const extraOperations = Array.from({ length: EXTRA_OPS_COUNT }, (_, i) => ({
+  query: parse(`query BgQuery${i} { node${i}(id: "${i}") { id value } }`),
+  data: { [`node${i}`]: { __typename: `Node${i}`, id: `${i}`, value: i } },
+}));
 
 const addWatchers = (
   watcherCount: number,
@@ -404,4 +414,32 @@ export const scenarios = [
       };
     },
   },
+  {
+    name: "write-many-ops-one-field-change",
+    prepare: (ctx: ScenarioContext) => {
+      const { operations, CacheFactory, configuration, watcherCount } = ctx;
+      const cache = new CacheFactory(configuration);
+
+      // Populate the cache with many unrelated operations (simulating real app)
+      for (const { query, data } of extraOperations) {
+        cache.writeQuery({ query, data });
+      }
+
+      // Write the main query, add watchers, then measure a write with a change.
+      // Each watcher's diff triggers readOperation → applyTransformations →
+      // createChunkMatcher → getCoveringOperationIds which scans all trees.
+      const { data, query } = operations["complex-nested"];
+      cache.writeQuery({ query, data: data["complex-nested"] });
+      addWatchers(watcherCount, cache, query);
+
+      return {
+        run() {
+          return cache.writeQuery({
+            query,
+            data: data["complex-nested-single-field-change"],
+          });
+        },
+      };
+    },
+  },
 ] as const satisfies Scenario[];

packages/apollo-forest-run/src/__tests__/helpers/forest.ts

@@ -36,6 +36,7 @@ export function createTestForest(): IndexedForest {
     operationsByNodes: new Map(),
     operationsWithErrors: new Set(),
     operationsByName: new Map(),
+    operationsByCoveredName: new Map(),
     operationsByPartitions: new Map(),
     deletedNodes: new Set(),
   };

packages/apollo-forest-run/src/cache/draftHelpers.ts

@@ -248,15 +248,15 @@ function getCoveringOperationIds(
   let ids: Set<OperationId> | undefined;
 
   for (const layer of layers) {
-    // Forward: find ops that cover us (their covers list includes our name)
+    // Forward: find ops that cover us via pre-built index
     if (opName) {
-      for (const tree of layer.trees.values()) {
-        if (
-          tree.operation.id !== operation.id &&
-          tree.operation.covers.includes(opName)
-        ) {
-          if (!ids) ids = new Set();
-          ids.add(tree.operation.id);
+      const coveringIds = layer.operationsByCoveredName.get(opName);
+      if (coveringIds) {
+        for (const id of coveringIds) {
+          if (id !== operation.id) {
+            if (!ids) ids = new Set();
+            ids.add(id);
+          }
         }
       }
     }

packages/apollo-forest-run/src/cache/store.ts

@@ -27,6 +27,7 @@ export function createStore(_: CacheEnv): Store {
     trees: new Map(),
     operationsByNodes: new Map<NodeKey, Set<OperationId>>(),
     operationsByName: new Map(),
+    operationsByCoveredName: new Map(),
     operationsByPartitions: new Map(),
     operationsWithErrors: new Set<OperationDescriptor>(),
     extraRootIds: new Map<NodeKey, TypeName>(),
@@ -193,6 +194,7 @@ export function createOptimisticLayer(
     operationsByNodes: new Map(),
     operationsWithErrors: new Set(),
     operationsByName: new Map(),
+    operationsByCoveredName: new Map(),
     operationsByPartitions: new Map(),
     extraRootIds: new Map(),
     readResults: new Map(),
@@ -361,6 +363,9 @@ function removeDataTree(
   dataForest.readResults.delete(operation);
   dataForest.operationsWithErrors.delete(operation);
   dataForest.operationsByName.get(operation.name ?? "")?.delete(operation.id);
+  for (const coveredName of operation.covers) {
+    dataForest.operationsByCoveredName.get(coveredName)?.delete(operation.id);
+  }
   dataForest.operationsByPartitions.get(partition)?.delete(operation.id);
   optimisticReadResults.delete(operation);
   partialReadResults.delete(operation);
@@ -380,6 +385,7 @@ export function resetStore(store: Store): void {
   dataForest.operationsByNodes.clear();
   dataForest.operationsWithErrors.clear();
   dataForest.operationsByName.clear();
+  dataForest.operationsByCoveredName.clear();
   dataForest.operationsWithDanglingRefs.clear();
   dataForest.readResults.clear();
   operations.clear();

packages/apollo-forest-run/src/forest/addTree.ts

@@ -50,6 +50,12 @@ function trackOperationName(forest: IndexedForest, tree: IndexedTree) {
   const name = tree.operation.name;
   if (!name) return;
   getOrCreate(forest.operationsByName, name, newSet).add(tree.operation.id);
+
+  for (const coveredName of tree.operation.covers) {
+    getOrCreate(forest.operationsByCoveredName, coveredName, newSet).add(
+      tree.operation.id,
+    );
+  }
 }
 
 function trackPartitions(

packages/apollo-forest-run/src/forest/types.ts

@@ -125,6 +125,7 @@ export type IndexedForest = {
   operationsByNodes: Map<NodeKey, Set<OperationId>>; // May contain false positives
   operationsWithErrors: Set<OperationDescriptor>; // May contain false positives
   operationsByName: Map<string, Set<OperationId>>; // operationName → operation IDs
+  operationsByCoveredName: Map<string, Set<OperationId>>; // coveredName → IDs of ops whose covers list includes it
   operationsByPartitions: Map<string, Set<OperationId>>; // partition key => operation IDs
   deletedNodes: Set<NodeKey>;
 };