Merge pull request #18 from ai-2070/perf/linear-bulk-fold

Performance

Author: LZL0

.npmignore

@@ -17,6 +17,9 @@ Thumbs.db
 # Dev dependencies config
 renovate.json
 
+# Benchmarks (dev-only)
+bench/
+
 # Lock files (users have their own)
 package-lock.json
 pnpm-lock.yaml

bench/memex.bench.ts

@@ -0,0 +1,159 @@
+// Performance benchmarks for the bulk fold paths.
+//
+// Run with: npm run bench
+//
+// These guard the O(N) behavior of replay / import / cascade. The "fold
+// strategy contrast" group makes the win explicit: folding commands with the
+// immutable `applyCommand` clones the whole graph per command (the old,
+// quadratic shape) while the in-place fold used by replay stays linear.
+
+import { bench, describe } from "vitest";
+import { applyCommand } from "../src/reducer.js";
+import { createGraphState } from "../src/graph.js";
+import { createMemoryItem, createEdge } from "../src/helpers.js";
+import { replayCommands, replayFromEnvelopes } from "../src/replay.js";
+import { importSlice } from "../src/transplant.js";
+import type { MemexExport } from "../src/transplant.js";
+import { createIntentState } from "../src/intent.js";
+import { createTaskState } from "../src/task.js";
+import { cascadeRetract, getDependents } from "../src/integrity.js";
+import type {
+  GraphState,
+  MemoryItem,
+  MemoryCommand,
+  Edge,
+  EventEnvelope,
+} from "../src/types.js";
+
+const BASE = 1_700_000_000_000;
+const padId = (i: number): string => `m-${i.toString().padStart(8, "0")}`;
+
+function item(i: number, parents?: string[]): MemoryItem {
+  return createMemoryItem({
+    id: padId(i),
+    scope: "bench",
+    kind: "observation",
+    content: { text: `item ${i}` },
+    author: "agent:bench",
+    source_kind: "observed",
+    authority: 0.5,
+    created_at: BASE + i,
+    ...(parents ? { parents } : {}),
+  });
+}
+
+function createCommands(n: number): MemoryCommand[] {
+  const cmds: MemoryCommand[] = [];
+  for (let i = 0; i < n; i++)
+    cmds.push({ type: "memory.create", item: item(i) });
+  return cmds;
+}
+
+function createEnvelopes(n: number): EventEnvelope<MemoryCommand>[] {
+  const envs: EventEnvelope<MemoryCommand>[] = [];
+  for (let i = 0; i < n; i++) {
+    envs.push({
+      id: `e-${i}`,
+      namespace: "memory",
+      type: "memory.create",
+      // Reverse the timestamps so replayFromEnvelopes pays for the sort.
+      ts: new Date(BASE + (n - i)).toISOString(),
+      payload: { type: "memory.create", item: item(i) },
+    });
+  }
+  return envs;
+}
+
+function chainCommands(n: number): MemoryCommand[] {
+  const cmds: MemoryCommand[] = [];
+  for (let i = 0; i < n; i++) {
+    cmds.push({
+      type: "memory.create",
+      item: item(i, i > 0 ? [padId(i - 1)] : undefined),
+    });
+  }
+  return cmds;
+}
+
+function sliceOf(n: number): MemexExport {
+  const memories: MemoryItem[] = [];
+  const edges: Edge[] = [];
+  for (let i = 0; i < n; i++) {
+    memories.push(item(i));
+    if (i > 0) {
+      edges.push(
+        createEdge({
+          edge_id: `edge-${i}`,
+          from: padId(i),
+          to: padId(i - 1),
+          kind: "DERIVED_FROM",
+          author: "agent:bench",
+          source_kind: "observed",
+          authority: 0.5,
+        }),
+      );
+    }
+  }
+  return { memories, edges, intents: [], tasks: [] };
+}
+
+describe("replayCommands", () => {
+  const c5k = createCommands(5_000);
+  const c20k = createCommands(20_000);
+  bench("5k creates", () => {
+    replayCommands(c5k);
+  });
+  bench("20k creates", () => {
+    replayCommands(c20k);
+  });
+});
+
+describe("replayFromEnvelopes (sort + fold)", () => {
+  const e5k = createEnvelopes(5_000);
+  bench("5k reverse-sorted creates", () => {
+    replayFromEnvelopes(e5k);
+  });
+});
+
+describe("importSlice (memories + edges into empty graph)", () => {
+  const s5k = sliceOf(5_000);
+  const s10k = sliceOf(10_000);
+  bench("5k memories", () => {
+    importSlice(
+      createGraphState(),
+      createIntentState(),
+      createTaskState(),
+      s5k,
+    );
+  });
+  bench("10k memories", () => {
+    importSlice(
+      createGraphState(),
+      createIntentState(),
+      createTaskState(),
+      s10k,
+    );
+  });
+});
+
+describe("cascadeRetract / getDependents (5k-deep chain)", () => {
+  const { state } = replayCommands(chainCommands(5_000));
+  const root = padId(0);
+  bench("cascadeRetract whole chain", () => {
+    cascadeRetract(state, root, "agent:bench");
+  });
+  bench("getDependents transitive", () => {
+    getDependents(state, root, true);
+  });
+});
+
+describe("fold strategy contrast (2k creates)", () => {
+  const cmds = createCommands(2_000);
+  bench("immutable applyCommand fold — clone per command (~O(N^2))", () => {
+    let state: GraphState = createGraphState();
+    for (const cmd of cmds) state = applyCommand(state, cmd).state;
+  });
+  bench("in-place fold via replayCommands (O(N))", () => {
+    replayCommands(cmds);
+  });
+});

