fix(workers): preserve global refIndex across shards

FuseWorker now tracks per-shard global index mappings so search
results report refIndex against the full collection, not the shard.
add() appends the new global index to the chosen shard; setCollection
rebuilds the mappings in lockstep with chunked reassignment.

Worker + index mapping are colocated in a Shard object to keep the
invariant structural rather than invariant-by-convention.

Adds test/fuse-worker.test.js with a MockWorker that runs real Fuse
instances per shard, exercising search, interleaved add, and
setCollection with exact refIndex assertions.

Author: krisk

src/workers/FuseWorker.ts

@@ -19,6 +19,11 @@ interface PendingCall {
   reject: (reason: any) => void
 }
 
+interface Shard {
+  worker: Worker
+  globalIndices: number[]
+}
+
 const DEFAULT_MAX_WORKERS = 8
 
 function getDefaultWorkerCount(): number {
@@ -31,8 +36,9 @@ function getDefaultWorkerCount(): number {
 export default class FuseWorker<T> {
   private _options: IFuseOptions<T>
   private _workerOptions: FuseWorkerOptions
-  private _docs: ReadonlyArray<T>
-  private _workers: Worker[] | null = null
+  private _docs: T[]
+  private _shards: Shard[] | null = null
+  private _addCursor = 0
   private _initPromise: Promise<void> | null = null
   private _pending: Map<number, PendingCall> = new Map()
   private _nextId = 0
@@ -43,7 +49,7 @@ export default class FuseWorker<T> {
     options?: IFuseOptions<T>,
     workerOptions?: FuseWorkerOptions
   ) {
-    this._docs = docs
+    this._docs = docs.slice()
     this._options = options || {} as IFuseOptions<T>
     this._workerOptions = workerOptions || {}
     // eslint-disable-next-line @typescript-eslint/ban-ts-comment
@@ -63,37 +69,48 @@ export default class FuseWorker<T> {
     return this._initPromise
   }
 
+  private _spawnWorker(): Worker {
+    const worker = new Worker(this._workerUrl, { type: 'module' })
+
+    worker.onmessage = (e: MessageEvent) => {
+      const { id, result, error } = e.data
+      const handler = this._pending.get(id)
+      if (!handler) return
+      this._pending.delete(id)
+      if (error) {
+        handler.reject(new Error(error))
+      } else {
+        handler.resolve(result)
+      }
+    }
+
+    worker.onerror = (e: ErrorEvent) => {
+      for (const [, handler] of this._pending) {
+        handler.reject(new Error(e.message))
+      }
+    }
+
+    return worker
+  }
+
   private async _init(): Promise<void> {
     const numWorkers = this._getNumWorkers()
     const chunkSize = Math.ceil(this._docs.length / numWorkers)
-    const initPromises: Promise<void>[] = []
 
-    this._workers = []
+    this._shards = []
+    this._addCursor = 0
 
+    const initPromises: Promise<void>[] = []
     for (let i = 0; i < numWorkers; i++) {
-      const chunk = this._docs.slice(i * chunkSize, (i + 1) * chunkSize)
-      const worker = new Worker(this._workerUrl, { type: 'module' })
-
-      worker.onmessage = (e: MessageEvent) => {
-        const { id, result, error } = e.data
-        const handler = this._pending.get(id)
-        if (!handler) return
-        this._pending.delete(id)
-        if (error) {
-          handler.reject(new Error(error))
-        } else {
-          handler.resolve(result)
-        }
-      }
-
-      worker.onerror = (e: ErrorEvent) => {
-        for (const [, handler] of this._pending) {
-          handler.reject(new Error(e.message))
-        }
-      }
-
-      this._workers.push(worker)
-      initPromises.push(this._call(worker, 'init', [chunk, this._options]))
+      const start = i * chunkSize
+      const end = Math.min(start + chunkSize, this._docs.length)
+      const chunk = this._docs.slice(start, end)
+      const globalIndices: number[] = []
+      for (let j = start; j < end; j += 1) globalIndices.push(j)
+
+      const shard: Shard = { worker: this._spawnWorker(), globalIndices }
+      this._shards.push(shard)
+      initPromises.push(this._call(shard.worker, 'init', [chunk, this._options]))
     }
 
     await Promise.all(initPromises)
@@ -113,14 +130,18 @@ export default class FuseWorker<T> {
   ): Promise<FuseResult<T>[]> {
     await this._ensureInit()
 
-    const results = await Promise.all(
-      this._workers!.map((worker) => this._call(worker, 'search', [query, options]))
+    const shards = this._shards!
+    const results: FuseResult<T>[][] = await Promise.all(
+      shards.map((s) => this._call(s.worker, 'search', [query, options]))
     )
 
-    // Merge results from all shards
+    // Merge results from all shards, rewriting refIndex from shard-local to global
     const merged: FuseResult<T>[] = []
-    for (const shardResults of results) {
-      merged.push(...shardResults)
+    for (let i = 0, len = results.length; i < len; i += 1) {
+      const { globalIndices } = shards[i]
+      for (const r of results[i]) {
+        merged.push({ ...r, refIndex: globalIndices[r.refIndex] })
+      }
     }
 
     // Sort by score (lower is better)
@@ -141,35 +162,50 @@ export default class FuseWorker<T> {
   async add(doc: T): Promise<void> {
     await this._ensureInit()
 
-    // Round-robin across workers
-    const idx = this._nextId % this._workers!.length
-    await this._call(this._workers![idx], 'add', [doc])
+    const shards = this._shards!
+    const shard = shards[this._addCursor % shards.length]
+    this._addCursor += 1
+
+    const globalIdx = this._docs.length
+    this._docs.push(doc)
+    shard.globalIndices.push(globalIdx)
+
+    await this._call(shard.worker, 'add', [doc])
   }
 
   async setCollection(docs: ReadonlyArray<T>): Promise<void> {
-    this._docs = docs
-
-    if (this._workers) {
-      const numWorkers = this._workers.length
-      const chunkSize = Math.ceil(docs.length / numWorkers)
-
-      await Promise.all(
-        this._workers.map((worker, i) => {
-          const chunk = docs.slice(i * chunkSize, (i + 1) * chunkSize)
-          return this._call(worker, 'setCollection', [chunk])
-        })
-      )
-    } else {
+    this._docs = docs.slice()
+
+    if (!this._shards) {
       this._initPromise = null
+      return
+    }
+
+    const shards = this._shards
+    const chunkSize = Math.ceil(this._docs.length / shards.length)
+    this._addCursor = 0
+
+    const tasks: Promise<void>[] = []
+    for (let i = 0, len = shards.length; i < len; i += 1) {
+      const start = i * chunkSize
+      const end = Math.min(start + chunkSize, this._docs.length)
+      const chunk = this._docs.slice(start, end)
+      const globalIndices: number[] = []
+      for (let j = start; j < end; j += 1) globalIndices.push(j)
+
+      shards[i].globalIndices = globalIndices
+      tasks.push(this._call(shards[i].worker, 'setCollection', [chunk]))
     }
+
+    await Promise.all(tasks)
   }
 
   terminate(): void {
-    if (this._workers) {
-      for (const worker of this._workers) {
+    if (this._shards) {
+      for (const { worker } of this._shards) {
         worker.terminate()
       }
-      this._workers = null
+      this._shards = null
     }
     this._initPromise = null
 

test/fuse-worker.test.js

@@ -0,0 +1,191 @@
+// Unit tests for FuseWorker — mocks the Worker global so we can verify
+// the sharding/refIndex rewrite logic without spawning real workers.
+
+import { describe, test, expect, beforeAll, afterAll } from 'vitest'
+
+process.env.EXTENDED_SEARCH_ENABLED = 'true'
+process.env.TOKEN_SEARCH_ENABLED = 'true'
+
+const { default: Fuse } = await import('../src/entry')
+const { default: FuseWorker } = await import('../src/workers/FuseWorker')
+
+// Mock Worker: each instance owns a real Fuse over its assigned chunk and
+// responds to init / search / add / setCollection messages synchronously.
+class MockWorker {
+  constructor() {
+    this._fuse = null
+    this.onmessage = null
+    this.onerror = null
+    MockWorker.instances.push(this)
+  }
+
+  postMessage(msg) {
+    queueMicrotask(() => this._handle(msg))
+  }
+
+  _handle({ id, method, args }) {
+    try {
+      let result
+      switch (method) {
+        case 'init': {
+          this._fuse = new Fuse(args[0], args[1])
+          result = true
+          break
+        }
+        case 'search':
+          result = this._fuse.search(args[0], args[1])
+          break
+        case 'add':
+          this._fuse.add(args[0])
+          result = true
+          break
+        case 'setCollection':
+          this._fuse.setCollection(args[0])
+          result = true
+          break
+      }
+      this.onmessage?.({ data: { id, result } })
+    } catch (err) {
+      this.onmessage?.({ data: { id, error: err.message } })
+    }
+  }
+
+  terminate() {}
+}
+
+MockWorker.instances = []
+
+const Books = [
+  { title: 'Old Man\'s War', author: 'Scalzi' },
+  { title: 'The Lock Artist', author: 'Hamilton' },
+  { title: 'HTML5', author: 'Sharp' },
+  { title: 'A Brief History of Time', author: 'Hawking' },
+  { title: 'The Shock of the Fall', author: 'Filer' },
+  { title: 'The Great Gatsby', author: 'Fitzgerald' },
+  { title: 'The DaVinci Code', author: 'Brown' },
+  { title: 'Angels & Demons', author: 'Brown' },
+  { title: 'The Rosie Project', author: 'Simsion' }
+]
+
+describe('FuseWorker sharding', () => {
+  let originalWorker
+
+  beforeAll(() => {
+    originalWorker = globalThis.Worker
+    globalThis.Worker = MockWorker
+  })
+
+  afterAll(() => {
+    globalThis.Worker = originalWorker
+  })
+
+  test('search results carry global refIndex, not shard-local', async () => {
+    MockWorker.instances = []
+    const fw = new FuseWorker(Books, { keys: ['title', 'author'], includeScore: true }, { numWorkers: 3 })
+
+    const results = await fw.search('brown')
+
+    // Both "Brown" author entries should be found
+    const titles = results.map(r => r.item.title).sort()
+    expect(titles).toEqual(['Angels & Demons', 'The DaVinci Code'])
+
+    // refIndex must point back to Books (global)
+    for (const r of results) {
+      expect(Books[r.refIndex]).toBe(r.item)
+    }
+
+    fw.terminate()
+  })
+
+  test('add() appends globally and subsequent search returns correct refIndex', async () => {
+    MockWorker.instances = []
+    const fw = new FuseWorker(Books.slice(), { keys: ['title', 'author'] }, { numWorkers: 3 })
+
+    // Prime workers
+    await fw.search('xyz')
+
+    const newDoc = { title: 'Brown Bear', author: 'Someone' }
+    await fw.add(newDoc)
+
+    const results = await fw.search('brown')
+    const added = results.find(r => r.item.title === 'Brown Bear')
+    expect(added).toBeDefined()
+    // Added doc's refIndex should equal its global append position
+    expect(added.refIndex).toBe(Books.length)
+
+    fw.terminate()
+  })
+
+  test('interleaved search/add/setCollection keeps refIndex consistent', async () => {
+    MockWorker.instances = []
+    const initial = Books.slice(0, 4)
+    const fw = new FuseWorker(initial, { keys: ['title', 'author'] }, { numWorkers: 2 })
+
+    const globalDocs = initial.slice()
+
+    // Round 1: search before any add
+    let results = await fw.search('war')
+    for (const r of results) {
+      expect(globalDocs[r.refIndex]).toBe(r.item)
+    }
+
+    // Round 2: add two docs, then search
+    const added1 = { title: 'War Stories', author: 'Writer A' }
+    const added2 = { title: 'Peace and War', author: 'Writer B' }
+    await fw.add(added1)
+    globalDocs.push(added1)
+    await fw.add(added2)
+    globalDocs.push(added2)
+
+    results = await fw.search('war')
+    for (const r of results) {
+      expect(globalDocs[r.refIndex]).toBe(r.item)
+    }
+    // Specifically, the added docs must map to their append positions
+    const a1 = results.find(r => r.item === added1)
+    const a2 = results.find(r => r.item === added2)
+    expect(a1?.refIndex).toBe(initial.length)
+    expect(a2?.refIndex).toBe(initial.length + 1)
+
+    // Round 3: setCollection resets, then add and search again
+    const reset = Books.slice(4, 7)
+    await fw.setCollection(reset)
+    const globalDocs2 = reset.slice()
+
+    const added3 = { title: 'Brown Ale', author: 'Brewer' }
+    await fw.add(added3)
+    globalDocs2.push(added3)
+
+    results = await fw.search('brown')
+    for (const r of results) {
+      expect(globalDocs2[r.refIndex]).toBe(r.item)
+    }
+    const a3 = results.find(r => r.item === added3)
+    expect(a3?.refIndex).toBe(reset.length)
+
+    fw.terminate()
+  })
+
+  test('setCollection rebuilds mapping so refIndex matches new collection', async () => {
+    MockWorker.instances = []
+    const fw = new FuseWorker(Books, { keys: ['title', 'author'] }, { numWorkers: 2 })
+    await fw.search('xyz') // force init
+
+    const newDocs = [
+      { title: 'Alpha', author: 'A' },
+      { title: 'Bravo', author: 'B' },
+      { title: 'Charlie Brown', author: 'C' },
+      { title: 'Delta', author: 'D' },
+      { title: 'Brown Fox', author: 'E' }
+    ]
+
+    await fw.setCollection(newDocs)
+
+    const results = await fw.search('brown')
+    for (const r of results) {
+      expect(newDocs[r.refIndex]).toBe(r.item)
+    }
+
+    fw.terminate()
+  })
+})