Cache BM25 search index and entity map build

src/lib/search/retrieval.test.ts

@@ -117,6 +117,45 @@ describe('BM25 Retrieval Engine', () => {
     expect(results).toHaveLength(0)
   })
 
+  it('serves consistent results from the reference cache on repeated searches', () => {
+    const first = search(entities, claims, 'react hooks')
+    const second = search(entities, claims, 'react hooks')
+    expect(second).toEqual(first)
+    expect(second[0].id).toBe('e1')
+  })
+
+  it('rebuilds the index when entities change referentially', () => {
+    const updatedEntities = [
+      ...entities,
+      makeEntity({
+        id: 'e4',
+        name: 'React Native',
+        description: 'Mobile framework built on React',
+        tags: ['react', 'mobile'],
+        content: '# Native\nJavaScript for iOS and Android',
+      }),
+    ]
+    const results = search(updatedEntities, claims, 'react native mobile')
+    expect(results.length).toBeGreaterThan(0)
+    expect(results.some((r) => r.id === 'e4')).toBe(true)
+    // Subsequent searches on the same reference still see the new entity.
+    expect(search(updatedEntities, claims, 'native mobile').some((r) => r.id === 'e4')).toBe(true)
+  })
+
+  it('rebuilds the index when claims change referentially', () => {
+    const updatedClaims = [
+      ...claims,
+      makeClaim({
+        id: 'c3',
+        entityId: 'e1',
+        statement: 'Hooks compose cleanly for custom logic',
+      }),
+    ]
+    const results = search(entities, updatedClaims, 'custom hooks compose')
+    expect(results.length).toBeGreaterThan(0)
+    expect(results.some((r) => r.id === 'c3')).toBe(true)
+  })
+
   it('performance benchmark with large dataset', () => {
     const largeEntities: Entity[] = Array.from({ length: 500 }, (_, i) =>
       makeEntity({

src/lib/search/retrieval.ts

@@ -134,27 +134,62 @@ function getSnippet(entry: IndexEntry, maxLength: number = 140): string {
   return `${text.slice(0, maxLength)}…`
 }
 
+// Reference-based cache to avoid rebuilding index and token maps on every search query change
+let lastEntities: Entity[] | null = null
+let lastClaims: Claim[] | null = null
+let cachedEntityMap = new Map<string, Entity>()
+let cachedEntries: IndexEntry[] = []
+let cachedAvgDl = 0
+
+interface SearchIndex {
+  entityMap: Map<string, Entity>
+  entries: IndexEntry[]
+  avgDl: number
+}
+
+/** Returns the cached index when inputs are referentially unchanged, otherwise rebuilds it. */
+const getIndex = (entities: Entity[], claims: Claim[]): SearchIndex => {
+  if (entities === lastEntities && claims === lastClaims) {
+    return {
+      entityMap: cachedEntityMap,
+      entries: cachedEntries,
+      avgDl: cachedAvgDl,
+    }
+  }
+
+  const entityMap = new Map<string, Entity>()
+  for (const e of entities) {
+    entityMap.set(e.id, e)
+  }
+
+  const entries = buildIndex(entities, claims, entityMap)
+  const totalLength = entries.reduce((sum, e) => sum + e.tokenCount, 0)
+  const avgDl = entries.length > 0 ? totalLength / entries.length : 0
+
+  lastEntities = entities
+  lastClaims = claims
+  cachedEntityMap = entityMap
+  cachedEntries = entries
+  cachedAvgDl = avgDl
+
+  return { entityMap, entries, avgDl }
+}
+
 /** Run a BM25 full-text search over entities and claims. */
 export const search = (
   entities: Entity[],
   claims: Claim[],
   query: string,
   limit = 5,
 ): SearchResult[] => {
-  const entityMap = new Map<string, Entity>()
-  for (const e of entities) {
-    entityMap.set(e.id, e)
-  }
+  const { entityMap, entries, avgDl } = getIndex(entities, claims)
 
-  const entries = buildIndex(entities, claims, entityMap)
   if (entries.length === 0) return []
 
   const queryTokens = tokenize(query)
   if (queryTokens.length === 0) return []
 
   const idf = computeIDF(entries, queryTokens)
-  const totalLength = entries.reduce((sum, e) => sum + e.tokenCount, 0)
-  const avgDl = totalLength / entries.length
 
   const scored = entries
     .map((entry) => ({