package.json

@@ -19,8 +19,9 @@
     "build": "tsc",
     "test": "vitest run",
     "test:watch": "vitest",
-    "prettier": "prettier --write \"src/**/*.ts\" \"tests/**/*.ts\"",
-    "prettier:check": "prettier --check \"src/**/*.ts\" \"tests/**/*.ts\""
+    "bench": "vitest bench --run",
+    "prettier": "prettier --write \"src/**/*.ts\" \"tests/**/*.ts\" \"bench/**/*.ts\"",
+    "prettier:check": "prettier --check \"src/**/*.ts\" \"tests/**/*.ts\" \"bench/**/*.ts\""
   },
   "license": "Apache-2.0",
   "dependencies": {

src/integrity.ts

@@ -6,8 +6,13 @@ import type {
   ScoreWeights,
   ScoredItem,
 } from "./types.js";
-import { applyCommand } from "./reducer.js";
-import { getEdges, getChildren, getScoredItems } from "./query.js";
+import { applyCommand, retractItemsInPlace } from "./reducer.js";
+import {
+  getEdges,
+  getChildren,
+  getScoredItems,
+  buildChildrenIndex,
+} from "./query.js";
 import { uuidv7 } from "uuidv7";
 
 // ---------------------------------------------------------------------------
@@ -181,19 +186,23 @@ export function getDependents(
   itemId: string,
   transitive = false,
 ): MemoryItem[] {
-  const direct = getChildren(state, itemId);
-  if (!transitive) return direct;
+  if (!transitive) return getChildren(state, itemId);
 
+  // Walk the dependency tree off a single children index instead of re-scanning
+  // the whole graph at every node (which made the transitive walk O(nodes x
+  // items)).
+  const childrenIndex = buildChildrenIndex(state);
   const visited = new Set<string>();
   const result: MemoryItem[] = [];
-  const queue = [...direct];
+  const queue = [...(childrenIndex.get(itemId) ?? [])];
 
   while (queue.length > 0) {
     const item = queue.pop()!;
     if (visited.has(item.id)) continue;
     visited.add(item.id);
     result.push(item);
-    queue.push(...getChildren(state, item.id));
+    const children = childrenIndex.get(item.id);
+    if (children) for (const child of children) queue.push(child);
   }
 
   return result;
@@ -215,12 +224,17 @@ export function cascadeRetract(
   // Pre-mark the root as visited so any cycle that points back to it is
   // ignored — the root is retracted separately at the end of this function,
   // never prematurely as part of the descendants list.
+  // Children index built once: getChildren is O(items), and walking it per node
+  // made the traversal O(nodes x items) on top of the per-retract clone below.
+  const childrenIndex = buildChildrenIndex(state);
+  const childrenOf = (id: string): MemoryItem[] => childrenIndex.get(id) ?? [];
+
   const visited = new Set<string>([itemId]);
   const order: string[] = [];
 
   type Frame = { id: string; phase: "enter" | "exit" };
   const stack: Frame[] = [];
-  for (const child of getChildren(state, itemId)) {
+  for (const child of childrenOf(itemId)) {
     stack.push({ id: child.id, phase: "enter" });
   }
 
@@ -234,43 +248,26 @@ export function cascadeRetract(
     visited.add(frame.id);
     // Push exit first so it's processed after all children (post-order).
     stack.push({ id: frame.id, phase: "exit" });
-    for (const child of getChildren(state, frame.id)) {
+    for (const child of childrenOf(frame.id)) {
       if (!visited.has(child.id)) {
         stack.push({ id: child.id, phase: "enter" });
       }
     }
   }
 
-  let current = state;
-  const allEvents: MemoryLifecycleEvent[] = [];
-  const retracted: string[] = [];
-
-  for (const depId of order) {
-    if (!current.items.has(depId)) continue;
-    const r = applyCommand(current, {
-      type: "memory.retract",
-      item_id: depId,
-      author,
-      reason: reason ?? `parent ${itemId} retracted`,
-    });
-    current = r.state;
-    allEvents.push(...r.events);
-    retracted.push(depId);
-  }
-
-  if (current.items.has(itemId)) {
-    const r = applyCommand(current, {
-      type: "memory.retract",
-      item_id: itemId,
-      author,
-      reason,
-    });
-    current = r.state;
-    allEvents.push(...r.events);
-    retracted.push(itemId);
-  }
-
-  return { state: current, events: allEvents, retracted };
+  // Retract descendants (post-order) then the root, in a single clone with one
+  // edge index, instead of cloning the whole graph once per retracted item.
+  const items = new Map(state.items);
+  const edges = new Map(state.edges);
+  if (state.items.has(itemId)) order.push(itemId); // root retracted last
+  const { events, retracted } = retractItemsInPlace(
+    items,
+    edges,
+    order,
+    "memory.retract",
+  );
+
+  return { state: { items, edges }, events, retracted };
 }
 
 // ---------------------------------------------------------------------------

src/query.ts

@@ -423,3 +423,33 @@ export function getChildren(state: GraphState, itemId: string): MemoryItem[] {
   }
   return results;
 }
+
+/**
+ * Build a parent-id -> children index in a single pass over the graph.
+ *
+ * `getChildren` is O(items) per call; callers that need every node's children
+ * (transitive dependents, cascade retraction) would otherwise re-scan the whole
+ * graph once per node, making the walk O(nodes x items). Building the index once
+ * up front turns those walks into O(items + edges-of-the-walk).
+ */
+export function buildChildrenIndex(
+  state: GraphState,
+): Map<string, MemoryItem[]> {
+  const index = new Map<string, MemoryItem[]>();
+  for (const item of state.items.values()) {
+    if (!item.parents) continue;
+    // Dedup an item's own parent list so a child listed twice under the same
+    // parent appears once, matching getChildren's `includes` semantics.
+    const seen = item.parents.length > 1 ? new Set<string>() : null;
+    for (const pid of item.parents) {
+      if (seen) {
+        if (seen.has(pid)) continue;
+        seen.add(pid);
+      }
+      let list = index.get(pid);
+      if (!list) index.set(pid, (list = []));
+      list.push(item);
+    }
+  }
+  return index;
+}