diff --git a/packages/memory-graph/scripts/bench-cluster-assignments-v8.mjs b/packages/memory-graph/scripts/bench-cluster-assignments-v8.mjs new file mode 100644 index 00000000..ff6b4ea1 --- /dev/null +++ b/packages/memory-graph/scripts/bench-cluster-assignments-v8.mjs @@ -0,0 +1,305 @@ +// Companion to bench-cluster-assignments.ts, for the V8 engine specifically +// (Node, and what Chrome/Edge/most browsers actually run -- i.e. where this +// component executes for real users). +// +// scripts/bench-cluster-assignments.ts imports the live computeClusterAssignments +// and is the benchmark to trust for correctness (it runs the real, current +// source). But run under Bun -- this repo's own dev/test runtime -- it shows +// close to NO difference between the old and new implementation. That is +// not a flaw in the fix: Bun's JavaScriptCore engine optimizes Array.shift() +// far better than V8 does, so the O(n^2) behavior this fix removes barely +// shows up there. Do not conclude from the .ts benchmark alone that this +// optimization is a no-op -- run this file too, with plain `node`. +// +// This script is a standalone, verbatim copy of both implementations (not +// an import), specifically so it can run under plain `node` without hitting +// Node's strict ESM extension-resolution rules on the rest of the source +// tree, and without adding a new dev dependency (e.g. tsx) just to make +// that import work. If computeClusterAssignments changes again, this file's +// copies should be refreshed to match. +// +// Usage (plain Node, not bun): +// node scripts/bench-cluster-assignments-v8.mjs + +function hashString(value) { + let hash = 0 + for (let i = 0; i < value.length; i++) { + hash = (Math.imul(31, hash) + value.charCodeAt(i)) | 0 + } + return hash >>> 0 +} + +const CLUSTER_COLORS = [ + "#58C7E8", + "#E7BC52", + "#74D680", + "#D47B75", + "#A789E8", + "#62C5A8", + "#74ABD8", + "#C78AC8", + "#D18A58", + "#8BCB6F", +] + +function getClusterColor(key) { + return CLUSTER_COLORS[hashString(key) % CLUSTER_COLORS.length] +} + +function ensureAdjacency(map, id) { + if (!map.has(id)) map.set(id, new Set()) +} + +function connect(map, a, b) { + ensureAdjacency(map, a) + ensureAdjacency(map, b) + map.get(a)?.add(b) + map.get(b)?.add(a) +} + +function getMemoryRelationTargets(mem) { + if ( + mem.memoryRelations && + typeof mem.memoryRelations === "object" && + Object.keys(mem.memoryRelations).length > 0 + ) { + return mem.memoryRelations + } + if (mem.parentMemoryId) return { [mem.parentMemoryId]: "updates" } + return {} +} + +// Verbatim (pre-fix): BFS frontier drained with Array.shift() -- O(remaining +// length) per call. +function previousImplementation(documents) { + const adjacency = new Map() + const docByMemory = new Map() + const orderByMemory = new Map() + const allMemoryIds = new Set() + let order = 0 + + for (const doc of documents) { + let firstMemoryId = null + for (const mem of doc.memories) { + allMemoryIds.add(mem.id) + docByMemory.set(mem.id, doc.id) + orderByMemory.set(mem.id, order++) + ensureAdjacency(adjacency, mem.id) + if (!firstMemoryId) firstMemoryId = mem.id + else connect(adjacency, firstMemoryId, mem.id) + } + } + + for (const doc of documents) { + for (const mem of doc.memories) { + for (const targetId of Object.keys(getMemoryRelationTargets(mem))) { + if (!allMemoryIds.has(targetId)) continue + connect(adjacency, mem.id, targetId) + } + } + } + + const assignments = new Map() + const visited = new Set() + const memoryIdsByOrder = [...allMemoryIds].sort( + (a, b) => (orderByMemory.get(a) ?? 0) - (orderByMemory.get(b) ?? 0), + ) + + for (const startId of memoryIdsByOrder) { + if (visited.has(startId)) continue + const component = [] + const queue = [startId] + visited.add(startId) + + while (queue.length > 0) { + const id = queue.shift() + component.push(id) + for (const nextId of adjacency.get(id) ?? []) { + if (visited.has(nextId)) continue + visited.add(nextId) + queue.push(nextId) + } + } + + component.sort( + (a, b) => (orderByMemory.get(a) ?? 0) - (orderByMemory.get(b) ?? 0), + ) + const firstId = component[0] ?? startId + const docIds = new Set(component.map((id) => docByMemory.get(id))) + const firstDocId = docByMemory.get(firstId) ?? "unknown" + const key = + docIds.size <= 1 + ? `doc:${firstDocId}` + : `relation:${firstDocId}:${firstId}` + const assignment = { + key, + color: getClusterColor(key), + size: component.length, + } + for (const id of component) assignments.set(id, assignment) + } + + return assignments +} + +// Optimized (current): index-pointer dequeue -- O(1) per call. Identical +// otherwise -- this is the actual diff applied to use-graph-data.ts. +function optimizedImplementation(documents) { + const adjacency = new Map() + const docByMemory = new Map() + const orderByMemory = new Map() + const allMemoryIds = new Set() + let order = 0 + + for (const doc of documents) { + let firstMemoryId = null + for (const mem of doc.memories) { + allMemoryIds.add(mem.id) + docByMemory.set(mem.id, doc.id) + orderByMemory.set(mem.id, order++) + ensureAdjacency(adjacency, mem.id) + if (!firstMemoryId) firstMemoryId = mem.id + else connect(adjacency, firstMemoryId, mem.id) + } + } + + for (const doc of documents) { + for (const mem of doc.memories) { + for (const targetId of Object.keys(getMemoryRelationTargets(mem))) { + if (!allMemoryIds.has(targetId)) continue + connect(adjacency, mem.id, targetId) + } + } + } + + const assignments = new Map() + const visited = new Set() + const memoryIdsByOrder = [...allMemoryIds].sort( + (a, b) => (orderByMemory.get(a) ?? 0) - (orderByMemory.get(b) ?? 0), + ) + + for (const startId of memoryIdsByOrder) { + if (visited.has(startId)) continue + const component = [] + const queue = [startId] + let head = 0 + visited.add(startId) + + while (head < queue.length) { + const id = queue[head++] + component.push(id) + for (const nextId of adjacency.get(id) ?? []) { + if (visited.has(nextId)) continue + visited.add(nextId) + queue.push(nextId) + } + } + + component.sort( + (a, b) => (orderByMemory.get(a) ?? 0) - (orderByMemory.get(b) ?? 0), + ) + const firstId = component[0] ?? startId + const docIds = new Set(component.map((id) => docByMemory.get(id))) + const firstDocId = docByMemory.get(firstId) ?? "unknown" + const key = + docIds.size <= 1 + ? `doc:${firstDocId}` + : `relation:${firstDocId}:${firstId}` + const assignment = { + key, + color: getClusterColor(key), + size: component.length, + } + for (const id of component) assignments.set(id, assignment) + } + + return assignments +} + +function makeMemory(id, relations) { + return { id, memoryRelations: relations ?? null, parentMemoryId: null } +} + +/** + * One hub memory + n-1 memories that `derives` from it, one per document -- + * a single connected component with a wide BFS frontier, matching the + * cross-document relation-merge path in use-graph-data.ts. + */ +function buildStarDataset(n) { + const memories = [makeMemory("mem-hub")] + for (let i = 1; i < n; i++) { + memories.push(makeMemory(`mem-${i}`, { "mem-hub": "derives" })) + } + return memories.map((mem, i) => ({ id: `doc-${i}`, memories: [mem] })) +} + +function timeOnce(fn) { + const start = performance.now() + fn() + return performance.now() - start +} + +console.log(`Engine: ${process.release?.name ?? "unknown"} ${process.version}`) +if (process.versions?.bun) { + console.log( + "WARNING: running under Bun. This is the wrong engine to judge this fix by -- run with plain `node` instead.", + ) +} + +console.log( + "\nCorrectness check: previous vs. optimized produce identical assignments", +) +{ + const dataset = buildStarDataset(2_000) + const prev = previousImplementation(dataset) + const curr = optimizedImplementation(dataset) + let mismatch = prev.size !== curr.size + if (!mismatch) { + for (const [id, assignment] of prev) { + const currAssignment = curr.get(id) + if ( + !currAssignment || + currAssignment.key !== assignment.key || + currAssignment.color !== assignment.color || + currAssignment.size !== assignment.size + ) { + mismatch = true + break + } + } + } + console.log( + mismatch + ? " MISMATCH -- do not trust these numbers" + : " OK, outputs are identical", + ) + if (mismatch) process.exit(1) +} + +const RUNS_PER_SIZE = 5 +const SIZES = [5_000, 10_000, 20_000, 40_000, 80_000, 100_000] + +console.log("\nn\tprevious (ms)\toptimized (ms)\tspeedup") +for (const n of SIZES) { + const dataset = buildStarDataset(n) + let prevBest = Number.POSITIVE_INFINITY + let currBest = Number.POSITIVE_INFINITY + for (let i = 0; i < RUNS_PER_SIZE; i++) { + let pTime + let cTime + // Alternate which implementation runs first each sample, to cancel + // out any heap-warmup/ordering bias between the two. + if (i % 2 === 0) { + pTime = timeOnce(() => previousImplementation(dataset)) + cTime = timeOnce(() => optimizedImplementation(dataset)) + } else { + cTime = timeOnce(() => optimizedImplementation(dataset)) + pTime = timeOnce(() => previousImplementation(dataset)) + } + if (pTime < prevBest) prevBest = pTime + if (cTime < currBest) currBest = cTime + } + console.log( + `${n}\t${prevBest.toFixed(2)}\t\t${currBest.toFixed(2)}\t\t${(prevBest / currBest).toFixed(2)}x`, + ) +} diff --git a/packages/memory-graph/scripts/bench-cluster-assignments.ts b/packages/memory-graph/scripts/bench-cluster-assignments.ts new file mode 100644 index 00000000..0a326310 --- /dev/null +++ b/packages/memory-graph/scripts/bench-cluster-assignments.ts @@ -0,0 +1,304 @@ +/** + * Benchmark for computeClusterAssignments (src/hooks/use-graph-data.ts). + * + * computeClusterAssignments runs inside a useMemo keyed on the full + * `documents` array, so it re-executes on the whole dataset every time the + * memory graph loads or its data changes, on the browser main thread. + * + * Its BFS drains the frontier with `queue.shift()`, which is O(n) per call. + * For a graph shaped like a wide star -- one "hub" memory that many other + * memories `derives`/`updates` from (e.g. a canonical/root memory + * referenced by a long history of later updates) -- the BFS frontier grows + * to ~n before draining, making the whole traversal O(n^2). + * + * This script measures wall-clock time across a range of dataset sizes for + * both the live implementation and a frozen pre-fix snapshot + * (previousImplementation, below). Scaling behavior (does cost grow + * linearly or quadratically with n?) is the signal we actually care about, + * and it is far more robust to GC/JIT noise than a single head-to-head + * timing at one size -- an earlier attempt at this benchmark using vitest's + * `bench` reported the *current, unmodified* implementation as "1.2x faster + * than itself" purely from GC noise on 40k-node allocations, which is why + * this script uses best-of-N sampling at multiple sizes instead. + * + * IMPORTANT: this repo's dev tooling runs on Bun, whose JavaScriptCore + * engine optimizes Array.shift() far better than V8 (what Chrome/Edge/most + * browsers -- i.e. this component's actual users -- run). Run under Bun, + * this benchmark will show close to NO difference between old and new. That + * does not mean the fix is a no-op -- see the companion script + * bench-cluster-assignments-v8.mjs, which is engine-independent (plain + * Node, zero dependencies) and shows the real, growing gap. + * + * Usage: + * bun run scripts/bench-cluster-assignments.ts (correctness + Bun numbers) + * node scripts/bench-cluster-assignments-v8.mjs (V8/real-world numbers) + */ +import { computeClusterAssignments } from "../src/hooks/use-graph-data" +import type { GraphApiDocument, GraphApiMemory } from "../src/types" + +function makeMemory( + id: string, + relations?: Record, +): GraphApiMemory { + return { + id, + memory: "test", + isStatic: false, + spaceId: "default", + isLatest: true, + isForgotten: false, + forgetAfter: null, + forgetReason: null, + version: 1, + parentMemoryId: null, + rootMemoryId: null, + createdAt: "2024-01-01", + updatedAt: "2024-01-01", + memoryRelations: relations ?? null, + } +} + +function makeDocument( + id: string, + memories: GraphApiMemory[], +): GraphApiDocument { + return { + id, + title: id, + summary: null, + documentType: "text", + createdAt: "2024-01-01", + updatedAt: "2024-01-01", + memories, + } +} + +/** + * One hub memory + n-1 memories that `derives` from it, spread one-per- + * document. Exercises the cross-document relation-merge path at + * use-graph-data.ts:140-147 and produces a single large connected component + * with a wide BFS frontier -- the shape that triggers the O(n^2) behavior. + */ +function buildStarDataset(n: number): GraphApiDocument[] { + const memories: GraphApiMemory[] = [makeMemory("mem-hub")] + for (let i = 1; i < n; i++) { + memories.push(makeMemory(`mem-${i}`, { "mem-hub": "derives" })) + } + return memories.map((mem, i) => makeDocument(`doc-${i}`, [mem])) +} + +// --- frozen pre-optimization snapshot (benchmark comparison only) --- +// Verbatim copy of computeClusterAssignments and its private helpers as +// they existed before the perf/cluster-bfs-queue fix. Kept here only so +// this benchmark keeps comparing old vs. new after the production code is +// optimized. + +function hashStringSnapshot(value: string): number { + let hash = 0 + for (let i = 0; i < value.length; i++) { + hash = (Math.imul(31, hash) + value.charCodeAt(i)) | 0 + } + return hash >>> 0 +} + +const CLUSTER_COLORS_SNAPSHOT = [ + "#58C7E8", + "#E7BC52", + "#74D680", + "#D47B75", + "#A789E8", + "#62C5A8", + "#74ABD8", + "#C78AC8", + "#D18A58", + "#8BCB6F", +] + +function getClusterColorSnapshot(key: string): string { + return CLUSTER_COLORS_SNAPSHOT[ + hashStringSnapshot(key) % CLUSTER_COLORS_SNAPSHOT.length + ] as string +} + +function ensureAdjacencySnapshot(map: Map>, id: string) { + if (!map.has(id)) map.set(id, new Set()) +} + +function connectSnapshot(map: Map>, a: string, b: string) { + ensureAdjacencySnapshot(map, a) + ensureAdjacencySnapshot(map, b) + map.get(a)?.add(b) + map.get(b)?.add(a) +} + +function getMemoryRelationTargetsSnapshot( + mem: GraphApiMemory, +): Record { + if ( + mem.memoryRelations && + typeof mem.memoryRelations === "object" && + Object.keys(mem.memoryRelations).length > 0 + ) { + return mem.memoryRelations + } + if (mem.parentMemoryId) return { [mem.parentMemoryId]: "updates" } + return {} +} + +function previousImplementation(documents: GraphApiDocument[]) { + const adjacency = new Map>() + const docByMemory = new Map() + const orderByMemory = new Map() + const allMemoryIds = new Set() + let order = 0 + + for (const doc of documents) { + let firstMemoryId: string | null = null + for (const mem of doc.memories) { + allMemoryIds.add(mem.id) + docByMemory.set(mem.id, doc.id) + orderByMemory.set(mem.id, order++) + ensureAdjacencySnapshot(adjacency, mem.id) + + if (!firstMemoryId) { + firstMemoryId = mem.id + } else { + connectSnapshot(adjacency, firstMemoryId, mem.id) + } + } + } + + for (const doc of documents) { + for (const mem of doc.memories) { + for (const targetId of Object.keys( + getMemoryRelationTargetsSnapshot(mem), + )) { + if (!allMemoryIds.has(targetId)) continue + connectSnapshot(adjacency, mem.id, targetId) + } + } + } + + const assignments = new Map() + const visited = new Set() + const memoryIdsByOrder = [...allMemoryIds].sort( + (a, b) => (orderByMemory.get(a) ?? 0) - (orderByMemory.get(b) ?? 0), + ) + + for (const startId of memoryIdsByOrder) { + if (visited.has(startId)) continue + + const component: string[] = [] + const queue = [startId] + visited.add(startId) + + while (queue.length > 0) { + const id = queue.shift() as string + component.push(id) + for (const nextId of adjacency.get(id) ?? []) { + if (visited.has(nextId)) continue + visited.add(nextId) + queue.push(nextId) + } + } + + component.sort( + (a, b) => (orderByMemory.get(a) ?? 0) - (orderByMemory.get(b) ?? 0), + ) + const firstId = component[0] ?? startId + const docIds = new Set(component.map((id) => docByMemory.get(id))) + const firstDocId = docByMemory.get(firstId) ?? "unknown" + const key = + docIds.size <= 1 + ? `doc:${firstDocId}` + : `relation:${firstDocId}:${firstId}` + const assignment = { + key, + color: getClusterColorSnapshot(key), + size: component.length, + } + + for (const id of component) assignments.set(id, assignment) + } + + return assignments +} + +// --- timing harness --- + +function bestOf(fn: () => void, runs: number): number { + let best = Number.POSITIVE_INFINITY + for (let i = 0; i < runs; i++) { + const start = performance.now() + fn() + const elapsed = performance.now() - start + if (elapsed < best) best = elapsed + } + return best +} + +const RUNS_PER_SIZE = 5 +const SIZES = [5_000, 10_000, 20_000, 40_000, 80_000, 100_000] + +if ((process.versions as { bun?: string }).bun) { + console.log( + "WARNING: running under Bun -- its JavaScriptCore engine optimizes\n" + + "Array.shift() well, so the numbers below will look flat regardless\n" + + "of the fix. Run `node scripts/bench-cluster-assignments-v8.mjs` for\n" + + "the engine-independent, real-world (V8) comparison.\n", + ) +} + +console.log( + "Correctness check: previous vs. current produce identical assignments", +) +{ + const dataset = buildStarDataset(2_000) + const prev = previousImplementation(dataset) + const curr = computeClusterAssignments(dataset) + let mismatch = false + if (prev.size !== curr.size) mismatch = true + for (const [id, assignment] of prev) { + const currAssignment = curr.get(id) + if ( + !currAssignment || + currAssignment.key !== assignment.key || + currAssignment.color !== assignment.color || + currAssignment.size !== assignment.size + ) { + mismatch = true + break + } + } + console.log( + mismatch + ? " MISMATCH -- do not trust these numbers" + : " OK, outputs are identical", + ) + if (mismatch) process.exit(1) +} + +console.log("\nn\tprevious (ms)\tcurrent (ms)\tspeedup") +for (const n of SIZES) { + const dataset = buildStarDataset(n) + // Alternate which implementation runs first across the repeated samples + // to cancel out any ordering/heap-warmup bias between the two. + let prevBest = Number.POSITIVE_INFINITY + let currBest = Number.POSITIVE_INFINITY + for (let i = 0; i < RUNS_PER_SIZE; i++) { + let pTime: number + let cTime: number + if (i % 2 === 0) { + pTime = bestOf(() => previousImplementation(dataset), 1) + cTime = bestOf(() => computeClusterAssignments(dataset), 1) + } else { + cTime = bestOf(() => computeClusterAssignments(dataset), 1) + pTime = bestOf(() => previousImplementation(dataset), 1) + } + if (pTime < prevBest) prevBest = pTime + if (cTime < currBest) currBest = cTime + } + console.log( + `${n}\t${prevBest.toFixed(2)}\t\t${currBest.toFixed(2)}\t\t${(prevBest / currBest).toFixed(2)}x`, + ) +} diff --git a/packages/memory-graph/src/__tests__/graph-data-utils.test.ts b/packages/memory-graph/src/__tests__/graph-data-utils.test.ts index f98df074..dffbd454 100644 --- a/packages/memory-graph/src/__tests__/graph-data-utils.test.ts +++ b/packages/memory-graph/src/__tests__/graph-data-utils.test.ts @@ -140,6 +140,27 @@ describe("cluster assignments", () => { expect(assignments.get("a1")?.key).toBe(assignments.get("b1")?.key) }) + + it("merges a wide fan-out of memories relating to one hub into a single cluster", () => { + // Regression test for the BFS queue drain: a hub with many direct + // relations produces a wide frontier, which previously interacted + // badly with an O(n) `Array.shift()` dequeue. + const hub = makeDocument("doc-hub", [makeMemory({ id: "hub" })]) + const spokes = Array.from({ length: 200 }, (_, i) => + makeDocument(`doc-${i}`, [ + makeMemory({ id: `spoke-${i}`, memoryRelations: { hub: "derives" } }), + ]), + ) + + const assignments = computeClusterAssignments([hub, ...spokes]) + + const hubKey = assignments.get("hub")?.key + expect(hubKey).toBeDefined() + for (let i = 0; i < spokes.length; i++) { + expect(assignments.get(`spoke-${i}`)?.key).toBe(hubKey) + } + expect(assignments.get("hub")?.size).toBe(spokes.length + 1) + }) }) describe("memory orbit placement", () => { diff --git a/packages/memory-graph/src/hooks/use-graph-data.ts b/packages/memory-graph/src/hooks/use-graph-data.ts index e0a3b1b3..3e9b0229 100644 --- a/packages/memory-graph/src/hooks/use-graph-data.ts +++ b/packages/memory-graph/src/hooks/use-graph-data.ts @@ -157,10 +157,15 @@ export function computeClusterAssignments( const component: string[] = [] const queue = [startId] + // Index pointer instead of Array.shift(): shift() is O(remaining + // length) per call, so draining a wide BFS frontier (e.g. a heavily + // cross-referenced "hub" memory with many direct relations) was + // O(n^2) for large components. + let head = 0 visited.add(startId) - while (queue.length > 0) { - const id = queue.shift() as string + while (head < queue.length) { + const id = queue[head++] as string component.push(id) for (const nextId of adjacency.get(id) ?? []) { if (visited.has(nextId)) continue