From 89bf26ba9dc41e49c0e9f7f7d7e6520db92860d8 Mon Sep 17 00:00:00 2001 From: codewarnab Date: Mon, 19 Oct 2026 07:54:11 +0000 Subject: [PATCH 1/2] test: cover emoji and maximum query browser parity --- docs/IMPROVEMENT-PLAN.md | 1 + scripts/test-regression.ts | 29 ++++++++++++++++++++++++++--- 2 files changed, 27 insertions(+), 3 deletions(-) diff --git a/docs/IMPROVEMENT-PLAN.md b/docs/IMPROVEMENT-PLAN.md index 201483c..de0a6f2 100644 --- a/docs/IMPROVEMENT-PLAN.md +++ b/docs/IMPROVEMENT-PLAN.md @@ -25,6 +25,7 @@ Low-risk housekeeping that clears the way for the bigger phases. There are about 300 tests, but several core modules have none of their own. This phase covers the guarantees that 1.0 commits to. ### 2.1 Tests that GPU and CPU return the same scores +- **Browser coverage extended:** the existing `scripts/test-regression.ts` differential gate now includes joined emoji and component queries, plus 128-token ASCII and case-fold-expanded queries in substring and fuzzy modes. Matching records and last-token near misses exercise the actual WGSL path against the CPU reference. CI already runs this gate on every PR. - Add `test/scoring-parity.test.ts`. It runs a fixed set of queries through the mock WebGPU path and through `cpu-reference.ts` / `exact-scorer.ts`. - Check that both return the same `SearchResultItem[]`: same order, same integer scores, scores in descending order. - Cover ASCII, CJK, emoji/ZWJ, mixed case, and queries at the maximum length. diff --git a/scripts/test-regression.ts b/scripts/test-regression.ts index 41c7d85..649d6f4 100644 --- a/scripts/test-regression.ts +++ b/scripts/test-regression.ts @@ -1,7 +1,7 @@ import puppeteer from 'puppeteer-core'; import path from 'path'; import fs from 'fs'; -import { normalizeText, scoreExactMatches } from '../packages/webgpu-search/src/index'; +import { normalizeText, scoreExactMatches, QUERY_TOKENS_MAX } from '../packages/webgpu-search/src/index'; import { getChromeExecutablePath, getChromeLaunchArgs, @@ -180,6 +180,12 @@ async function main() { // agree on NFC for this ASCII-heavy corpus (negligible risk here; // cross-runtime NFC is probed/reported, not gated). const FCP = String.fromCodePoint; + const joinedEmoji = FCP(0x1f469, 0x200d, 0x1f4bb); + // A distinct final token catches dispatch/scoring paths that truncate + // the query at 127 tokens. The folded query reaches the same cap with + // only 64 input characters, so the limit is measured after folding. + const maxQuery = 'a'.repeat(QUERY_TOKENS_MAX - 1) + 'z'; + const foldedMaxQuery = FCP(0xdf).repeat(QUERY_TOKENS_MAX / 2); const parityStrings: string[] = []; for (let i = 0; i < 200; i++) { parityStrings.push(`src/services/Auth${i % 7 === 0 ? FCP(0xdf) : 's'}Controller_${i}.ts`); @@ -189,7 +195,16 @@ async function main() { FCP(0x5317) + FCP(0x4eac) + 'city', FCP(0x1f600) + 'smile', 'caf' + FCP(0xe9), - 'plain' + 'plain', + joinedEmoji, + 'Dev ' + joinedEmoji + ' at work', + FCP(0x1f469) + ' ' + FCP(0x1f4bb), // Separate scalars are not the ZWJ sequence. + maxQuery, + 'prefix ' + maxQuery + ' suffix', + 'a'.repeat(QUERY_TOKENS_MAX - 1) + 'x', // Differs only at the last token. + foldedMaxQuery, + 'ss'.repeat(QUERY_TOKENS_MAX / 2), + 'ss'.repeat(QUERY_TOKENS_MAX / 2 - 1) + 'sx' ); const parityCells = [ { query: 'auth', mode: 'substring', limit: 50 }, @@ -200,6 +215,14 @@ async function main() { { query: 'SS', mode: 'substring', limit: 50 }, { query: FCP(0x5317), mode: 'substring', limit: 10 }, { query: FCP(0x1f600), mode: 'substring', limit: 10 }, + { query: joinedEmoji, mode: 'substring', limit: 10 }, + { query: joinedEmoji, mode: 'fuzzy', limit: 10 }, + { query: FCP(0x1f4bb), mode: 'substring', limit: 10 }, + { query: FCP(0x1f4bb), mode: 'fuzzy', limit: 10 }, + { query: maxQuery, mode: 'substring', limit: 10 }, + { query: maxQuery, mode: 'fuzzy', limit: 10 }, + { query: foldedMaxQuery, mode: 'substring', limit: 10 }, + { query: foldedMaxQuery, mode: 'fuzzy', limit: 10 }, { query: 'zzz-no-match', mode: 'substring', limit: 50 }, { query: 'zzz-no-match', mode: 'fuzzy', limit: 50 }, ] as Array<{ query: string; mode: 'substring' | 'fuzzy'; limit: number }>; @@ -258,7 +281,7 @@ async function main() { } } (testResults as any).results.push({ - name: `parity ${cell.mode} ${JSON.stringify(cell.query.slice(0, 12))} (ordered index/score/text)`, + name: `parity ${cell.mode} ${JSON.stringify(cell.query.slice(0, 12))} (${qT.length} tokens; ordered index/score/text)`, passed: pass, details: { totalMatches: got.totalMatches, candidateCount: got.candidateCount, expected: expected.totalMatches } }); From 83b0e511e440a205aaf836c7fbcb14b8e6d7476b Mon Sep 17 00:00:00 2001 From: codewarnab Date: Wed, 21 Oct 2026 19:26:37 +0000 Subject: [PATCH 2/2] refactor: extract document and GPU engine modules --- packages/webgpu-search/src/document-index.ts | 3175 +---------------- packages/webgpu-search/src/document/build.ts | 1147 ++++++ .../src/document/filter-facet.ts | 140 + packages/webgpu-search/src/document/query.ts | 1513 ++++++++ .../webgpu-search/src/document/snapshot.ts | 522 +++ packages/webgpu-search/src/document/state.ts | 106 + packages/webgpu-search/src/gpu/buffers.ts | 276 ++ packages/webgpu-search/src/gpu/dispatch.ts | 226 ++ packages/webgpu-search/src/gpu/pipelines.ts | 62 + packages/webgpu-search/src/gpu/state.ts | 74 + packages/webgpu-search/src/webgpu-engine.ts | 520 +-- .../test/engine-pool-regressions.test.ts | 40 +- scripts/check-public-api.ts | 48 +- scripts/test-document-records.ts | 148 +- scripts/test-faceting.ts | 172 +- scripts/test-parity-harness.ts | 240 +- scripts/test-vgpu-mock.ts | 37 +- 17 files changed, 4724 insertions(+), 3722 deletions(-) create mode 100644 packages/webgpu-search/src/document/build.ts create mode 100644 packages/webgpu-search/src/document/filter-facet.ts create mode 100644 packages/webgpu-search/src/document/query.ts create mode 100644 packages/webgpu-search/src/document/snapshot.ts create mode 100644 packages/webgpu-search/src/document/state.ts create mode 100644 packages/webgpu-search/src/gpu/buffers.ts create mode 100644 packages/webgpu-search/src/gpu/dispatch.ts create mode 100644 packages/webgpu-search/src/gpu/pipelines.ts create mode 100644 packages/webgpu-search/src/gpu/state.ts diff --git a/packages/webgpu-search/src/document-index.ts b/packages/webgpu-search/src/document-index.ts index 6cfb517..91c776c 100644 --- a/packages/webgpu-search/src/document-index.ts +++ b/packages/webgpu-search/src/document-index.ts @@ -1,104 +1,51 @@ -import { WebGPUEngine } from './webgpu-engine'; -import { CPUEngine } from './cpu-engine'; -import { GpuDevicePool, assertValidPowerPreference } from './gpu-device-pool'; -import { packDataset, checkMemoryBudget } from './dataset-packing'; -import { normalizeText } from './text-normalization'; -import { - initialEngineState, - transitionEngineState, - type EngineState, -} from './engine-state'; -import { - scoreExactMatchesMultiField, - scoreFuzzyTokens, - scoreSubstringTypoTokens, - type FieldScoreDefinition, - type CpuModeOptions -} from './exact-scorer'; -import { - normalizeTypoTolerance, - type NormalizedTypoOptions -} from './search/typo-tolerance'; -import { - normalizeTokenMatchOptions, - type NormalizedTokenMatchOptions -} from './search/token-search'; -import { - normalizePrefixOptions, - assertPrefixLengthForQuery, - scorePrefixTokens, - type NormalizedPrefixOptions -} from './search/prefix-search'; -import { - compareRanked, - isExactTokenMatch, - normalizeTieBreakers, - type RankableCandidate -} from './ranking'; -import { - aggregateDocMatches as aggregateDocMatchesPure, - accumulateDocMatch, -} from './document-aggregate'; -import { - normalizeAutocompleteOptions, - type NormalizedAutocompleteOptions -} from './autocomplete'; -import { - normalizeSearchHooks, - resolveEffectiveHooks, - getTokenTermsForQuery, - applyScoringHook, - applyPostProcess, - assertHooksSatisfied -} from './hooks'; -import { - clampLimit, - throwIfAborted, - abortError, - nowMs, - assertValidMode -} from './guard'; -import { - normalizeCostBudgetOptions, - throwIfBudgetAborted, - assertTimeBudget, - assertCandidateBudget, - computeFilterSelectivity, - isBroadQueryHeuristic, - isBroadSelectivity, - broadQueryRouteWarning, - broadSelectivityWarning, - candidateOverflowWarning -} from './diagnostics'; +/** + * High-level multi-field document search index with dynamic mutations and highlighting. + * + * This file is the public `DocumentIndex` facade. It owns: + * - construction-time fail-closed validation (fields, removed `extensions`, + * power preference, text profile) and state construction, + * - lifecycle gating (destroyed checks, mutation mutex scheduling, the + * search-generation snapshot), + * - the public getter/serializer surface, + * - snapshot restore orchestration (`restore`, `fromSnapshot`, + * `fromSnapshotData`). + * + * All implementation lives in the `document/` sibling modules operating on + * the internal `DocumentIndexState` bundle: + * - `document/build.ts` : ingestion, batch mutations, GPU bootstrap, + * compaction, destroy. + * - `document/query.ts` : search validation/execution + autocomplete. + * - `document/filter-facet.ts`: filter compilation + facet aggregation. + * - `document/snapshot.ts` : snapshot merge/apply + restore guards. + */ + +import { assertValidPowerPreference } from './gpu-device-pool'; import { - QUERY_TOKENS_MAX, + IncompatibleOptionError, + ProfileMismatchError, SCORING_VERSION, SNAPSHOT_FORMAT_VERSION, UNICODE_VERSION, - normalizeCpuScorer, - DuplicateIdError, - DocumentNotFoundError, - IncompatibleIndexError, - IncompatibleOptionError, - ProfileMismatchError, - QueryTooLongError, - type TextProfileId + type TextProfileId, } from './text-profile'; -import { DocumentBitset } from './filtering/doc-bitset'; -import { ColumnarStore, readOwnField, type PreparedColumnarRow } from './filtering/columnar-store'; -import { compileFilter } from './filtering/compile-filter'; +import { nowMs, throwIfAborted } from './guard'; import { - FacetEngine, - excludeFieldFromFilter, - normalizeFacetRequests, - type NormalizedFacet -} from './faceting/facet-engine'; -import { alignHighlights, renderHighlightedText } from './highlight'; + applyBatchOps, + createDocumentIndexState, + compactSync, + destroyIndex, + enqueue, + initializeDocumentIndex, + rebuildGpuEngine, + compactCpu, +} from './document/build'; +import { runAutocomplete, runSearch } from './document/query'; +import { applySnapshotData, mergeSnapshotRestoreOptions } from './document/snapshot'; import { decodeSnapshot, - restoreSnapshot, encodeSnapshot, - type RestoredDocumentSnapshot + restoreSnapshot, + type RestoredDocumentSnapshot, } from './snapshot-codec'; import type { AddOptions, @@ -107,123 +54,28 @@ import type { DocumentIndexStats, DocumentSearchOptions, DocumentSearchResponse, - DocumentSearchResultItem, - EngineType, - HighlightRange, + FilterFieldDefinition, MutationBatch, MutationResult, RestoreDocumentIndexOptions, SearchHooks, - SearchMode, - SearchTimings, SerializeDocumentIndexOptions, AutocompleteOptions, - SuggestionItem, SuggestResponse, - TieBreakerCriterion, - DocumentIndexSchema, - FacetResult, - FilterExpression, - FilterFieldDefinition, - FilterFieldType, - QueryDiagnostics, - QueryDiagnosticsTimings } from './types'; +import type { ColumnarStore } from './filtering/columnar-store'; +import type { DocumentIndexState, InternalField } from './document/state'; -/** - * Own-property write that is safe for arbitrary user field names: a plain - * `obj[key] = v` with key `"__proto__"` would replace the prototype instead - * of creating an entry. - */ -function setOwnKey(obj: Record, key: string, value: V): void { - Object.defineProperty(obj, key, { value, enumerable: true, configurable: true, writable: true }); -} - -export interface InternalField extends FieldScoreDefinition { - name: string; - weight: number; - getter: (doc: TDoc) => string | string[] | undefined | null; - originalIndex: number; -} - -/** - * fail-closed custom filter-getter restore guard (mirrors the - * hookIds idiom). Snapshots recording `hasGetter: true` for a filter field - * require the caller to supply a matching getter override via - * `options.options.filterFields`; otherwise columnar rebuild via default - * `doc[name]` would silently change filter semantics. - */ -export function assertSnapshotFilterGettersSatisfied( - snapshotFilterFields: Array<{ name: string; hasGetter?: boolean }> | undefined, - userFilterFields: Array | undefined -): void { - if (!snapshotFilterFields || snapshotFilterFields.length === 0) return; - const needed = snapshotFilterFields.filter((ff) => ff && ff.hasGetter === true); - if (needed.length === 0) return; - const userGetterNames = new Set(); - if (Array.isArray(userFilterFields)) { - for (const uf of userFilterFields) { - if (uf && typeof uf === 'object' && typeof (uf as { name?: unknown }).name === 'string') { - const u = uf as { name: string; getter?: unknown }; - if (typeof u.getter === 'function') userGetterNames.add(u.name); - } - } - } - for (const ff of needed) { - if (!userGetterNames.has(ff.name)) { - throw new IncompatibleIndexError( - `filter-getter:${ff.name}`, - 'missing custom getter override for snapshot filter field (supply options.options.filterFields with getter)' - ); - } - } -} +export type { InternalField } from './document/state'; +export { assertSnapshotFilterGettersSatisfied } from './document/snapshot'; /** * High-level multi-field document search index with dynamic mutations and highlighting. - * Full document-index implementation. + * Facade over the `document/` implementation modules. */ export class DocumentIndex> { readonly options: DocumentIndexOptions; - private readonly normalized: boolean; - private readonly preferGpu: boolean; - private gpuEngine: WebGPUEngine | null = null; - private engineState: EngineState = initialEngineState(); - private vramAllocatedBytes: number = 0; - private unsubscribeDeviceLost?: () => void; - private isDestroyed: boolean = false; - private readonly profileId: TextProfileId = 'unicode-default'; - - // Active documents state - private records: TDoc[] = []; - private docIds: DocumentId[] = []; - private idToDocIndex: Map = new Map(); - private getId: (doc: TDoc) => DocumentId; - - // Fields and weights (stratified: sorted by weight desc) - private sortedFields: InternalField[] = []; - private readonly fieldNameToIndex: Map = new Map(); - - // Tokens and row mapping - private rowTokens: Uint32Array[] = []; - private rowToDocIndex: number[] = []; - private rowToFieldIndex: number[] = []; - private docToRowIndices: number[][] = []; // [docIndex][fieldIndex] -> row index - private rawFieldStrings: string[][] = []; // [docIndex][fieldIndex] - private totalTokens: number = 0; - private candidateCapacity: number = 8192; - private buildTimeMs: number = 0; - private restoreTimeMs?: number; - private lastMutationTimeMs?: number; - private mutationEpoch: number = 0; - private tombstones: Set = new Set(); // tombstone candidate filtering - private searchMutex: Promise = Promise.resolve(); - private generation: number = 0; - private readonly initialCapacity: number; - private readonly growthFactor: number; - private columnarStore: ColumnarStore; - private filterFieldDefinitions: FilterFieldDefinition[] = []; - private indexHooks: SearchHooks | undefined = undefined; + private readonly state: DocumentIndexState; constructor(options: DocumentIndexOptions) { if (!options || !Array.isArray(options.fields) || options.fields.length === 0) { @@ -241,136 +93,7 @@ export class DocumentIndex> { throw new ProfileMismatchError('unicode-default', options.textProfile, 'textProfile'); } this.options = options; - this.normalized = !(options.caseSensitive ?? false); - this.preferGpu = options.preferGpu ?? false; - // fail-closed extension hook validation at construction. - this.indexHooks = normalizeSearchHooks(options.hooks); - this.initialCapacity = typeof options.initialCapacity === 'number' && Number.isFinite(options.initialCapacity) && options.initialCapacity > 0 - ? Math.floor(options.initialCapacity) - : 0; - this.growthFactor = typeof options.growthFactor === 'number' && Number.isFinite(options.growthFactor) && options.growthFactor >= 1.0 - ? options.growthFactor - : 1.5; - - if (typeof options.idField === 'function') { - this.getId = options.idField; - } else if (typeof options.idField === 'string') { - const prop = options.idField; - this.getId = (doc: any) => doc[prop]; - } else if (options.idField === undefined) { - this.getId = (doc: any) => doc.id; - } else { - throw new TypeError('[webgpu-search] idField must be a string property name or function.'); - } - - const seenNames = new Set(); - const normalized: InternalField[] = options.fields.map((f, originalIndex) => { - let name: string; - let weight = 1.0; - let getter: (doc: TDoc) => string | string[] | undefined | null; - - if (typeof f === 'string') { - if (f.trim().length === 0) { - throw new TypeError('[webgpu-search] Field name string must not be empty.'); - } - name = f; - getter = (doc: any) => readOwnField(doc, name); - } else if (f && typeof f === 'object') { - if (typeof f.name !== 'string' || f.name.length === 0) { - throw new TypeError('[webgpu-search] Field definition requires a non-empty string name.'); - } - name = f.name; - if (f.weight !== undefined) { - if (typeof f.weight !== 'number' || !Number.isFinite(f.weight) || f.weight <= 0) { - throw new RangeError( - `[webgpu-search] Field "${name}" weight must be a positive finite number, got ${f.weight}.` - ); - } - weight = f.weight; - } - if (f.getter !== undefined) { - if (typeof f.getter !== 'function') { - throw new TypeError(`[webgpu-search] Field "${name}" getter must be a function.`); - } - getter = f.getter; - } else { - getter = (doc: any) => readOwnField(doc, name); - } - } else { - throw new TypeError('[webgpu-search] Field must be a string or FieldDefinition object.'); - } - - if (seenNames.has(name)) { - throw new Error(`[webgpu-search] Duplicate field name: "${name}".`); - } - seenNames.add(name); - - return { name, weight, getter, originalIndex }; - }); - - // Field-stratified packing: sort fields descending by weight - this.sortedFields = [...normalized].sort((a, b) => { - if (b.weight !== a.weight) return b.weight > a.weight ? 1 : -1; - if (a.originalIndex !== b.originalIndex) return a.originalIndex > b.originalIndex ? 1 : -1; - return 0; - }); - - for (let i = 0; i < this.sortedFields.length; i++) { - this.fieldNameToIndex.set(this.sortedFields[i].name, i); - } - - const seenFilterNames = new Set(); - const normalizedFilterFields: FilterFieldDefinition[] = []; - if (options.filterFields !== undefined) { - if (!Array.isArray(options.filterFields)) { - throw new TypeError('[webgpu-search] DocumentIndex expects options.filterFields to be an array.'); - } - for (let i = 0; i < options.filterFields.length; i++) { - const ff = options.filterFields[i]; - let name: string; - let type: FilterFieldType | undefined = undefined; - let getter: ((doc: TDoc) => any) | undefined = undefined; - let hasGetter = false; - - if (typeof ff === 'string') { - if (ff.trim().length === 0) { - throw new TypeError('[webgpu-search] Filter field name string must not be empty.'); - } - name = ff; - getter = (doc: any) => readOwnField(doc, name); - hasGetter = false; - } else if (ff && typeof ff === 'object') { - if (typeof ff.name !== 'string' || ff.name.trim().length === 0) { - throw new TypeError('[webgpu-search] Filter field definition requires a non-empty string name.'); - } - name = ff.name; - type = ff.type; - if (ff.getter !== undefined) { - if (typeof ff.getter !== 'function') { - throw new TypeError(`[webgpu-search] Filter field "${name}" getter must be a function.`); - } - getter = ff.getter; - hasGetter = true; - } else { - getter = (doc: any) => readOwnField(doc, name); - hasGetter = false; - } - } else { - throw new TypeError('[webgpu-search] Filter field must be a string or FilterFieldDefinition object.'); - } - - if (seenFilterNames.has(name)) { - throw new Error(`[webgpu-search] Duplicate filter field name: "${name}".`); - } - seenFilterNames.add(name); - normalizedFilterFields.push(hasGetter ? { name, type, getter, hasGetter: true } : { name, type, getter }); - } - } - this.filterFieldDefinitions = normalizedFilterFields; - this.columnarStore = new ColumnarStore(normalizedFilterFields, { - initialCapacity: this.initialCapacity, - growthFactor: this.growthFactor - }); + this.state = createDocumentIndexState(options); } static async create>( @@ -396,1753 +119,35 @@ export class DocumentIndex> { } private async initialize(records: TDoc[]): Promise { - const t0 = nowMs(); - const docCount = records.length; - const fieldCount = this.sortedFields.length; - const totalRows = docCount * fieldCount; - - this.records = records.slice(); - this.docIds = new Array(docCount); - this.idToDocIndex = new Map(); - - for (let i = 0; i < docCount; i++) { - const doc = this.records[i]; - if (doc === null || typeof doc !== 'object') { - throw new TypeError(`[webgpu-search] Document at index ${i} must be an object.`); - } - const id = this.getId(doc); - if ( - id === null || - id === undefined || - (typeof id !== 'string' && typeof id !== 'number') || - (typeof id === 'number' && !Number.isFinite(id)) - ) { - throw new TypeError( - `[webgpu-search] Document at index ${i} has invalid id: ${String(id)}. DocumentId must be a non-empty string or finite number.` - ); - } - if (typeof id === 'string' && id.length === 0) { - throw new TypeError( - `[webgpu-search] Document at index ${i} has an empty string id. DocumentId must not be empty.` - ); - } - if (this.idToDocIndex.has(id)) { - throw new DuplicateIdError(id); - } - this.docIds[i] = id; - this.idToDocIndex.set(id, i); - } - - // Dynamic Candidate Pool Scaling (Section 3.1): - // candidateCapacity = min(32768, max(8192, docCount * fieldCount * 0.1)) - const rawCap = this.options.candidateCapacity; - const validUserCap = typeof rawCap === 'number' && Number.isFinite(rawCap) && rawCap > 0 - ? Math.floor(rawCap) - : Math.floor(docCount * fieldCount * 0.1); - this.candidateCapacity = Math.min(32768, Math.max(8192, validUserCap)); - - // Field-Stratified Row Packing: - // Rows 0..N-1: Primary fields (highest weight). - // Rows N..2N-1: Secondary fields. - // Rows 2N..3N-1: Lower-weight fields. - this.rowTokens = new Array(totalRows); - this.rowToDocIndex = new Array(totalRows); - this.rowToFieldIndex = new Array(totalRows); - this.docToRowIndices = Array.from({ length: docCount }, () => new Array(fieldCount)); - this.rawFieldStrings = Array.from({ length: docCount }, () => new Array(fieldCount)); - - let corpusTokens = 0; - let rowIdx = 0; - - for (let f = 0; f < fieldCount; f++) { - const field = this.sortedFields[f]; - for (let d = 0; d < docCount; d++) { - const doc = this.records[d]; - const rawVal = field.getter(doc); - let rawStr = ''; - if (rawVal !== null && rawVal !== undefined) { - if (Array.isArray(rawVal)) { - rawStr = rawVal.filter((x) => x !== null && x !== undefined).join(' '); - } else { - rawStr = typeof rawVal === 'string' ? rawVal : String(rawVal); - } - } - this.rawFieldStrings[d][f] = rawStr; - const norm = normalizeText(rawStr, this.normalized); - this.rowTokens[rowIdx] = norm.tokens; - this.rowToDocIndex[rowIdx] = d; - this.rowToFieldIndex[rowIdx] = f; - this.docToRowIndices[d][f] = rowIdx; - corpusTokens += norm.tokenCount; - rowIdx++; - } - } - - this.totalTokens = corpusTokens; - this.columnarStore.init(this.records); - - if (docCount === 0) { - this.engineState = transitionEngineState( - this.engineState, - 'cpu', - this.options.preferGpu === false ? 'prefer-cpu' : 'below-threshold', - ); - this.buildTimeMs = nowMs() - t0; - return; - } - - const threshold = this.options.threshold ?? 30_000; - const shouldAttemptGpu = this.options.preferGpu === false - ? false - : this.preferGpu || totalRows >= threshold; - - if (this.options.preferGpu === false) { - this.engineState = transitionEngineState(this.engineState, 'cpu', 'prefer-cpu'); - } else if (!this.preferGpu && totalRows < threshold) { - this.engineState = transitionEngineState(this.engineState, 'cpu', 'below-threshold'); - } - - if (shouldAttemptGpu) { - const measuredAvgBytes = totalRows > 0 ? (corpusTokens * 4) / totalRows : 0; - const budget = checkMemoryBudget(totalRows, measuredAvgBytes, this.options.device, this.candidateCapacity); - if (!budget.allowed) { - console.warn(`[webgpu-search] ${budget.reason} Falling back to CPU.`); - this.engineState = transitionEngineState(this.engineState, 'cpu', 'memory-budget-exceeded'); - } else { - let gpu: WebGPUEngine | null = null; - try { - gpu = new WebGPUEngine(); - const initialized = await gpu.init( - this.options.device !== undefined || this.options.powerPreference !== undefined - ? { device: this.options.device, powerPreference: this.options.powerPreference } - : undefined - ); - - if (initialized && gpu.isReady) { - gpu.ensureCandidateCapacity(this.candidateCapacity); - const packed = packDataset(this.rowTokens, { - normalized: this.normalized, - totalTokens: corpusTokens - }); - const effectiveCap = this.initialCapacity > 0 ? this.initialCapacity : Math.max(16, docCount); - const rowCap = effectiveCap * fieldCount; - const avgTokens = totalRows > 0 ? corpusTokens / totalRows : 16; - const tokenCap = Math.max(Math.floor(rowCap * avgTokens), 64); - await gpu.loadDataset(packed, { - rowCapacity: rowCap, - tokenCapacity: tokenCap, - growthFactor: this.initialCapacity > 0 ? 1.0 : this.growthFactor - }); - - this.gpuEngine = gpu; - this.engineState = transitionEngineState(this.engineState, 'webgpu'); - this.vramAllocatedBytes = gpu.vramAllocatedBytes; - gpu = null; // Ownership transferred - - this.subscribeDeviceLost(); - } else { - this.engineState = transitionEngineState( - this.engineState, - 'cpu', - typeof navigator === 'undefined' || !('gpu' in navigator) || !navigator.gpu - ? 'webgpu-unsupported' - : 'device-request-failed', - ); - } - } catch (gpuErr) { - console.warn('[webgpu-search] WebGPU initialization failed, falling back to CPU:', gpuErr); - this.engineState = transitionEngineState( - this.engineState, - 'cpu', - typeof navigator === 'undefined' || !('gpu' in navigator) || !navigator.gpu - ? 'webgpu-unsupported' - : 'device-request-failed', - ); - } finally { - if (gpu) { - gpu.destroy(); - } - } - } - } else { - this.engineState = transitionEngineState( - this.engineState, - 'cpu', - this.engineState.fallbackReason, - ); - } - - this.buildTimeMs = nowMs() - t0; - } - - /** - * Validate search options fail-closed (extracted from `search`). - * Covers query shape, mode/scorer/profile checks, token/prefix/typo - * normalization, hooks, budget/diagnostics, query-length gates, facet, - * ranking/autocomplete, field scoping, and tokenizer expansion. - * No behavior change: code moved verbatim from `search`. - */ - private validateSearchOptions(query: string, options: DocumentSearchOptions) { - if (typeof query !== 'string') { - throw new TypeError(`[webgpu-search] search expects query: string, got ${typeof query}.`); - } - if ((options as unknown as Record).cpuAlgorithm !== undefined) { - throw new IncompatibleOptionError( - 'cpuScorer', - "[webgpu-search] cpuAlgorithm was removed; use cpuScorer: 'exact' | 'ufuzzy'." - ); - } - if ((options as unknown as Record).extensions !== undefined) { - throw new IncompatibleOptionError( - 'hooks', - "[webgpu-search] extensions was removed; use hooks." - ); - } - if ((options as unknown as Record).suggest !== undefined) { - throw new IncompatibleOptionError( - 'autocomplete', - "[webgpu-search] suggest was removed; use autocomplete." - ); - } - - const { - mode = 'fuzzy', - caseSensitive = false, - signal, - cpuScorer: cpuScorerOpt, - onQueryTooLong = 'throw', - fields: searchFields, - facets, - faceting - } = options; - - const cpuScorer = normalizeCpuScorer(cpuScorerOpt) ?? 'exact'; - if (caseSensitive === this.normalized) { - throw new ProfileMismatchError(!this.normalized, caseSensitive); - } - if (this.preferGpu && cpuScorer === 'ufuzzy') { - throw new IncompatibleOptionError( - 'cpuScorer', - "cpuScorer:'ufuzzy' is CPU-only; use preferGpu:false or cpuScorer:'exact'." - ); - } - assertValidMode(mode); - // fail-closed option validation up front so malformed - // token/prefix/typo shapes throw identically on GPU and CPU paths. - // 'fuzzy' validates typo shape but ignores it (subsequence matching is - // inherently typo-tolerant); 'substring'/'token'/'prefix' honor it. - const tokenOpts: NormalizedTokenMatchOptions = normalizeTokenMatchOptions(options.tokenMatch); - const prefixOpts: NormalizedPrefixOptions = normalizePrefixOptions(options.prefixMatch); - const typo: NormalizedTypoOptions = normalizeTypoTolerance(options.typoTolerance); - if (mode === 'prefix' && options.prefixMatch?.exactCase !== undefined && prefixOpts.exactCase !== caseSensitive) { - throw new ProfileMismatchError(caseSensitive, prefixOpts.exactCase, 'prefixMatch.exactCase'); - } - // Legacy ufuzzy/native scorers only implement fuzzy/substring-exact. - if (cpuScorer === 'ufuzzy' && (mode === 'token' || mode === 'prefix' || typo.enabled)) { - throw new IncompatibleOptionError( - 'cpuScorer', - `cpuScorer:'ufuzzy' supports only exact 'fuzzy'/'substring' modes without typo tolerance (got mode '${mode}'${typo.enabled ? ' with typoTolerance' : ''}). Use cpuScorer:'exact'.` - ); - } - // fail-closed extension hook validation up front so malformed - // hook shapes throw identically on GPU and CPU paths (before early exits). - const effectiveHooks = resolveEffectiveHooks(this.indexHooks, options.hooks); - // Threaded into every exact CPU call below (single normalization). - const cpuModeOptions: CpuModeOptions = { - tokenMatch: { operator: tokenOpts.operator, minMatchCount: tokenOpts.minMatchCount }, - prefixMatch: prefixOpts.prefixLength !== undefined - ? { prefixLength: prefixOpts.prefixLength, exactCase: prefixOpts.exactCase } - : { exactCase: prefixOpts.exactCase }, - typoTolerance: { - enabled: typo.enabled, - maxDistance: typo.maxDistance, - minWordLengthForOneTypo: typo.minWordLengthForOneTypo, - minWordLengthForTwoTypos: typo.minWordLengthForTwoTypos, - prefixExactLength: typo.prefixExactLength - } - }; - // fail-closed cost-budget + diagnostics validation up front so - // malformed budgets throw identically on GPU and CPU paths and on empty - // corpora/queries (before early exits). Enforcement is independent of the - // diagnostics flag: budgets constrain even when telemetry is off. - const budget = normalizeCostBudgetOptions(options.budget); - if (options.diagnostics !== undefined && typeof options.diagnostics !== 'boolean') { - throw new TypeError('[webgpu-search] options.diagnostics must be a boolean.'); - } - const wantsDiagnostics = options.diagnostics === true; - const needsClock = wantsDiagnostics || budget?.maxExecutionTimeMs !== undefined; - const queryStartMs = needsClock ? nowMs() : 0; - const diagWarnings: string[] = []; - - const rawTrimmed = query.trim(); - let forceCpu = false; - if (rawTrimmed.length > QUERY_TOKENS_MAX * 4) { - if (rawTrimmed.length > 1_000_000) { - let cpCount = 0; - for (const _ch of rawTrimmed) cpCount++; - const est = cpCount * 3; - if (est > QUERY_TOKENS_MAX) { - if (onQueryTooLong !== 'cpu-fallback') { - throw new QueryTooLongError(QUERY_TOKENS_MAX, est, this.profileId); - } - forceCpu = true; - } - } - } - - const normalizedQuery = normalizeText(query, this.normalized); - const queryTokenCount = normalizedQuery.tokenCount; - if (queryTokenCount > QUERY_TOKENS_MAX) { - if (onQueryTooLong === 'cpu-fallback') { - forceCpu = true; - } else { - throw new QueryTooLongError(QUERY_TOKENS_MAX, queryTokenCount, this.profileId); - } - } - - const clampedLimit = clampLimit(options.limit ?? options.maxResults ?? 50); - throwIfAborted(signal); - throwIfBudgetAborted(budget); - - // Hoisted prefixLength check: fail-closed even on empty queries/corpora. - // Skipped for empty queries to match the scorer's early noMatch. - if (mode === 'prefix' && normalizedQuery.tokens.length > 0) { - assertPrefixLengthForQuery(prefixOpts, normalizedQuery.tokens.length); - } - - // custom tokenizer terms for 'token' mode (CPU-only). - // Other modes ignore the tokenizer. Empty queries skip hook invocation - // (no-match by design); non-empty token queries expand via the hook so - // scoring and highlights share one term set (score-highlight symmetry). - let customTokenTerms: Uint32Array[] | undefined = undefined; - if (effectiveHooks?.tokenizer !== undefined && mode === 'token' && !normalizedQuery.isEmpty) { - customTokenTerms = getTokenTermsForQuery( - query, - normalizedQuery.tokens, - this.normalized, - effectiveHooks.tokenizer - ); - cpuModeOptions.tokenTermsOverride = customTokenTerms; - } - - // (): validate facet contracts fail-fast, before early exits. - // An empty facet list is equivalent to not requesting facets (absent key). - // `faceting` is ignored unless facets are requested (no throw on absent). - const facetSpec = normalizeFacetRequests(facets); - const wantsFacets = facetSpec !== undefined && facetSpec.length > 0; - let facetingMode: 'auto' | 'force-exact' = 'auto'; - if (wantsFacets && faceting !== undefined) { - if (faceting !== 'auto' && faceting !== 'force-exact') { - throw new TypeError(`[webgpu-search] Invalid faceting option: "${String(faceting)}". Must be 'auto' or 'force-exact'.`); - } - facetingMode = faceting; - } - const facetEngine = wantsFacets ? new FacetEngine(this.columnarStore) : null; - if (wantsFacets && facetEngine && facetSpec) { - // Fail-closed field validation up front: unknown facet fields and - // range-on-non-number must throw even when an empty match set would - // otherwise take an early noHits exit below. - facetEngine.validateRequests(facetSpec); - } - - // Inline autocomplete config: deterministic ranking hierarchy. - // All scope validation (ranking, autocomplete, fields) is fail-closed before - // early exits so malformed shapes throw identically on empty and - // non-empty corpora/queries. - if (options.ranking !== undefined && (typeof options.ranking !== 'object' || options.ranking === null || Array.isArray(options.ranking))) { - throw new TypeError('[webgpu-search] options.ranking must be an object.'); - } - const tieBreakers: TieBreakerCriterion[] = normalizeTieBreakers( - options.ranking?.tieBreakers - ); - let suggestSpec: NormalizedAutocompleteOptions | undefined = undefined; - const suggestRaw = options.autocomplete; - if (suggestRaw !== undefined && suggestRaw !== false) { - suggestSpec = normalizeAutocompleteOptions(suggestRaw as AutocompleteOptions | boolean); - // Inline autocomplete inherits the search ranking hierarchy unless the - // autocomplete object carries its own explicit tieBreakers. - const rawSuggest = suggestRaw as AutocompleteOptions; - if ( - typeof rawSuggest === 'object' && - rawSuggest !== null && - rawSuggest.tieBreakers === undefined && - options.ranking?.tieBreakers !== undefined - ) { - suggestSpec.tieBreakers = [...tieBreakers]; - } - if (suggestSpec.field !== undefined && !this.fieldNameToIndex.has(suggestSpec.field)) { - throw new Error(`[webgpu-search] Unknown autocomplete field: "${suggestSpec.field}".`); - } - } - // Hoisted search-field validation (fail-closed before empty-query, - // empty-corpus, and empty-filter early exits below). - let allowedFieldIndices: Set | undefined = undefined; - if (searchFields !== undefined) { - if (!Array.isArray(searchFields)) { - throw new TypeError('[webgpu-search] options.fields must be an array of string field names.'); - } - if (searchFields.length !== 0) { - allowedFieldIndices = new Set(); - for (const fName of searchFields) { - if (typeof fName !== 'string') { - throw new TypeError('[webgpu-search] options.fields elements must be strings.'); - } - const idx = this.fieldNameToIndex.get(fName); - if (idx === undefined) { - throw new Error(`[webgpu-search] Unknown search field: "${fName}".`); - } - allowedFieldIndices.add(idx); - } - } - } - const hasEmptyFieldList = searchFields !== undefined && searchFields.length === 0; - - return { - mode, - caseSensitive, - signal, - cpuScorer, - tokenOpts, - prefixOpts, - typo, - effectiveHooks, - cpuModeOptions, - budget, - wantsDiagnostics, - queryStartMs, - diagWarnings, - normalizedQuery, - clampedLimit, - forceCpu, - facetSpec, - wantsFacets, - facetingMode, - facetEngine, - tieBreakers, - suggestSpec, - allowedFieldIndices, - hasEmptyFieldList, - customTokenTerms - }; - } - - /** - * Compile structured/predicate filters (extracted from `search`). - * Compiles `FilterExpression` to a bitset, measures `filteringMs`, - * enforces abort/budget gates, and composes extension predicates. - * Empty bitsets return early with `filteringMs: 0` so callers can take - * the `noHits` path with identical telemetry. No behavior change. - */ - private compileFilterPhase( - filter: DocumentSearchOptions['filter'], - effectiveHooks: SearchHooks | undefined, - wantsDiagnostics: boolean, - signal: AbortSignal | undefined, - budget: ReturnType, - queryStartMs: number - ) { - let filterBitset: DocumentBitset | undefined = undefined; - let filterPredicate: ((doc: TDoc) => boolean) | undefined = undefined; - let structuredFilter: FilterExpression | undefined = undefined; - let filteringMs = 0; - - const tFilter0 = wantsDiagnostics ? nowMs() : 0; - if (filter !== undefined) { - if (typeof filter === 'function') { - filterPredicate = filter; - } else if (typeof filter === 'object' && filter !== null) { - structuredFilter = filter as FilterExpression; - filterBitset = compileFilter(structuredFilter, this.columnarStore); - if (filterBitset.isEmpty()) { - return { filterBitset, filterPredicate, structuredFilter, filteringMs }; - } - } else { - throw new TypeError('[webgpu-search] options.filter must be a function or FilterExpression.'); - } - } - if (wantsDiagnostics) filteringMs = nowMs() - tFilter0; - throwIfAborted(signal); - throwIfBudgetAborted(budget); - assertTimeBudget(queryStartMs, budget); - - // extension filter predicates compose conjunctively (AND) with - // `options.filter` functions. Index-level vs per-query hooks already - // merged via `resolveEffectiveHooks` (per-query wins per key). - const extensionPredicate = effectiveHooks?.filterPredicate; - if (extensionPredicate !== undefined) { - const basePredicate = filterPredicate; - if (basePredicate === undefined) { - filterPredicate = extensionPredicate; - } else { - const optFn = basePredicate; - const extFn = extensionPredicate; - filterPredicate = (doc: TDoc) => { - const d = doc as TDoc; - return optFn(d) && extFn(d); - }; - } - } - - return { filterBitset, filterPredicate, structuredFilter, filteringMs }; - } - - /** - * Route to WebGPU vs CPU (extracted from `search`). - * Applies candidate-budget enforcement, broad-query CPU routing, and - * field-restriction routing. Mutates `diagWarnings` for the broad-query - * warning exactly as before. No behavior change. - */ - private routeEngine(params: { - mode: SearchMode; - typoEnabled: boolean; - cpuScorer: string; - forceCpu: boolean; - allowedFieldIndices: Set | undefined; - normalizedQueryTokenCount: number; - filterBitset: DocumentBitset | undefined; - budget: ReturnType; - queryStartMs: number; - wantsDiagnostics: boolean; - diagWarnings: string[]; - }) { - const { - mode, - typoEnabled, - cpuScorer, - forceCpu, - allowedFieldIndices, - normalizedQueryTokenCount, - filterBitset, - budget, - queryStartMs, - wantsDiagnostics, - diagWarnings - } = params; - // token/prefix modes and typo-tolerant queries are CPU-only - // (exact-only WGSL kernels) and skip GPU dispatch with fallbackReason - // 'unsupported-mode' (parity boundary). - const isGpuSupportedMode: boolean = - (mode === 'fuzzy' || mode === 'substring') && !typoEnabled; - - // broad-query pre-dispatch guard + candidate ceiling. - // `filteredCandidateCount` is the exact post-filter population to score; - // exceeding `maxCandidates` throws before any scoring work. Short queries - // over massive corpora route to the CPU streaming scan to avoid GPU - // buffer saturation and driver timeouts (TDR). With function-predicate - // filters the pre-predicate population is checked (conservative — the - // true post-predicate count is unknowable pre-scan). - const activeDocCount = this.idToDocIndex.size; - const filteredCandidateCount = filterBitset !== undefined - ? filterBitset.popcount() - : activeDocCount; - assertCandidateBudget(filteredCandidateCount, budget); - let broadQueryCpuRoute = false; - if (isBroadQueryHeuristic(normalizedQueryTokenCount, activeDocCount)) { - broadQueryCpuRoute = true; - // Suppress the GPU-avoidance warning when already CPU-by-design - // (token/prefix/typo modes, explicit ufuzzy, query-too-long fallback): - // routing outcome is correct but the stated cause would misattribute. - const cpuByDesign = !isGpuSupportedMode || cpuScorer === 'ufuzzy' || forceCpu; - if (wantsDiagnostics && !cpuByDesign) { - diagWarnings.push(broadQueryRouteWarning(activeDocCount, normalizedQueryTokenCount)); - } - } - assertTimeBudget(queryStartMs, budget); - - // When field-restricted, route to CPU to prevent unselected high-priority - // fields from saturating GPU candidate buffer before low-priority allowed fields. - const isFieldRestricted = allowedFieldIndices !== undefined && allowedFieldIndices.size < this.sortedFields.length; - - // 1. WebGPU execution path (exact fuzzy/substring only; token/prefix - // and typo queries skip dispatch via isGpuSupportedMode above). - const gpuHandle = this.gpuEngine; - const useGpu = !forceCpu && - !isFieldRestricted && - !broadQueryCpuRoute && - isGpuSupportedMode && - cpuScorer !== 'ufuzzy' && - this.engineState.engine === 'webgpu' && - gpuHandle !== null && - gpuHandle.isReady; - - return { - isGpuSupportedMode, - broadQueryCpuRoute, - isFieldRestricted, - filteredCandidateCount, - useGpu, - gpuHandle - }; - } - - /** - * Aggregate per-row matches into per-document best-field hits (thin - * delegator over the pure `document-aggregate` helper so GPU/CPU paths - * share one ranking implementation; unit-tested without a device). - * No behavior change. - */ - private aggregateDocMatches( - docMatches: Map }>, - normalizedQueryTokens: Uint32Array, - tieBreakers: TieBreakerCriterion[] - ) { - return aggregateDocMatchesPure(docMatches, { - sortedFields: this.sortedFields, - docToRowIndices: this.docToRowIndices, - rowTokens: this.rowTokens, - docIds: this.docIds, - records: this.records, - normalizedQueryTokens, - tieBreakers - }); - } - - /** - * Enrich Top-K results with highlights + post-match hooks (extracted). - * Runs highlight enrichment, scoring-hook boost + deterministic re-sort, - * and postProcess transform. Returns updated results with phase timings. - * No behavior change. - */ - private enrichResults( - results: DocumentSearchResultItem[], - query: string, - mode: SearchMode, - options: DocumentSearchOptions, - customTokenTerms: Uint32Array[] | undefined, - normalizedQueryTokens: Uint32Array, - effectiveHooks: SearchHooks | undefined, - tieBreakers: TieBreakerCriterion[], - wantsDiagnostics: boolean - ) { - const tHl0 = wantsDiagnostics ? nowMs() : 0; - this.enrichHighlights(results, query, mode, options, customTokenTerms, normalizedQueryTokens); - const highlightMs = wantsDiagnostics ? nowMs() - tHl0 : 0; - - // post-match extension pipeline (Top-K only). Scoring boosts - // apply identically on GPU and CPU paths (parity preserved), followed - // by a deterministic re-sort; postProcess is the final transform. - // hook time joins the scoring bucket on all paths for comparability. - const tHook0 = wantsDiagnostics ? nowMs() : 0; - if (effectiveHooks?.scoringHook !== undefined) { - applyScoringHook(results, effectiveHooks.scoringHook, query); - this.resortResultsAfterScoring(results, normalizedQueryTokens, tieBreakers); - } - let out = results; - if (effectiveHooks?.postProcess !== undefined) { - out = applyPostProcess(results, effectiveHooks.postProcess); - } - const hookMs = wantsDiagnostics ? nowMs() - tHook0 : 0; - return { results: out, highlightMs, hookMs }; + await initializeDocumentIndex(this.state, records); } async search( query: string, options: DocumentSearchOptions = {} ): Promise> { - if (this.isDestroyed) { + if (this.state.isDestroyed) { throw new Error('[webgpu-search] DocumentIndex has been destroyed.'); } throwIfAborted(options.signal); - await this.searchMutex; - if (this.isDestroyed) { + await this.state.searchMutex; + if (this.state.isDestroyed) { throw new Error('[webgpu-search] DocumentIndex has been destroyed.'); } throwIfAborted(options.signal); - const gen = this.generation; - - const validated = this.validateSearchOptions(query, options); - const { - mode, - caseSensitive, - signal, - cpuScorer, - typo, - effectiveHooks, - cpuModeOptions, - budget, - wantsDiagnostics, - queryStartMs, - diagWarnings, - normalizedQuery, - clampedLimit, - forceCpu, - facetSpec, - wantsFacets, - facetingMode, - facetEngine, - tieBreakers, - suggestSpec, - allowedFieldIndices, - hasEmptyFieldList, - customTokenTerms - } = validated; - - const filterState = this.compileFilterPhase( - options.filter, - effectiveHooks, - wantsDiagnostics, - signal, - budget, - queryStartMs - ); - const { filterBitset, filterPredicate, structuredFilter, filteringMs } = filterState; - /** - * assemble response diagnostics (undefined unless requested). - * `filteringMs` is read at call time so early exits after compilation - * report measured filter cost; `totalMs` is wall-clock from query entry - * (end-to-end including inline autocomplete; see `autocompleteMs` for the slice). - * Scoring includes post-match scoring-hook time on all paths. - */ - const buildDiagnostics = ( - routedEngine: EngineType, - scoringMs: number, - highlightMs: number, - facetingMs: number | undefined, - hasOverflow: boolean, - autocompleteMs?: number - ): QueryDiagnostics | undefined => { - if (!wantsDiagnostics) return undefined; - const scanned = this.idToDocIndex.size; - const selectivity = filterBitset !== undefined - ? computeFilterSelectivity(filterBitset.popcount(), scanned) - : 1.0; - const totalMs = nowMs() - queryStartMs; - const timings: QueryDiagnosticsTimings = { - filteringMs, - scoringMs, - highlightMs, - ...(facetingMs !== undefined ? { facetingMs } : {}), - ...(autocompleteMs !== undefined ? { autocompleteMs } : {}), - totalMs - }; - const diag: QueryDiagnostics = { - scannedCandidates: scanned, - filterSelectivity: selectivity, - routedEngine, - hasOverflow, - timings - }; - if (diagWarnings.length > 0) diag.warnings = [...diagWarnings]; - return diag; - }; - /** - * post-hoc broad-query + overflow warnings (non-fatal, gated on - * `diagnostics:true` — routing decisions stay ungated for safety). - * `rawTotalMatches` carries the pre-filter pool count on the GPU path so - * the overflow sentence stays accurate; `facetsExact` suppresses the - * stale "approximate" claim after a force-exact rescan. - */ - const pushPostHocWarnings = ( - totalMatches: number, - hasOverflow: boolean, - opts?: { rawTotalMatches?: number; facetsRequested?: boolean; facetsExact?: boolean } - ): void => { - if (!wantsDiagnostics) return; - const scanned = this.idToDocIndex.size; - const selectivity = scanned > 0 ? totalMatches / scanned : 0; - if (isBroadSelectivity(selectivity, scanned)) { - diagWarnings.push(broadSelectivityWarning(selectivity, scanned)); - } - if (hasOverflow) { - diagWarnings.push( - candidateOverflowWarning(totalMatches, this.candidateCapacity, { - facetsRequested: opts?.facetsRequested ?? false, - facetsExact: opts?.facetsExact ?? false, - rawTotalMatches: opts?.rawTotalMatches, - }) - ); - } - }; - - const noHits = (q: string): DocumentSearchResponse => { - // even trivial exits honor caller aborts and time budgets. - throwIfAborted(signal); - throwIfBudgetAborted(budget); - assertTimeBudget(queryStartMs, budget); - const diag = buildDiagnostics('cpu', 0, 0, wantsFacets ? 0 : undefined, false); - // Public totalMs is scorer wall-clock (0 on no-hit: nothing scored); - // diagnostics.totalMs carries the wall-clock including validation. - return { - query: q, - mode, - engine: 'cpu', - totalMatches: 0, - candidateCount: 0, - hasOverflow: false, - results: [], - timings: { - queryUploadMs: 0, - encodeSubmitMs: 0, - gpuExecutionMs: null, - readbackMs: 0, - totalMs: 0, - gpuDispatchMs: 0 - }, - profileId: this.profileId, - scoringVersion: SCORING_VERSION, - cpuScorer, - fallbackReason: forceCpu ? 'query-too-long' : this.engineState.fallbackReason, - ...(wantsFacets && facetEngine && facetSpec - ? { facets: facetEngine.emptyResults(facetSpec) } - : {}), - ...(diag ? { diagnostics: diag } : {}) - }; - }; - - if (filterBitset !== undefined && filterBitset.isEmpty()) { - return noHits(query); - } - - if (normalizedQuery.isEmpty) { - return noHits(''); - } - if (this.idToDocIndex.size === 0) { - return noHits(query); - } - if (hasEmptyFieldList) { - return noHits(query); - } - - const route = this.routeEngine({ - mode, - typoEnabled: typo.enabled, - cpuScorer, - forceCpu, - allowedFieldIndices, - normalizedQueryTokenCount: normalizedQuery.tokens.length, - filterBitset, - budget, - queryStartMs, - wantsDiagnostics, - diagWarnings - }); - const { broadQueryCpuRoute, isGpuSupportedMode, useGpu, gpuHandle } = route; - - if (useGpu && gpuHandle !== null) { - // narrow GPU try to dispatch only so extension hook - // errors (filterPredicate / scoringHook / postProcess) propagate - // directly without destroying the engine or double-invoking via CPU - // fallback. Only dispatch failures fall through to CPU. - let gpuResult: Awaited> | undefined = undefined; - try { - gpuResult = await gpuHandle.search(query, { - ...options, - mode, - limit: this.candidateCapacity, - caseSensitive, - signal - }); - - if (this.isDestroyed || this.generation !== gen) { - throw abortError(); - } - throwIfAborted(signal); - // GPU dispatch is async wall-clock work — enforce the caller - // deadline and budget abort before touching the readback. - throwIfBudgetAborted(budget); - assertTimeBudget(queryStartMs, budget); - } catch (err: any) { - if (err.name === 'AbortError') { - throw err; - } - // Mirror search-index: mode-routed rejections (token/prefix/typo via - // direct or raced calls) are not engine failures — keep a healthy GPU - // up and fall through to exact CPU. - const modeRouted: boolean = - err instanceof IncompatibleOptionError && (err.option === 'mode' || err.option === 'typoTolerance'); - if (!modeRouted) { - console.warn('[webgpu-search] GPU document search failed, falling back to CPU:', err); - this.engineState = transitionEngineState(this.engineState, 'cpu', 'gpu-execution-error'); - if (this.gpuEngine) { - try { this.gpuEngine.destroy(); } catch {} - this.gpuEngine = null; - } - } - gpuResult = undefined; - } - - if (gpuResult !== undefined) { - const gpuResultOk = gpuResult; - - // WebGPU candidate readback & fixed-point score enrichment - const docMatches = new Map; - }>(); - // unfiltered query match set for facet aggregation (filter applied below). - const queryMatchedAll = wantsFacets ? new Set() : null; - - // hook errors propagate (no fallback / no double-invoke). - // Only GPU dispatch failures fall back; readback + hook pipeline runs - // outside the dispatch try above. - for (let i = 0; i < gpuResultOk.results.length; i++) { - const item = gpuResultOk.results[i]; - const r = item.index; - if (this.tombstones.has(r)) continue; - const fIdx = this.rowToFieldIndex[r]; - if (allowedFieldIndices && !allowedFieldIndices.has(fIdx)) continue; - const dIdx = this.rowToDocIndex[r]; - if (queryMatchedAll !== null && this.records[dIdx]) { - queryMatchedAll.add(dIdx); - } - if (filterBitset && !filterBitset.has(dIdx)) continue; - const doc = this.records[dIdx]; - if (!doc) continue; - if (filterPredicate && !filterPredicate(doc)) continue; - - const rawScore = item.score; - const fDef = this.sortedFields[fIdx]; - const weightedScore = Math.round(rawScore * fDef.weight); - - accumulateDocMatch(docMatches, dIdx, fIdx, weightedScore); - } - - const hits = this.aggregateDocMatches(docMatches, normalizedQuery.tokens, tieBreakers); - - const totalMatches = hits.length; - const candidateCount = Math.min(totalMatches, this.candidateCapacity); - let results = hits.slice(0, clampedLimit).map((h) => h.item); - const gpuEnriched = this.enrichResults(results, query, mode, options, customTokenTerms, normalizedQuery.tokens, effectiveHooks, tieBreakers, wantsDiagnostics); - results = gpuEnriched.results; - const gpuHighlightMs = gpuEnriched.highlightMs; - const gpuHookMs = gpuEnriched.hookMs; - - // facet aggregation. Exact when the GPU pool covered all matches; - // approximate over the top pool on overflow unless force-exact rescan. - let gpuFacets: Record | undefined = undefined; - let gpuFacetMs: number | undefined = undefined; - let gpuFacetExact = true; - if (wantsFacets && facetEngine && facetSpec && queryMatchedAll !== null) { - const facetT0 = nowMs(); - let facetCandidates: Iterable = queryMatchedAll; - let facetIsApproximate = gpuResultOk.hasOverflow; - if (gpuResultOk.hasOverflow && facetingMode === 'force-exact') { - const exact = scoreExactMatchesMultiField( - this.records.length, - this.sortedFields, - this.rowTokens, - this.rowToDocIndex, - this.rowToFieldIndex, - normalizedQuery.tokens, - mode, - clampedLimit, - this.candidateCapacity, - allowedFieldIndices, - this.tombstones, - undefined, - true, - cpuModeOptions - ); - facetCandidates = exact.allMatchedDocIndices ?? exact.results.map((r) => r.docIndex); - facetIsApproximate = false; - } - gpuFacets = this.buildFacetResults( - facetEngine, facetSpec, facetCandidates, structuredFilter, filterPredicate, facetIsApproximate - ); - const facetMs = nowMs() - facetT0; - gpuFacetMs = facetMs; - gpuFacetExact = !facetIsApproximate; - } - - // deadline + post-hoc warnings before returning GPU results. - // Inline autocomplete runs before the diagnostics snapshot so - // diagnostics.totalMs covers end-to-end latency (autocomplete bucket). - const tGpuSuggest0 = wantsDiagnostics && suggestSpec ? nowMs() : 0; - const gpuSuggestions = suggestSpec - ? this.computeSuggestions(query, normalizedQuery.tokens, suggestSpec, allowedFieldIndices) - : undefined; - const gpuSuggestMs = wantsDiagnostics && suggestSpec ? nowMs() - tGpuSuggest0 : undefined; - throwIfAborted(signal); - throwIfBudgetAborted(budget); - assertTimeBudget(queryStartMs, budget); - pushPostHocWarnings(totalMatches, gpuResultOk.hasOverflow, { - rawTotalMatches: gpuResultOk.totalMatches, - facetsRequested: wantsFacets, - facetsExact: gpuFacetExact, - }); - // Scoring = engine scan + post-match hooks (facet time excluded — - // no double-counting, no mutation of the engine timings object). - const gpuEngineTotal = gpuResultOk.timings.totalMs; - const gpuDiag = buildDiagnostics( - 'webgpu', - gpuEngineTotal + gpuHookMs, - gpuHighlightMs, - gpuFacetMs, - gpuResultOk.hasOverflow, - gpuSuggestMs - ); - - return { - query: gpuResultOk.query, - mode: gpuResultOk.mode, - engine: 'webgpu', - totalMatches, - candidateCount, - hasOverflow: gpuResultOk.hasOverflow, - results, - timings: { - ...gpuResultOk.timings, - totalMs: gpuEngineTotal + (gpuFacetMs ?? 0) + gpuHighlightMs + gpuHookMs, - }, - profileId: this.profileId, - scoringVersion: SCORING_VERSION, - cpuScorer, - ...(gpuFacets ? { facets: gpuFacets } : {}), - ...(gpuDiag ? { diagnostics: gpuDiag } : {}), - ...(gpuSuggestions ? { suggestions: gpuSuggestions } : {}) - }; - } - } - - // 2. CPU execution path - if (this.isDestroyed || this.generation !== gen) { - throw abortError(); - } - throwIfAborted(signal); - throwIfBudgetAborted(budget); - assertTimeBudget(queryStartMs, budget); - const t0 = nowMs(); - - if (cpuScorer === 'ufuzzy') { - const cpuEngine = new CPUEngine(); - const docMatches = new Map; - }>(); - // unfiltered query match set for exact facet aggregation. - const legacyMatchedAll = wantsFacets ? new Set() : null; - - for (let fIdx = 0; fIdx < this.sortedFields.length; fIdx++) { - if (allowedFieldIndices && !allowedFieldIndices.has(fIdx)) continue; - const fDef = this.sortedFields[fIdx]; - const fieldStrings = this.records.map((doc, dIdx) => (doc && this.rawFieldStrings[dIdx] ? this.rawFieldStrings[dIdx][fIdx] : '')); - const legacyResult = mode === 'fuzzy' - ? cpuEngine.searchWithUFuzzy(fieldStrings, query, this.records.length, caseSensitive) - : cpuEngine.searchNaiveScan(fieldStrings, query, this.records.length, caseSensitive); - - for (const hit of legacyResult.results) { - const dIdx = hit.index; - // Explicit tombstone guard for symmetry with GPU/exact paths: - // removed docs have null records (columnar presence also cleared). - if (!this.records[dIdx] || this.docIds[dIdx] === null || this.docIds[dIdx] === undefined) continue; - if (legacyMatchedAll !== null) { - legacyMatchedAll.add(dIdx); - } - if (filterBitset && !filterBitset.has(dIdx)) continue; - const doc = this.records[dIdx]; - if (!doc) continue; - if (filterPredicate && !filterPredicate(doc)) continue; - const weightedScore = Math.round(hit.score * fDef.weight); - accumulateDocMatch(docMatches, dIdx, fIdx, weightedScore); - } - } - - const hits = this.aggregateDocMatches(docMatches, normalizedQuery.tokens, tieBreakers); - - const totalMatches = hits.length; - const candidateCount = Math.min(totalMatches, this.candidateCapacity); - let results = hits.slice(0, clampedLimit).map((h) => h.item); - const legacyEnriched = this.enrichResults(results, query, mode, options, customTokenTerms, normalizedQuery.tokens, effectiveHooks, tieBreakers, wantsDiagnostics); - results = legacyEnriched.results; - const legacyHighlightMs = legacyEnriched.highlightMs; - const preFacetMs = nowMs() - t0; - // scoring bucket covers scan + ranking + post-match hooks; - // highlight enrichment is metered separately above. - const legacyScoringMs = Math.max(0, preFacetMs - legacyHighlightMs); - - // CPU evaluates the full match set, so legacy facets are exact - // w.r.t. the serving (ufuzzy/native) match set; see types for caveat. - let legacyFacets: Record | undefined = undefined; - let legacyFacetMs = 0; - if (wantsFacets && facetEngine && facetSpec && legacyMatchedAll !== null) { - const fT0 = nowMs(); - legacyFacets = this.buildFacetResults( - facetEngine, facetSpec, legacyMatchedAll, structuredFilter, filterPredicate, false - ); - legacyFacetMs = nowMs() - fT0; - } - - // inline autocomplete precedes the diagnostics snapshot so - // diagnostics.totalMs covers end-to-end latency (autocompleteMs bucket). - const tLegacySuggest0 = wantsDiagnostics && suggestSpec ? nowMs() : 0; - const legacySuggestions = suggestSpec - ? this.computeSuggestions(query, normalizedQuery.tokens, suggestSpec, allowedFieldIndices) - : undefined; - const legacySuggestMs = wantsDiagnostics && suggestSpec ? nowMs() - tLegacySuggest0 : undefined; - throwIfAborted(signal); - throwIfBudgetAborted(budget); - assertTimeBudget(queryStartMs, budget); - pushPostHocWarnings(totalMatches, totalMatches > this.candidateCapacity, { - facetsRequested: wantsFacets, - facetsExact: true, - }); - const legacyDiag = buildDiagnostics( - 'cpu', - legacyScoringMs, - legacyHighlightMs, - wantsFacets ? legacyFacetMs : undefined, - totalMatches > this.candidateCapacity, - legacySuggestMs - ); - - return { - query, - mode, - engine: 'cpu', - totalMatches, - candidateCount, - hasOverflow: totalMatches > this.candidateCapacity, - results, - timings: { - queryUploadMs: 0, - encodeSubmitMs: 0, - gpuExecutionMs: null, - readbackMs: 0, - totalMs: preFacetMs + legacyFacetMs, - gpuDispatchMs: 0 - }, - profileId: this.profileId, - scoringVersion: SCORING_VERSION, - cpuScorer, - fallbackReason: 'cpu-algorithm-requested', - ...(legacyFacets ? { facets: legacyFacets } : {}), - ...(legacyDiag ? { diagnostics: legacyDiag } : {}), - ...(legacySuggestions ? { suggestions: legacySuggestions } : {}) - }; - } - - // Default exact CPU algorithm (collect full match set only when faceting). - // passes the deterministic hierarchy + doc IDs so pre-truncation order - // already reflects (score, weight, exact, length, id). - const parityResult = scoreExactMatchesMultiField( - this.records.length, - this.sortedFields, - this.rowTokens, - this.rowToDocIndex, - this.rowToFieldIndex, - normalizedQuery.tokens, - mode, - clampedLimit, - this.candidateCapacity, - allowedFieldIndices, - this.tombstones, - (filterBitset || filterPredicate) ? (dIdx) => { - if (filterBitset && !filterBitset.has(dIdx)) return false; - if (filterPredicate) { - const doc = this.records[dIdx]; - return doc !== null && doc !== undefined && filterPredicate(doc); - } - return true; - } : undefined, - wantsFacets, - cpuModeOptions, - { tieBreakers, docIds: this.docIds as unknown as ReadonlyArray } - ); - - if (this.isDestroyed || this.generation !== gen) { - throw abortError(); - } - throwIfAborted(signal); - throwIfBudgetAborted(budget); - assertTimeBudget(queryStartMs, budget); - - const idProp = typeof this.options.idField === 'string' ? this.options.idField : 'id'; - let enrichedResults: DocumentSearchResultItem[] = []; - for (let i = 0; i < parityResult.results.length; i++) { - const hit = parityResult.results[i]; - const id = this.docIds[hit.docIndex]; - if (id === null || id === undefined) continue; - const doc = this.records[hit.docIndex] ?? ({ [idProp]: id } as any); - enrichedResults.push({ - id, - doc, - score: hit.score, - matchedField: hit.matchedField, - matches: hit.matches - }); - } - const parityEnriched = this.enrichResults(enrichedResults, query, mode, options, customTokenTerms, normalizedQuery.tokens, effectiveHooks, tieBreakers, wantsDiagnostics); - enrichedResults = parityEnriched.results; - const parityHighlightMs = parityEnriched.highlightMs; - const parityHookMs = parityEnriched.hookMs; - const parityScoringMs = parityResult.durationMs + parityHookMs; - - // exact facet aggregation. Disjunctive facets need the unfiltered - // query match set, so a structured filter triggers one extra unfiltered - // parity scan (results above stay filtered and correctly ranked; O(F*n) - // per-facet re-evaluation + F bitset allocs by design for disjunctive - // semantics). Wall-clock facet work is included in totalMs; future - // attribute facet work to diagnostics.timings.facetingMs. - let parityFacets: Record | undefined = undefined; - const facetT0 = nowMs(); - let facetScanMs = 0; - if (wantsFacets && facetEngine && facetSpec) { - let facetCandidates: Iterable = - parityResult.allMatchedDocIndices ?? parityResult.results.map((r) => r.docIndex); - if (structuredFilter !== undefined) { - const unfiltered = scoreExactMatchesMultiField( - this.records.length, - this.sortedFields, - this.rowTokens, - this.rowToDocIndex, - this.rowToFieldIndex, - normalizedQuery.tokens, - mode, - clampedLimit, - this.candidateCapacity, - allowedFieldIndices, - this.tombstones, - undefined, - true, - cpuModeOptions - ); - facetCandidates = unfiltered.allMatchedDocIndices ?? unfiltered.results.map((r) => r.docIndex); - facetScanMs = unfiltered.durationMs; - } - parityFacets = this.buildFacetResults( - facetEngine, facetSpec, facetCandidates, structuredFilter, filterPredicate, false - ); - } - const facetMs = wantsFacets ? Math.max(0, nowMs() - facetT0 - facetScanMs) : 0; - // the disjunctive unfiltered rescan serves facets, so it joins the - // faceting bucket (scoring stays scan + post-match hooks). - const parityFacetingMs = wantsFacets ? facetScanMs + facetMs : undefined; - - // Public totalMs is scorer wall-clock: scan + hooks + highlight + facets - // (excludes inline autocomplete; diagnostics.totalMs is end-to-end). - const timings: SearchTimings = { - queryUploadMs: 0, - encodeSubmitMs: 0, - gpuExecutionMs: null, - readbackMs: 0, - totalMs: parityScoringMs + parityHighlightMs + facetScanMs + facetMs, - gpuDispatchMs: 0 - }; - - let effectiveFallbackReason = this.engineState.fallbackReason; - if (forceCpu) { - effectiveFallbackReason = 'query-too-long'; - } else if (!isGpuSupportedMode) { - // token/prefix/typo queries are CPU-by-design (exact-only - // WGSL kernels) — recorded per the exact boundary. - effectiveFallbackReason = 'unsupported-mode'; - } else if (useGpu && gpuHandle !== null) { - effectiveFallbackReason = 'gpu-execution-error'; - } else if (broadQueryCpuRoute) { - // broad-query CPU routing is a routing decision (like - // field-restriction below), not a scorer request — leave the reason - // as-is and surface the decision in diagnostics.warnings instead. - } - // Note: field-restricted routing (isFieldRestricted) intentionally leaves - // fallbackReason as-is — it is a routing decision, not a scorer request, - // so reusing 'cpu-algorithm-requested' would mislead telemetry. - - // inline autocomplete precedes the diagnostics snapshot so - // diagnostics.totalMs covers end-to-end latency (autocompleteMs bucket). - const tParitySuggest0 = wantsDiagnostics && suggestSpec ? nowMs() : 0; - const paritySuggestions = suggestSpec - ? this.computeSuggestions(query, normalizedQuery.tokens, suggestSpec, allowedFieldIndices) - : undefined; - const paritySuggestMs = wantsDiagnostics && suggestSpec ? nowMs() - tParitySuggest0 : undefined; - throwIfAborted(signal); - throwIfBudgetAborted(budget); - assertTimeBudget(queryStartMs, budget); - pushPostHocWarnings(parityResult.totalMatches, parityResult.hasOverflow, { - facetsRequested: wantsFacets, - facetsExact: true, - }); - const parityDiag = buildDiagnostics( - 'cpu', - parityScoringMs, - parityHighlightMs, - parityFacetingMs, - parityResult.hasOverflow, - paritySuggestMs - ); + const gen = this.state.generation; - return { - query, - mode, - engine: 'cpu', - totalMatches: parityResult.totalMatches, - candidateCount: parityResult.candidateCount, - hasOverflow: parityResult.hasOverflow, - results: enrichedResults, - timings, - profileId: this.profileId, - scoringVersion: SCORING_VERSION, - cpuScorer, - fallbackReason: effectiveFallbackReason, - ...(parityFacets ? { facets: parityFacets } : {}), - ...(parityDiag ? { diagnostics: parityDiag } : {}), - ...(paritySuggestions ? { suggestions: paritySuggestions } : {}) - }; + return runSearch(this.state, query, options, gen); } - /** - * First-party autocomplete / did-you-mean primitive. - * - * Scans active documents over the autocomplete field (or all search fields), - * scores each field row with the prefix scorer (`mode: 'prefix'`) or the - * fuzzy scorer (`mode: 'fuzzy'`; typo-tolerant substring when - * `fuzzyDistance > 0`), ranks with the deterministic comparator, and - * returns the top `limit` completions with Unicode-safe highlight ranges. - * `prefix` yields `type: 'completion'` (including typo-tolerant prefix); - * `fuzzy` yields `type: 'did-you-mean'`. - * - * Suggestions are index-wide by design: search filters (structured or - * predicate) narrow search results but never suggestion candidates, so - * inline `search({ filter, autocomplete })` suggestions match standalone - * `autocomplete()` output for the same query and options. `search.fields` - * does scope suggestions unless the autocomplete spec carries an explicit - * `field`, which wins silently — pass `autocomplete.field` to pin the scan. - * Suggestion ranking uses `autocomplete.tieBreakers` when provided, else the - * inline search `ranking.tieBreakers`, else the default; granularity - * is per (doc, field) row (a document may appear multiple times), unlike - * per-document best-field-wins search ranking. - * - * Autocomplete uses default `prefixMatch` options and only `fuzzyDistance` for - * typo tolerance — the search `prefixMatch`/`typoTolerance` are ignored. - * - * Latency is a single O(docs x fields) scan plus a full sort of the match - * set (unbounded pre-truncation by design; highlights are bounded to - * top-K). Measured ~35ms prefix / ~95ms fuzzy+d2 over 50k rows (25k docs - * x 2 fields); inline `search({ autocomplete })` pays both scans (~2x). There - * is no `AbortSignal` support on this path; per-keystroke callers should - * debounce. Queries longer than `QUERY_TOKENS_MAX` throw `QueryTooLongError`. - */ async autocomplete( query: string, options: AutocompleteOptions = {} ): Promise> { - if (this.isDestroyed) { + if (this.state.isDestroyed) { throw new Error('[webgpu-search] DocumentIndex has been destroyed.'); } - if (typeof query !== 'string') { - throw new TypeError(`[webgpu-search] autocomplete expects query: string, got ${typeof query}.`); - } - const spec = normalizeAutocompleteOptions(options); - if (spec.field !== undefined && !this.fieldNameToIndex.has(spec.field)) { - throw new Error(`[webgpu-search] Unknown autocomplete field: "${spec.field}".`); - } - const t0 = nowMs(); - const normalized = normalizeText(query, this.normalized); - if (normalized.tokenCount > QUERY_TOKENS_MAX) { - throw new QueryTooLongError(QUERY_TOKENS_MAX, normalized.tokenCount, this.profileId); - } - if (normalized.isEmpty || this.idToDocIndex.size === 0) { - return { suggestions: [], queryDurationMs: nowMs() - t0 }; - } - const allowed = spec.field !== undefined - ? new Set([this.fieldNameToIndex.get(spec.field) as number]) - : undefined; - const suggestions = this.computeSuggestions(query, normalized.tokens, spec, allowed); - return { suggestions, queryDurationMs: nowMs() - t0 }; - } - - /** - * Synchronous suggestion enumeration shared by `autocomplete()` and inline - * `search({ autocomplete })`. Callers pass pre-normalized query tokens and a - * validated spec; `searchFields` scopes the scan when the autocomplete spec - * carries no explicit field. - */ - private computeSuggestions( - query: string, - queryTokens: Uint32Array, - spec: NormalizedAutocompleteOptions, - searchFields?: ReadonlySet - ): SuggestionItem[] { - let fieldIndices: number[]; - if (spec.field !== undefined) { - fieldIndices = [this.fieldNameToIndex.get(spec.field) as number]; - } else if (searchFields !== undefined) { - fieldIndices = [...searchFields]; - } else { - fieldIndices = this.sortedFields.map((_, i) => i); - } - - const typoForSuggest = normalizeTypoTolerance( - spec.fuzzyDistance > 0 - ? { enabled: true, maxDistance: spec.fuzzyDistance as 1 | 2 } - : undefined - ); - const prefixOpts = normalizePrefixOptions(undefined); - - interface SuggestHit { - keys: RankableCandidate; - text: string; - dIdx: number; - fIdx: number; - rawStr: string; - } - const candidates: SuggestHit[] = []; - - for (let d = 0; d < this.records.length; d++) { - const doc = this.records[d]; - if (doc === null || doc === undefined) continue; - const id = this.docIds[d]; - if (id === null || id === undefined) continue; - const rows = this.docToRowIndices[d]; - if (!rows) continue; - for (let k = 0; k < fieldIndices.length; k++) { - const fIdx = fieldIndices[k] as number; - const rowIdx = rows[fIdx]; - if (rowIdx === undefined) continue; - if (this.tombstones.has(rowIdx)) continue; - const rowTok = this.rowTokens[rowIdx]; - if (!rowTok || rowTok.length === 0) continue; - const rawStr = this.rawFieldStrings[d]?.[fIdx] ?? ''; - if (rawStr.length === 0) continue; - - let matched = false; - let rawScore = 0; - if (spec.mode === 'prefix') { - const r = scorePrefixTokens(rowTok, queryTokens, prefixOpts, typoForSuggest); - matched = r.matched; - rawScore = r.score; - } else if (spec.fuzzyDistance > 0) { - const r = scoreSubstringTypoTokens(rowTok, queryTokens, typoForSuggest); - matched = r.matched; - rawScore = r.score; - } else { - const r = scoreFuzzyTokens(rowTok, queryTokens); - matched = r.matched; - rawScore = r.score; - } - if (!matched) continue; - const weight = this.sortedFields[fIdx].weight; - candidates.push({ - keys: { - score: Math.round(rawScore * weight), - fieldWeight: weight, - isExactMatch: isExactTokenMatch(rowTok, queryTokens), - matchedLength: rowTok.length, - id, - docIndex: d - }, - text: rawStr, - dIdx: d, - fIdx, - rawStr - }); - } - } - - candidates.sort((a, b) => compareRanked(a.keys, b.keys, spec.tieBreakers)); - const top = candidates.length > spec.limit ? candidates.slice(0, spec.limit) : candidates; - const out: SuggestionItem[] = []; - for (let i = 0; i < top.length; i++) { - const c = top[i] as SuggestHit; - const highlightMode = spec.mode === 'prefix' ? 'prefix' : (spec.fuzzyDistance > 0 ? 'substring' : 'fuzzy'); - let matchedRanges: HighlightRange[] = []; - try { - matchedRanges = alignHighlights(c.rawStr, query, { - mode: highlightMode, - normalized: this.normalized, - queryTokens, - typoTolerance: spec.fuzzyDistance > 0 - ? { enabled: true, maxDistance: spec.fuzzyDistance as 1 | 2 } - : undefined - }); - } catch (err) { - // Only tolerate expected highlight-alignment failures for a single - // candidate (scorer/highlight length gates, polarity mismatch); - // programmer errors and I/O faults must propagate. - if ( - err instanceof IncompatibleOptionError || - err instanceof ProfileMismatchError || - err instanceof RangeError - ) { - matchedRanges = []; - } else { - throw err; - } - } - out.push({ - text: c.text, - score: c.keys.score, - type: spec.mode === 'prefix' ? 'completion' : 'did-you-mean', - matchedRanges, - docId: c.keys.id, - doc: this.records[c.dIdx] as TDoc - }); - } - return out; - } - - /** - * (): aggregates facet buckets over an explicit query-matched - * doc set with disjunctive filter exclusion. Each facet ignores structured - * filter clauses on its own field (`excludeFieldFromFilter`) while keeping - * all other clauses; function predicates stay conjunctive. `or`/`not` - * branches referencing the facet field are kept verbatim (conservative, - * aliased). Cost is O(F*n) with F bitset allocs plus an optional extra - * unfiltered parity scan when a structured filter is present. - */ - private buildFacetResults( - facetEngine: FacetEngine, - facetSpec: NormalizedFacet[], - queryMatched: Iterable, - structuredFilter: FilterExpression | undefined, - filterPredicate: ((doc: TDoc) => boolean) | undefined, - isApproximate: boolean - ): Record { - const base: number[] = []; - for (const dIdx of queryMatched) { - const doc = this.records[dIdx]; - if (doc === null || doc === undefined) continue; - if (filterPredicate && !filterPredicate(doc)) continue; - base.push(dIdx); - } - const out: Record = {}; - for (const facet of facetSpec) { - let candidates: number[] = base; - if (structuredFilter !== undefined) { - const excluded = excludeFieldFromFilter(structuredFilter, facet.request.field); - if (excluded !== undefined) { - const mask = compileFilter(excluded, this.columnarStore); - candidates = base.filter((d) => mask.has(d)); - } - } - const result = facetEngine.aggregateOne(facet, candidates, isApproximate); - if (facet.name === '__proto__') { - Object.defineProperty(out, facet.name, { value: result, enumerable: true, configurable: true, writable: true }); - } else { - out[facet.name] = result; - } - } - return out; - } - - /** - * deterministic re-sort after `scoringHook` boosts. - * Reconstructs rank keys from index state (field weight, exactness via - * `isExactTokenMatch` on the winning row, matched length, id, docIndex) - * with updated scores, then sorts via `compareRanked` + `tieBreakers`. - * Preserves the total-order contract (deterministic across engines). - */ - private resortResultsAfterScoring( - results: DocumentSearchResultItem[], - queryTokens: Uint32Array, - tieBreakers: TieBreakerCriterion[] - ): void { - if (results.length <= 1) return; - const keys: RankableCandidate[] = new Array(results.length); - for (let i = 0; i < results.length; i++) { - const item = results[i] as DocumentSearchResultItem; - const dIdx = this.idToDocIndex.get(item.id); - const fIdx = this.fieldNameToIndex.get(item.matchedField); - if (dIdx !== undefined && fIdx !== undefined) { - const rowIdx = this.docToRowIndices[dIdx]?.[fIdx]; - const rowTok = rowIdx !== undefined ? this.rowTokens[rowIdx] : undefined; - const weight = this.sortedFields[fIdx]?.weight ?? 1.0; - keys[i] = { - score: item.score, - fieldWeight: weight, - isExactMatch: rowTok !== undefined ? isExactTokenMatch(rowTok, queryTokens) : false, - matchedLength: rowTok !== undefined ? rowTok.length : 0, - id: item.id, - docIndex: dIdx - }; - } else { - // Defensive fallback (unreachable in normal flow: resort runs before - // postProcess, so IDs/fields resolve). Positional docIndex preserves - // the total-order guarantee via unique positions. - keys[i] = { - score: item.score, - fieldWeight: 1.0, - isExactMatch: false, - matchedLength: 0, - id: item.id, - docIndex: i - }; - } - } - const order = results.map((_, i) => i); - order.sort((ia, ib) => compareRanked(keys[ia] as RankableCandidate, keys[ib] as RankableCandidate, tieBreakers)); - const sorted = order.map((i) => results[i] as DocumentSearchResultItem); - for (let i = 0; i < results.length; i++) { - results[i] = sorted[i] as DocumentSearchResultItem; - } - } - - private enrichHighlights( - results: DocumentSearchResultItem[], - query: string, - mode: SearchMode, - options: DocumentSearchOptions, - customTokenTerms?: Uint32Array[], - precomputedTokens?: Uint32Array - ): void { - const shouldHighlight = options.highlightOptions?.highlight ?? options.highlight ?? true; - if (!shouldHighlight) return; - - const tag = options.highlightOptions?.tag ?? options.tag; - const highlightFieldsRule = options.highlightOptions?.fields ?? options.highlightFields ?? 'all-matched'; - const escapeHtml = options.highlightOptions?.escapeHtml ?? options.escapeHtml ?? false; - - if ( - typeof highlightFieldsRule === 'string' && - highlightFieldsRule !== 'all-matched' && - highlightFieldsRule !== 'matched-field' && - highlightFieldsRule !== 'all-fields' - ) { - throw new TypeError( - `[webgpu-search] Invalid highlightFields option: "${highlightFieldsRule}". Must be 'all-matched', 'matched-field', 'all-fields', or an array of field names.` - ); - } - - const queryTokens = precomputedTokens ?? normalizeText(query, this.normalized).tokens; - // thread token/prefix/typo options so highlight ranges track - // the serving scorer (score-highlight symmetry across all modes). - // custom tokenizer terms replace the default split when supplied. - const alignOpts = { - mode, - normalized: this.normalized, - queryTokens, - tokenMatch: options.tokenMatch, - prefixMatch: options.prefixMatch, - typoTolerance: options.typoTolerance, - ...(customTokenTerms !== undefined ? { tokenTermsOverride: customTokenTerms } : {}) - }; - - for (let i = 0; i < results.length; i++) { - const item = results[i]; - const dIdx = this.idToDocIndex.get(item.id); - if (dIdx === undefined) continue; - - const highlights: Record = {}; - const highlightedText: Record = {}; - - if (highlightFieldsRule === 'all-fields') { - for (let f = 0; f < this.sortedFields.length; f++) { - const fieldDef = this.sortedFields[f]; - const rawStr = this.rawFieldStrings[dIdx][f] ?? ''; - const ranges = alignHighlights(rawStr, query, alignOpts); - if (ranges.length > 0) { - setOwnKey(highlights, fieldDef.name, ranges); - if (tag) { - setOwnKey(highlightedText, fieldDef.name, renderHighlightedText(rawStr, ranges, tag, escapeHtml)); - } - } - } - } else if (Array.isArray(highlightFieldsRule)) { - for (let k = 0; k < highlightFieldsRule.length; k++) { - const fName = highlightFieldsRule[k]; - const fIdx = this.fieldNameToIndex.get(fName); - if (fIdx !== undefined) { - const rawStr = this.rawFieldStrings[dIdx][fIdx] ?? ''; - const ranges = alignHighlights(rawStr, query, alignOpts); - if (ranges.length > 0) { - setOwnKey(highlights, fName, ranges); - if (tag) { - setOwnKey(highlightedText, fName, renderHighlightedText(rawStr, ranges, tag, escapeHtml)); - } - } - } - } - } else { - // 'matched-field' or 'all-matched' (default) - const primaryFIdx = this.fieldNameToIndex.get(item.matchedField); - if (primaryFIdx !== undefined) { - const primaryRaw = this.rawFieldStrings[dIdx][primaryFIdx] ?? ''; - const ranges = alignHighlights(primaryRaw, query, alignOpts); - setOwnKey(highlights, item.matchedField, ranges); - if (tag) { - setOwnKey(highlightedText, item.matchedField, renderHighlightedText(primaryRaw, ranges, tag, escapeHtml)); - } - } - - if (highlightFieldsRule === 'all-matched' && item.matches) { - for (let m = 0; m < item.matches.length; m++) { - const aux = item.matches[m]; - const auxFIdx = this.fieldNameToIndex.get(aux.field); - if (auxFIdx !== undefined) { - const auxRaw = this.rawFieldStrings[dIdx][auxFIdx] ?? ''; - const ranges = alignHighlights(auxRaw, query, alignOpts); - aux.highlights = ranges; - setOwnKey(highlights, aux.field, ranges); - if (tag) { - setOwnKey(highlightedText, aux.field, renderHighlightedText(auxRaw, ranges, tag, escapeHtml)); - } - } - } - } - } - - // Synchronize auxiliary match highlights across all modes if highlights exist - if (item.matches) { - for (let m = 0; m < item.matches.length; m++) { - const aux = item.matches[m]; - if (Object.prototype.hasOwnProperty.call(highlights, aux.field)) { - aux.highlights = highlights[aux.field]; - } - } - } - - item.highlights = highlights; - if (tag) { - item.highlightedText = highlightedText; - } - } - } - - private validateDocumentId(id: unknown, context: string): void { - if ( - id === null || - id === undefined || - (typeof id !== 'string' && typeof id !== 'number') || - (typeof id === 'number' && !Number.isFinite(id)) - ) { - throw new TypeError( - `[webgpu-search] Invalid document ID in ${context}: ${String(id)}. DocumentId must be a non-empty string or finite number.` - ); - } - if (typeof id === 'string' && id.length === 0) { - throw new TypeError( - `[webgpu-search] Document ID in ${context} must not be an empty string.` - ); - } - } - - private prepareDocFields(doc: TDoc, id: DocumentId): { - id: DocumentId; - doc: TDoc; - rawStrings: string[]; - tokens: Uint32Array[]; - tokenCount: number; - } { - const fieldCount = this.sortedFields.length; - const rawStrings = new Array(fieldCount); - const tokens = new Array(fieldCount); - let tokenCount = 0; - - for (let f = 0; f < fieldCount; f++) { - const field = this.sortedFields[f]; - const rawVal = field.getter(doc); - let rawStr = ''; - if (rawVal !== null && rawVal !== undefined) { - if (Array.isArray(rawVal)) { - rawStr = rawVal.filter((x) => x !== null && x !== undefined).join(' '); - } else { - rawStr = typeof rawVal === 'string' ? rawVal : String(rawVal); - } - } - rawStrings[f] = rawStr; - const norm = normalizeText(rawStr, this.normalized); - tokens[f] = norm.tokens; - tokenCount += norm.tokenCount; - } - - return { id, doc, rawStrings, tokens, tokenCount }; - } - - private queued(fn: () => Promise): Promise { - const p = this.searchMutex.then(fn, fn); - this.searchMutex = p.catch(() => {}); - return p; + return runAutocomplete(this.state, query, options); } async add(docs: TDoc | TDoc[], options?: AddOptions): Promise { @@ -2161,467 +166,53 @@ export class DocumentIndex> { } async applyBatch(batch: MutationBatch, options?: AddOptions): Promise { - if (this.isDestroyed) { + if (this.state.isDestroyed) { throw new Error('[webgpu-search] DocumentIndex has been destroyed.'); } if (!batch || typeof batch !== 'object') { throw new TypeError('[webgpu-search] applyBatch expects batch object.'); } - return this.queued(async () => { - if (this.isDestroyed) { - throw new Error('[webgpu-search] DocumentIndex has been destroyed.'); - } - const tStart = nowMs(); - const upsert = options?.upsert ?? false; - - // PHASE 1: ATOMIC TWO-PHASE VALIDATION (State remains untouched on error) - const removeIds: DocumentId[] = []; - const willRemoveSet = new Set(); - if (batch.remove !== undefined) { - if (!Array.isArray(batch.remove)) { - throw new TypeError('[webgpu-search] batch.remove must be an array of document IDs.'); - } - for (let i = 0; i < batch.remove.length; i++) { - const id = batch.remove[i]; - this.validateDocumentId(id, `remove at index ${i}`); - removeIds.push(id); - if (this.idToDocIndex.has(id)) { - willRemoveSet.add(id); - } - } - } - - interface PreparedDocUpdate { - id: DocumentId; - doc: TDoc; - rawStrings: string[]; - tokens: Uint32Array[]; - tokenCount: number; - columnar: PreparedColumnarRow; - } - const preparedUpdates: PreparedDocUpdate[] = []; - const updatedIdsInBatch = new Set(); - if (batch.update !== undefined) { - if (!Array.isArray(batch.update)) { - throw new TypeError('[webgpu-search] batch.update must be an array of documents.'); - } - for (let i = 0; i < batch.update.length; i++) { - const doc = batch.update[i]; - if (doc === null || typeof doc !== 'object') { - throw new TypeError(`[webgpu-search] batch.update document at index ${i} must be an object.`); - } - const id = this.getId(doc); - this.validateDocumentId(id, `update at index ${i}`); - if (!this.idToDocIndex.has(id) || willRemoveSet.has(id)) { - throw new DocumentNotFoundError(id); - } - if (updatedIdsInBatch.has(id)) { - throw new DuplicateIdError(id); - } - updatedIdsInBatch.add(id); - // Filter getters run here (before any mutation) so a throwing - // getter rejects the whole batch with the index untouched. - const prepared = this.prepareDocFields(doc, id); - preparedUpdates.push({ ...prepared, columnar: this.columnarStore.prepare(doc) }); - } - } - - interface PreparedDocAdd { - id: DocumentId; - doc: TDoc; - rawStrings: string[]; - tokens: Uint32Array[]; - tokenCount: number; - columnar: PreparedColumnarRow; - } - const preparedAdds: PreparedDocAdd[] = []; - const addedIdsInBatch = new Set(); - - if (batch.add !== undefined) { - if (!Array.isArray(batch.add)) { - throw new TypeError('[webgpu-search] batch.add must be an array of documents.'); - } - for (let i = 0; i < batch.add.length; i++) { - const doc = batch.add[i]; - if (doc === null || typeof doc !== 'object') { - throw new TypeError(`[webgpu-search] batch.add document at index ${i} must be an object.`); - } - const id = this.getId(doc); - this.validateDocumentId(id, `add at index ${i}`); - - const existsInIndex = this.idToDocIndex.has(id) && !willRemoveSet.has(id); - const existsInBatch = addedIdsInBatch.has(id); - - if (existsInIndex || existsInBatch) { - if (!upsert) { - throw new DuplicateIdError(id); - } - } - addedIdsInBatch.add(id); - - const prepared = this.prepareDocFields(doc, id); - preparedAdds.push({ ...prepared, columnar: this.columnarStore.prepare(doc) }); - } - } - - // PHASE 2: DETERMINISTIC EXECUTION (remove -> update -> add) - this.generation++; - let removedCount = 0; - let updatedCount = 0; - let addedCount = 0; - - // A. Remove - for (let i = 0; i < removeIds.length; i++) { - const id = removeIds[i]; - const d = this.idToDocIndex.get(id); - if (d !== undefined) { - const rows = this.docToRowIndices[d]; - if (rows) { - for (let f = 0; f < rows.length; f++) { - this.tombstones.add(rows[f]); - } - } - this.idToDocIndex.delete(id); - this.records[d] = null as any; - this.docIds[d] = null as any; - this.rawFieldStrings[d] = null as any; - this.docToRowIndices[d] = null as any; - this.columnarStore.remove(d); - removedCount++; - } - } - - // B. Update - const newRowTokensToAppend: Uint32Array[] = []; - const newRowDocIndicesToAppend: number[] = []; - const newRowFieldIndicesToAppend: number[] = []; - let newTokensToAppendCount = 0; - - for (let i = 0; i < preparedUpdates.length; i++) { - const upd = preparedUpdates[i]; - const oldD = this.idToDocIndex.get(upd.id)!; - const oldRows = this.docToRowIndices[oldD]; - if (oldRows) { - for (let f = 0; f < oldRows.length; f++) { - this.tombstones.add(oldRows[f]); - } - } - this.records[oldD] = null as any; - this.docIds[oldD] = null as any; - this.rawFieldStrings[oldD] = null as any; - this.docToRowIndices[oldD] = null as any; - this.columnarStore.remove(oldD); - - const newD = this.records.length; - this.records.push(upd.doc); - this.docIds.push(upd.id); - this.idToDocIndex.set(upd.id, newD); - this.rawFieldStrings.push(upd.rawStrings); - this.columnarStore.addPrepared(newD, upd.columnar); - - const docRows: number[] = new Array(this.sortedFields.length); - for (let f = 0; f < this.sortedFields.length; f++) { - const rowIdx = this.rowTokens.length + newRowTokensToAppend.length; - newRowTokensToAppend.push(upd.tokens[f]); - newRowDocIndicesToAppend.push(newD); - newRowFieldIndicesToAppend.push(f); - docRows[f] = rowIdx; - newTokensToAppendCount += upd.tokens[f].length; - } - this.docToRowIndices.push(docRows); - updatedCount++; - } - - // C. Add - for (let i = 0; i < preparedAdds.length; i++) { - const ad = preparedAdds[i]; - if (this.idToDocIndex.has(ad.id)) { - const oldD = this.idToDocIndex.get(ad.id)!; - const oldRows = this.docToRowIndices[oldD]; - if (oldRows) { - for (let f = 0; f < oldRows.length; f++) { - this.tombstones.add(oldRows[f]); - } - } - this.records[oldD] = null as any; - this.docIds[oldD] = null as any; - this.rawFieldStrings[oldD] = null as any; - this.docToRowIndices[oldD] = null as any; - this.columnarStore.remove(oldD); - updatedCount++; - } else { - addedCount++; - } - - const newD = this.records.length; - this.records.push(ad.doc); - this.docIds.push(ad.id); - this.idToDocIndex.set(ad.id, newD); - this.rawFieldStrings.push(ad.rawStrings); - this.columnarStore.addPrepared(newD, ad.columnar); - - const docRows: number[] = new Array(this.sortedFields.length); - for (let f = 0; f < this.sortedFields.length; f++) { - const rowIdx = this.rowTokens.length + newRowTokensToAppend.length; - newRowTokensToAppend.push(ad.tokens[f]); - newRowDocIndicesToAppend.push(newD); - newRowFieldIndicesToAppend.push(f); - docRows[f] = rowIdx; - newTokensToAppendCount += ad.tokens[f].length; - } - this.docToRowIndices.push(docRows); - } - - // Append newly prepared rows to rowTokens, rowToDocIndex, rowToFieldIndex - if (newRowTokensToAppend.length > 0) { - for (let r = 0; r < newRowTokensToAppend.length; r++) { - this.rowTokens.push(newRowTokensToAppend[r]); - this.rowToDocIndex.push(newRowDocIndicesToAppend[r]); - this.rowToFieldIndex.push(newRowFieldIndicesToAppend[r]); - } - this.totalTokens += newTokensToAppendCount; - } - - // D. Check Compaction / GPU Headroom - const totalRows = this.rowTokens.length; - const tombstoneCount = this.tombstones.size; - const tombstoneRatio = totalRows > 0 ? tombstoneCount / totalRows : 0; - - let compacted = false; - const isGpu = this.engineState.engine === 'webgpu' && this.gpuEngine !== null && this.gpuEngine.isReady; - const gpuHeadroomExhausted = isGpu && !this.gpuEngine!.canFitHeadroom(totalRows, this.totalTokens); - - const activeDocCount = this.idToDocIndex.size; - const newCandidateCap = Math.min(32768, Math.max(8192, Math.floor(activeDocCount * this.sortedFields.length * 0.1))); - if (newCandidateCap > this.candidateCapacity) { - this.candidateCapacity = newCandidateCap; - this.gpuEngine?.ensureCandidateCapacity(newCandidateCap); - } - - try { - if (tombstoneRatio >= 0.25 || gpuHeadroomExhausted) { - await this.compact(); - compacted = true; - } else if (isGpu && newRowTokensToAppend.length > 0) { - await this.syncAppendedRowsToGpu(newRowTokensToAppend, totalRows, newTokensToAppendCount); - } else if (!isGpu && (this.preferGpu || totalRows >= (this.options.threshold ?? 30_000)) && this.options.preferGpu !== false && activeDocCount > 0) { - await this.tryInitializeGpuEngine(); - } - } catch (err) { - console.warn('[webgpu-search] GPU sync failed during mutation; falling back to CPU:', err); - if (this.gpuEngine) { - try { this.gpuEngine.destroy(); } catch {} - this.gpuEngine = null; - } - this.engineState = transitionEngineState(this.engineState, 'cpu', 'gpu-execution-error'); - this.vramAllocatedBytes = 0; - } - - if (this.engineState.engine === 'cpu' && !this.engineState.fallbackReason) { - this.engineState = transitionEngineState( - this.engineState, - 'cpu', - this.options.preferGpu === false ? 'prefer-cpu' : 'below-threshold', - ); - } - - this.mutationEpoch++; - const durationMs = nowMs() - tStart; - this.lastMutationTimeMs = durationMs; - - return { - added: addedCount, - updated: updatedCount, - removed: removedCount, - mutationEpoch: this.mutationEpoch, - compacted, - durationMs - }; - }); - } - - private async syncAppendedRowsToGpu( - newRows: Uint32Array[], - totalRows: number, - newTokensCount: number - ): Promise { - if (!this.gpuEngine || !this.gpuEngine.isReady) return; - const numNewRows = newRows.length; - const oldTotalTokens = this.totalTokens - newTokensCount; - - const newOffsets = new Uint32Array(numNewRows); - let runningOffset = oldTotalTokens; - for (let i = 0; i < numNewRows; i++) { - runningOffset += newRows[i].length; - newOffsets[i] = runningOffset; - } - - const newTokens = new Uint32Array(newTokensCount); - let pos = 0; - for (let i = 0; i < numNewRows; i++) { - newTokens.set(newRows[i], pos); - pos += newRows[i].length; - } - - await this.gpuEngine.appendRows(newTokens, newOffsets, totalRows); - this.vramAllocatedBytes = this.gpuEngine.vramAllocatedBytes; + return enqueue(this.state, () => applyBatchOps(this.state, batch, options)); } compactCpu(): void { - const activeDocCount = this.idToDocIndex.size; - const fieldCount = this.sortedFields.length; - const totalRows = activeDocCount * fieldCount; - - const activeOldDocIndices: number[] = []; - for (let d = 0; d < this.records.length; d++) { - const id = this.docIds[d]; - if (id !== null && id !== undefined && this.idToDocIndex.get(id) === d) { - activeOldDocIndices.push(d); - } - } - - const newRecords: TDoc[] = new Array(activeDocCount); - const newDocIds: DocumentId[] = new Array(activeDocCount); - const newIdToDocIndex = new Map(); - const newRawFieldStrings: string[][] = Array.from({ length: activeDocCount }, () => new Array(fieldCount)); - - for (let newD = 0; newD < activeDocCount; newD++) { - const oldD = activeOldDocIndices[newD]; - const doc = this.records[oldD]; - const id = this.docIds[oldD]; - newRecords[newD] = doc; - newDocIds[newD] = id; - newIdToDocIndex.set(id, newD); - for (let f = 0; f < fieldCount; f++) { - newRawFieldStrings[newD][f] = this.rawFieldStrings[oldD][f]; - } - } - - const newRowTokens: Uint32Array[] = new Array(totalRows); - const newRowToDocIndex: number[] = new Array(totalRows); - const newRowToFieldIndex: number[] = new Array(totalRows); - const newDocToRowIndices: number[][] = Array.from({ length: activeDocCount }, () => new Array(fieldCount)); - - let corpusTokens = 0; - let rowIdx = 0; - - for (let f = 0; f < fieldCount; f++) { - for (let newD = 0; newD < activeDocCount; newD++) { - const oldD = activeOldDocIndices[newD]; - const oldRowIdx = this.docToRowIndices[oldD][f]; - const tokens = this.rowTokens[oldRowIdx]; - newRowTokens[rowIdx] = tokens; - newRowToDocIndex[rowIdx] = newD; - newRowToFieldIndex[rowIdx] = f; - newDocToRowIndices[newD][f] = rowIdx; - corpusTokens += tokens.length; - rowIdx++; - } - } - - this.records = newRecords; - this.docIds = newDocIds; - this.idToDocIndex = newIdToDocIndex; - this.rawFieldStrings = newRawFieldStrings; - this.rowTokens = newRowTokens; - this.rowToDocIndex = newRowToDocIndex; - this.rowToFieldIndex = newRowToFieldIndex; - this.docToRowIndices = newDocToRowIndices; - this.totalTokens = corpusTokens; - const oldToNewDocIndexMap = new Map(); - for (let newD = 0; newD < activeDocCount; newD++) { - oldToNewDocIndexMap.set(activeOldDocIndices[newD], newD); - } - this.columnarStore.compact(oldToNewDocIndexMap, activeDocCount); - this.tombstones.clear(); - } - - private isStratified(): boolean { - const docCount = this.records.length; - if (docCount === 0) return true; - for (let r = 0; r < this.rowTokens.length; r++) { - if ( - this.rowToFieldIndex[r] !== Math.floor(r / docCount) || - this.rowToDocIndex[r] !== (r % docCount) - ) { - return false; - } - } - return true; + compactCpu(this.state); } compactSync(): void { - if (this.tombstones.size > 0 || !this.isStratified()) { - this.compactCpu(); - if (this.gpuEngine) { - this.gpuEngine.destroy(); - this.gpuEngine = null; - this.engineState = transitionEngineState( - this.engineState, - 'cpu', - this.options.preferGpu === false ? 'prefer-cpu' : 'below-threshold', - ); - } - } - } - - private async compact(): Promise { - this.compactCpu(); - - if (this.engineState.engine === 'webgpu' && this.gpuEngine && this.gpuEngine.isReady) { - const activeDocCount = this.idToDocIndex.size; - const totalRows = this.rowTokens.length; - if (activeDocCount === 0) { - const packed = packDataset([], { normalized: this.normalized }); - await this.gpuEngine.loadDataset(packed); - this.vramAllocatedBytes = this.gpuEngine.vramAllocatedBytes; - } else { - const packed = packDataset(this.rowTokens, { - normalized: this.normalized, - totalTokens: this.totalTokens - }); - const rowCap = Math.max(totalRows, 16); - const tokenCap = Math.max(this.totalTokens, 64); - await this.gpuEngine.loadDataset(packed, { - rowCapacity: rowCap, - tokenCapacity: tokenCap, - growthFactor: this.growthFactor - }); - this.vramAllocatedBytes = this.gpuEngine.vramAllocatedBytes; - } - } + compactSync(this.state); } getRecords(): TDoc[] { - return this.records.slice(); + return this.state.records.slice(); } getDocIds(): DocumentId[] { - return this.docIds.slice(); + return this.state.docIds.slice(); } getSortedFields(): InternalField[] { - return this.sortedFields; + return this.state.sortedFields; } getRowTokens(): Uint32Array[] { - return this.rowTokens; + return this.state.rowTokens; } getTotalTokens(): number { - return this.totalTokens; + return this.state.totalTokens; } getTombstoneCount(): number { - return this.tombstones.size; + return this.state.tombstones.size; } isNormalized(): boolean { - return this.normalized; + return this.state.normalized; } getProfileId(): TextProfileId { - return this.profileId; + return this.state.profileId; } getUnicodeVersion(): string { @@ -2633,11 +224,11 @@ export class DocumentIndex> { } getColumnarStore(): ColumnarStore { - return this.columnarStore; + return this.state.columnarStore; } getFilterFieldDefinitions(): FilterFieldDefinition[] { - return this.filterFieldDefinitions.slice(); + return this.state.filterFieldDefinitions.slice(); } /** @@ -2645,25 +236,25 @@ export class DocumentIndex> { * are shared with the index). */ getHooks(): SearchHooks | undefined { - return this.indexHooks === undefined ? undefined : { ...this.indexHooks }; + return this.state.indexHooks === undefined ? undefined : { ...this.state.indexHooks }; } serialize(options?: SerializeDocumentIndexOptions): ArrayBuffer { - if (this.isDestroyed) { + if (this.state.isDestroyed) { throw new Error('[webgpu-search] DocumentIndex has been destroyed.'); } return encodeSnapshot(this, options); } async restore(buffer: ArrayBuffer, options?: RestoreDocumentIndexOptions): Promise { - if (this.isDestroyed) { + if (this.state.isDestroyed) { throw new Error('[webgpu-search] DocumentIndex has been destroyed.'); } const t0 = nowMs(); - await this.queued(async () => { + await enqueue(this.state, async () => { const snapshot = decodeSnapshot(buffer, options); - await this.applySnapshotData(snapshot, options); - this.restoreTimeMs = nowMs() - t0; + await applySnapshotData(this.state, snapshot, options); + this.state.restoreTimeMs = nowMs() - t0; }); } @@ -2679,458 +270,13 @@ export class DocumentIndex> { options?: RestoreDocumentIndexOptions ): Promise> { const t0 = nowMs(); - // Fail-closed profile guards (never read under the wrong profile). - // `textProfile` only admits 'unicode-default'; any other value throws - // `ProfileMismatchError`. An explicit `caseSensitive` override must agree - // with the snapshot polarity (`!header.normalized`); otherwise the - // retained token stream would be queried under the wrong normalization. - const rawTextProfile = (options?.options as { textProfile?: unknown } | undefined)?.textProfile; - if (rawTextProfile !== undefined && rawTextProfile !== 'unicode-default') { - throw new ProfileMismatchError('unicode-default', rawTextProfile, 'textProfile'); - } - const overrideCase = options?.options?.caseSensitive; - if (overrideCase !== undefined && overrideCase === snapshot.header.normalized) { - throw new ProfileMismatchError(!snapshot.header.normalized, overrideCase, 'caseSensitive'); - } - // fail-closed custom filter-getter guard. Snapshot `hasGetter` - // entries require a matching getter override; otherwise restore would - // silently rebuild columnar via default `doc[name]` (presence cleared). - assertSnapshotFilterGettersSatisfied( - snapshot.schema.filterFields as Array<{ name: string; hasGetter?: boolean }> | undefined, - options?.options?.filterFields as Array | undefined - ); - const userFields = options?.options?.fields; - const userFieldMap = new Map(); - if (Array.isArray(userFields)) { - for (const uf of userFields) { - if (typeof uf === 'string') { - userFieldMap.set(uf, uf); - } else if (uf && typeof uf === 'object' && typeof uf.name === 'string') { - userFieldMap.set(uf.name, uf); - } - } - } - - const resolvedFields = snapshot.schema.fields.map((sf) => { - const uf = userFieldMap.get(sf.name); - const getter = typeof uf === 'object' && uf !== null && typeof uf.getter === 'function' - ? uf.getter - : undefined; - return { - name: sf.name, - weight: sf.weight, - getter - }; - }); - - const mergedOptions: DocumentIndexOptions = { - ...(options?.options as any), - fields: resolvedFields, - idField: options?.options?.idField ?? snapshot.schema.idField ?? 'id', - caseSensitive: options?.options?.caseSensitive ?? snapshot.schema.caseSensitive ?? !(snapshot.header.normalized), - device: (options?.device ?? options?.options?.device) as GPUDevice | undefined, - preferGpu: options?.options?.preferGpu ?? snapshot.schema.preferGpu, - threshold: options?.options?.threshold ?? snapshot.schema.threshold, - candidateCapacity: options?.options?.candidateCapacity ?? snapshot.schema.candidateCapacity, - initialCapacity: options?.options?.initialCapacity ?? snapshot.schema.initialCapacity, - growthFactor: options?.options?.growthFactor ?? snapshot.schema.growthFactor, - filterFields: options?.options?.filterFields ?? (snapshot.schema.filterFields as any), - hooks: options?.options?.hooks - }; - - // fail-closed hook restore guard. Snapshots recording hookIds - // require matching handlers via `options.options.hooks`; closures - // are never serialized, only declarative IDs. - assertHooksSatisfied( - snapshot.schema.hookIds, - normalizeSearchHooks(mergedOptions.hooks) - ); - + const mergedOptions = mergeSnapshotRestoreOptions(snapshot, options); const index = new DocumentIndex(mergedOptions); - await index.applySnapshotData(snapshot, options); - index.restoreTimeMs = nowMs() - t0; + await applySnapshotData(index.state, snapshot, options); + index.state.restoreTimeMs = nowMs() - t0; return index; } - private configureFromSchema( - schema: DocumentIndexSchema, - options?: RestoreDocumentIndexOptions - ): void { - const userFields = options?.options?.fields; - const userFieldMap = new Map(); - if (Array.isArray(userFields)) { - for (const uf of userFields) { - if (typeof uf === 'string') { - userFieldMap.set(uf, uf); - } else if (uf && typeof uf === 'object' && typeof uf.name === 'string') { - userFieldMap.set(uf.name, uf); - } - } - } - - const normalized: InternalField[] = schema.fields.map( - (f: { name: string; weight: number }, originalIndex: number) => { - const name = f.name; - if (typeof name !== 'string' || name.length === 0) { - throw new TypeError('[webgpu-search] Restored schema field requires a non-empty string name.'); - } - const weight = f.weight ?? 1.0; - if (typeof weight !== 'number' || !Number.isFinite(weight) || weight <= 0) { - throw new RangeError( - `[webgpu-search] Restored field "${name}" weight must be a positive finite number, got ${String(weight)}.` - ); - } - const uf = userFieldMap.get(name); - let getter: (doc: TDoc) => string | string[] | undefined | null; - if (uf && typeof uf === 'object' && typeof uf.getter === 'function') { - getter = uf.getter; - } else { - getter = (doc: any) => readOwnField(doc, name); - } - return { name, weight, getter, originalIndex }; - }); - - this.sortedFields = [...normalized].sort((a, b) => { - if (b.weight !== a.weight) return b.weight > a.weight ? 1 : -1; - if (a.originalIndex !== b.originalIndex) return a.originalIndex > b.originalIndex ? 1 : -1; - return 0; - }); - - this.fieldNameToIndex.clear(); - for (let i = 0; i < this.sortedFields.length; i++) { - this.fieldNameToIndex.set(this.sortedFields[i].name, i); - } - - if (typeof options?.options?.idField === 'function') { - this.getId = options.options.idField; - } else if (typeof options?.options?.idField === 'string') { - const prop = options.options.idField; - this.getId = (doc: any) => doc[prop]; - } else if (typeof schema.idField === 'string') { - const prop = schema.idField; - this.getId = (doc: any) => doc[prop]; - } else if (options?.options?.idField === undefined && schema.idField === undefined) { - this.getId = (doc: any) => doc.id; - } - } - - private async applySnapshotData( - snapshot: RestoredDocumentSnapshot, - options?: RestoreDocumentIndexOptions - ): Promise { - if ((options?.options as unknown as Record | undefined)?.extensions !== undefined) { - throw new IncompatibleOptionError( - 'hooks', - "[webgpu-search] extensions was removed; use hooks." - ); - } - // Fail-closed profile guards (never read under the wrong profile). - // The live index normalization is fixed at construction; a snapshot - // built under the opposite polarity must be rejected, not reinterpreted. - // An explicit `caseSensitive`/`textProfile` override must agree with both - // the live index and the snapshot. - const rawTextProfile = (options?.options as { textProfile?: unknown } | undefined)?.textProfile; - if (rawTextProfile !== undefined && rawTextProfile !== 'unicode-default') { - throw new ProfileMismatchError('unicode-default', rawTextProfile, 'textProfile'); - } - const snapshotCaseSensitive = !snapshot.header.normalized; - const liveCaseSensitive = !this.normalized; - if (snapshotCaseSensitive !== liveCaseSensitive) { - throw new ProfileMismatchError(liveCaseSensitive, snapshotCaseSensitive, 'caseSensitive'); - } - const overrideCaseApply = options?.options?.caseSensitive; - if (overrideCaseApply !== undefined && overrideCaseApply !== snapshotCaseSensitive) { - throw new ProfileMismatchError(snapshotCaseSensitive, overrideCaseApply, 'caseSensitive'); - } - // fail-closed hook restore guard for instance `restore()`. - // Restore-supplied handlers (if any) replace the live index hooks; - // otherwise the live hooks must satisfy the snapshot. - const restoreHooks = normalizeSearchHooks(options?.options?.hooks); - if (restoreHooks !== undefined) { - this.indexHooks = restoreHooks; - (this.options as { hooks?: SearchHooks }).hooks = restoreHooks; - } - assertHooksSatisfied(snapshot.schema.hookIds, this.indexHooks); - - // fail-closed custom filter-getter guard for instance restore. - // Live filter defs (plus any restore override) must supply getters for - // snapshot `hasGetter` fields; otherwise columnar rebuild would silently - // change semantics. When the override supplies the getters, adopt it so - // `columnarStore.init` uses the correct accessors. - { - const overrideFF = options?.options?.filterFields as - | Array any }> - | undefined; - const liveCustom = new Set(); - for (const ff of this.filterFieldDefinitions) { - if ((ff as { hasGetter?: boolean }).hasGetter === true) liveCustom.add(ff.name as string); - } - const overrideGetterNames = new Set(); - if (Array.isArray(overrideFF)) { - for (const uf of overrideFF) { - if (uf && typeof uf === 'object' && typeof (uf as { getter?: unknown }).getter === 'function') { - overrideGetterNames.add((uf as { name: string }).name); - } - } - } - const needed = ((snapshot.schema.filterFields as Array<{ name: string; hasGetter?: boolean }> | undefined) ?? []).filter( - (ff) => ff && ff.hasGetter === true - ); - for (const ff of needed) { - if (!liveCustom.has(ff.name) && !overrideGetterNames.has(ff.name)) { - throw new IncompatibleIndexError( - `filter-getter:${ff.name}`, - 'missing custom getter override for snapshot filter field (supply options.options.filterFields with getter)' - ); - } - } - if (Array.isArray(overrideFF) && overrideFF.length > 0) { - const normalized: FilterFieldDefinition[] = []; - const seen = new Set(); - for (const uf of overrideFF) { - if (typeof uf === 'string') { - if (seen.has(uf)) continue; - seen.add(uf); - normalized.push({ name: uf as never, getter: ((doc: any) => (doc as any)[uf]) as never }); - } else if (uf && typeof uf === 'object' && typeof (uf as { name?: unknown }).name === 'string') { - const o = uf as { name: string; type?: FilterFieldType; getter?: (doc: TDoc) => any }; - if (seen.has(o.name)) continue; - seen.add(o.name); - const hasCustom = typeof o.getter === 'function'; - normalized.push( - hasCustom - ? { name: o.name as never, type: o.type, getter: o.getter as never, hasGetter: true } - : { name: o.name as never, type: o.type, getter: ((doc: any) => (doc as any)[o.name]) as never } - ); - } - } - if (normalized.length > 0) { - this.filterFieldDefinitions = normalized; - this.columnarStore = new ColumnarStore(normalized, { - initialCapacity: this.initialCapacity, - growthFactor: this.growthFactor - }); - } - } - } - - this.configureFromSchema(snapshot.schema, options); - - if (this.unsubscribeDeviceLost) { - this.unsubscribeDeviceLost(); - this.unsubscribeDeviceLost = undefined; - } - if (this.gpuEngine) { - this.gpuEngine.destroy(); - this.gpuEngine = null; - } - this.tombstones.clear(); - this.mutationEpoch = - typeof snapshot.schema.mutationEpoch === 'number' && - Number.isFinite(snapshot.schema.mutationEpoch) && - snapshot.schema.mutationEpoch >= 0 - ? Math.floor(snapshot.schema.mutationEpoch) - : 0; - - const docCount = snapshot.header.docCount; - const fieldCount = this.sortedFields.length; - const totalRows = snapshot.header.rowCount; - - this.records = snapshot.records.slice(); - this.docIds = snapshot.docIds.slice(); - this.idToDocIndex.clear(); - - if (this.records.length > 0 && this.docIds.length === 0) { - this.docIds = new Array(docCount); - for (let i = 0; i < docCount; i++) { - const id = this.getId(this.records[i]); - this.docIds[i] = id; - this.idToDocIndex.set(id, i); - } - } else { - for (let i = 0; i < this.docIds.length; i++) { - this.idToDocIndex.set(this.docIds[i], i); - } - } - - // If both records and docIds are present, ensure exact ID correspondence alignment - if (this.records.length === docCount && this.docIds.length === docCount) { - const recMap = new Map(); - for (const rec of this.records) { - recMap.set(this.getId(rec), rec); - } - if (recMap.size === docCount) { - const aligned: TDoc[] = new Array(docCount); - let allFound = true; - for (let i = 0; i < docCount; i++) { - const rec = recMap.get(this.docIds[i]); - if (rec === undefined) { - allFound = false; - break; - } - aligned[i] = rec; - } - if (allFound) { - this.records = aligned; - } - } - } - - if (this.records.length === 0 && docCount > 0) { - const idProp = typeof this.options.idField === 'string' ? this.options.idField : 'id'; - this.records = new Array(docCount); - for (let i = 0; i < docCount; i++) { - this.records[i] = { [idProp]: this.docIds[i] } as any; - } - } - - this.columnarStore.init(this.records); - - const rawCap = this.options.candidateCapacity; - const validUserCap = typeof rawCap === 'number' && Number.isFinite(rawCap) && rawCap > 0 - ? Math.floor(rawCap) - : Math.floor(docCount * fieldCount * 0.1); - this.candidateCapacity = Math.min(32768, Math.max(8192, validUserCap)); - - this.rowTokens = new Array(totalRows); - this.rowToDocIndex = new Array(totalRows); - this.rowToFieldIndex = new Array(totalRows); - this.docToRowIndices = Array.from({ length: docCount }, () => new Array(fieldCount)); - this.rawFieldStrings = Array.from({ length: docCount }, () => new Array(fieldCount)); - - for (let r = 0; r < totalRows; r++) { - const start = snapshot.offsets[r]; - const end = snapshot.offsets[r + 1]; - this.rowTokens[r] = snapshot.tokens.subarray(start, end); - - const f = fieldCount > 0 && docCount > 0 ? Math.floor(r / docCount) : 0; - const d = docCount > 0 ? r % docCount : 0; - this.rowToDocIndex[r] = d; - this.rowToFieldIndex[r] = f; - if (this.docToRowIndices[d]) { - this.docToRowIndices[d][f] = r; - } - } - - this.totalTokens = snapshot.header.tokenCount; - - if (this.records.length === docCount && docCount > 0) { - for (let d = 0; d < docCount; d++) { - const doc = this.records[d]; - for (let f = 0; f < fieldCount; f++) { - const field = this.sortedFields[f]; - const rawVal = field.getter(doc); - let rawStr = ''; - if (rawVal !== null && rawVal !== undefined) { - if (Array.isArray(rawVal)) { - rawStr = rawVal.filter((x) => x !== null && x !== undefined).join(' '); - } else { - rawStr = typeof rawVal === 'string' ? rawVal : String(rawVal); - } - } - this.rawFieldStrings[d][f] = rawStr; - } - } - } - - this.engineState = transitionEngineState( - this.engineState, - 'cpu', - this.engineState.fallbackReason, - ); - this.vramAllocatedBytes = 0; - - if (docCount > 0) { - const threshold = this.options.threshold ?? 30_000; - const shouldAttemptGpu = this.options.preferGpu === false - ? false - : this.preferGpu || totalRows >= threshold; - - if (this.options.preferGpu === false) { - this.engineState = transitionEngineState(this.engineState, 'cpu', 'prefer-cpu'); - } else if (!this.preferGpu && totalRows < threshold) { - this.engineState = transitionEngineState(this.engineState, 'cpu', 'below-threshold'); - } - - if (shouldAttemptGpu) { - const measuredAvgBytes = totalRows > 0 ? (this.totalTokens * 4) / totalRows : 0; - const deviceToUse = options?.device ?? this.options.device; - const powerPreference = options?.options?.powerPreference ?? this.options.powerPreference; - assertValidPowerPreference(powerPreference); - const budget = checkMemoryBudget(totalRows, measuredAvgBytes, deviceToUse, this.candidateCapacity); - if (!budget.allowed) { - this.engineState = transitionEngineState(this.engineState, 'cpu', 'memory-budget-exceeded'); - } else { - try { - const gpu = new WebGPUEngine(); - const initialized = await gpu.init( - deviceToUse !== undefined || powerPreference !== undefined - ? { device: deviceToUse ?? undefined, powerPreference } - : undefined - ); - if (initialized && gpu.isReady) { - gpu.ensureCandidateCapacity(this.candidateCapacity); - const packed = packDataset(this.rowTokens, { - normalized: this.normalized, - totalTokens: this.totalTokens - }); - const effectiveCap = this.initialCapacity > 0 ? this.initialCapacity : Math.max(16, docCount); - const rowCap = effectiveCap * fieldCount; - const avgTokens = totalRows > 0 ? this.totalTokens / totalRows : 16; - const tokenCap = Math.max(Math.floor(rowCap * avgTokens), 64); - await gpu.loadDataset(packed, { - rowCapacity: rowCap, - tokenCapacity: tokenCap, - growthFactor: this.initialCapacity > 0 ? 1.0 : this.growthFactor - }); - - this.gpuEngine = gpu; - this.engineState = transitionEngineState(this.engineState, 'webgpu'); - this.vramAllocatedBytes = gpu.vramAllocatedBytes; - - this.subscribeDeviceLost(); - } else { - this.engineState = transitionEngineState(this.engineState, 'cpu', 'device-request-failed'); - } - } catch { - this.engineState = transitionEngineState( - this.engineState, - 'cpu', - typeof navigator === 'undefined' || !('gpu' in navigator) || !navigator.gpu - ? 'webgpu-unsupported' - : 'device-request-failed', - ); - } - } - } - } - - this.generation++; - } - - /** - * Single device-loss subscription (re-subscribes replace the old one so - * rebuild cycles never leak listeners). On loss the GPU engine is torn - * down, VRAM accounting is zeroed, and the state moves to CPU with - * `fallbackReason: 'device-lost'` — CPU corpora (`rowTokens`, - * `records`) are retained so `rebuildGpu()` can re-upload identically. - */ - private subscribeDeviceLost(): void { - if (this.unsubscribeDeviceLost) { - this.unsubscribeDeviceLost(); - this.unsubscribeDeviceLost = undefined; - } - this.unsubscribeDeviceLost = GpuDevicePool.onDeviceLost(() => { - console.warn('[webgpu-search] GPU device lost, falling back to CPU.'); - if (this.gpuEngine) { - this.gpuEngine.destroy(); - this.gpuEngine = null; - } - this.vramAllocatedBytes = 0; - this.engineState = transitionEngineState(this.engineState, 'cpu', 'device-lost'); - }); - } - /** * Explicit device-loss rebuild path (Phase 2 reliability proof). * @@ -3150,179 +296,56 @@ export class DocumentIndex> { * - Throws only when the index itself is destroyed. */ async rebuildGpu(options?: { device?: GPUDevice; powerPreference?: GPUPowerPreference }): Promise { - if (this.isDestroyed) { - throw new Error('[webgpu-search] DocumentIndex has been destroyed.'); - } - if (this.gpuEngine !== null && this.gpuEngine.isReady && this.engineState.engine === 'webgpu') { - return true; - } - if (this.options.preferGpu === false) { - return false; - } - if (this.idToDocIndex.size === 0) { - return false; - } - if (options?.device !== undefined) { - (this.options as { device?: GPUDevice }).device = options.device; - } - if (options?.powerPreference !== undefined) { - (this.options as { powerPreference?: GPUPowerPreference }).powerPreference = options.powerPreference; - } - if (this.gpuEngine) { - try { this.gpuEngine.destroy(); } catch {} - this.gpuEngine = null; - } - this.vramAllocatedBytes = 0; - await this.tryInitializeGpuEngine(); - return this.engineState.engine === 'webgpu'; - } - - private async tryInitializeGpuEngine(): Promise { - const totalRows = this.rowTokens.length; - const threshold = this.options.threshold ?? 30_000; - const shouldAttempt = this.options.preferGpu === false - ? false - : this.preferGpu || totalRows >= threshold; - - if (!shouldAttempt) { - this.engineState = transitionEngineState( - this.engineState, - 'cpu', - this.options.preferGpu === false ? 'prefer-cpu' : 'below-threshold', - ); - return; - } - - const measuredAvgBytes = totalRows > 0 ? (this.totalTokens * 4) / totalRows : 0; - const budget = checkMemoryBudget(totalRows, measuredAvgBytes, this.options.device, this.candidateCapacity); - if (!budget.allowed) { - this.engineState = transitionEngineState(this.engineState, 'cpu', 'memory-budget-exceeded'); - return; - } - - try { - const gpu = new WebGPUEngine(); - const initialized = await gpu.init( - this.options.device !== undefined || this.options.powerPreference !== undefined - ? { device: this.options.device, powerPreference: this.options.powerPreference } - : undefined - ); - if (initialized && gpu.isReady) { - gpu.ensureCandidateCapacity(this.candidateCapacity); - const packed = packDataset(this.rowTokens, { - normalized: this.normalized, - totalTokens: this.totalTokens - }); - const effectiveCap = this.initialCapacity > 0 ? this.initialCapacity : Math.max(16, this.records.length); - const rowCap = effectiveCap * this.sortedFields.length; - const avgTokens = totalRows > 0 ? this.totalTokens / totalRows : 16; - const tokenCap = Math.max(Math.floor(rowCap * avgTokens), 64); - await gpu.loadDataset(packed, { - rowCapacity: rowCap, - tokenCapacity: tokenCap, - growthFactor: this.growthFactor - }); - this.gpuEngine = gpu; - this.engineState = transitionEngineState(this.engineState, 'webgpu'); - this.vramAllocatedBytes = gpu.vramAllocatedBytes; - this.subscribeDeviceLost(); - } else { - this.engineState = transitionEngineState( - this.engineState, - 'cpu', - typeof navigator === 'undefined' || !('gpu' in navigator) || !navigator.gpu - ? 'webgpu-unsupported' - : 'device-request-failed', - ); - } - } catch { - this.engineState = transitionEngineState( - this.engineState, - 'cpu', - typeof navigator === 'undefined' || !('gpu' in navigator) || !navigator.gpu - ? 'webgpu-unsupported' - : 'device-request-failed', - ); - } + return rebuildGpuEngine(this.state, options); } getStats(): DocumentIndexStats { - const adapter = this.gpuEngine?.adapterInfo; - const docCount = this.idToDocIndex.size; - const rowCount = this.rowTokens.length; - const tombstoneCount = this.tombstones.size; + const s = this.state; + const adapter = s.gpuEngine?.adapterInfo; + const docCount = s.idToDocIndex.size; + const rowCount = s.rowTokens.length; + const tombstoneCount = s.tombstones.size; const tombstoneRatio = rowCount > 0 ? tombstoneCount / rowCount : 0; - const tokenRamBytes = this.totalTokens * 4; + const tokenRamBytes = s.totalTokens * 4; const offsetRamBytes = (rowCount + 1) * 4; const ramBytes = tokenRamBytes + offsetRamBytes; return { size: docCount, - engine: this.engineState.engine, - vramAllocatedBytes: this.vramAllocatedBytes, + engine: s.engineState.engine, + vramAllocatedBytes: s.vramAllocatedBytes, adapterVendor: adapter?.vendor, adapterRenderer: adapter?.renderer, - profileId: this.profileId, + profileId: s.profileId, unicodeVersion: UNICODE_VERSION, scoringVersion: SCORING_VERSION, - tokenCount: this.totalTokens, - normalized: this.normalized, + tokenCount: s.totalTokens, + normalized: s.normalized, formatVersion: SNAPSHOT_FORMAT_VERSION, docCount, rowCount, tombstoneCount, tombstoneRatio, - buildTimeMs: this.buildTimeMs, - restoreTimeMs: this.restoreTimeMs, - lastMutationTimeMs: this.lastMutationTimeMs, - mutationEpoch: this.mutationEpoch, + buildTimeMs: s.buildTimeMs, + restoreTimeMs: s.restoreTimeMs, + lastMutationTimeMs: s.lastMutationTimeMs, + mutationEpoch: s.mutationEpoch, memory: { - vramBytes: this.vramAllocatedBytes, + vramBytes: s.vramAllocatedBytes, ramBytes, - totalBytes: this.vramAllocatedBytes + ramBytes, + totalBytes: s.vramAllocatedBytes + ramBytes, tokenRamBytes, offsetRamBytes }, - fallbackReason: this.engineState.fallbackReason + fallbackReason: s.engineState.fallbackReason }; } destroy(): void { - if (this.isDestroyed) { + if (this.state.isDestroyed) { return; } - this.isDestroyed = true; - this.generation++; - if (this.unsubscribeDeviceLost) { - this.unsubscribeDeviceLost(); - this.unsubscribeDeviceLost = undefined; - } - if (this.gpuEngine) { - this.gpuEngine.destroy(); - this.gpuEngine = null; - } - this.records = []; - this.docIds = []; - this.idToDocIndex.clear(); - this.rowTokens = []; - this.rowToDocIndex = []; - this.rowToFieldIndex = []; - this.docToRowIndices = []; - this.rawFieldStrings = []; - this.totalTokens = 0; - this.tombstones.clear(); - this.columnarStore = new ColumnarStore(this.filterFieldDefinitions, { - initialCapacity: this.initialCapacity, - growthFactor: this.growthFactor - }); - this.engineState = transitionEngineState( - this.engineState, - 'cpu', - this.engineState.fallbackReason, - ); - this.vramAllocatedBytes = 0; - const p = this.searchMutex.then(() => {}, () => {}); - this.searchMutex = p.catch(() => {}); + destroyIndex(this.state); } [Symbol.dispose](): void { diff --git a/packages/webgpu-search/src/document/build.ts b/packages/webgpu-search/src/document/build.ts new file mode 100644 index 0000000..36a77fe --- /dev/null +++ b/packages/webgpu-search/src/document/build.ts @@ -0,0 +1,1147 @@ +/** + * Document-index ingestion, mutation, engine-bootstrap, and compaction + * implementation (extracted verbatim from the pre-refactor monolithic + * `document-index.ts`). The public `DocumentIndex` facade delegates here. + * + * Scope of this module: + * - Option normalization for search fields, filter fields, and id accessors + * (construction path, fail-closed error order preserved). + * - Initial corpus ingestion (`initializeDocumentIndex`). + * - Two-phase atomic batch mutations (`applyBatchOps`). + * - CPU compaction + GPU (re)initialization and device-loss subscription. + * - Index destruction. + * + * All functions operate on the shared `DocumentIndexState` bundle and must + * only be invoked while the index-level invariants hold (e.g. mutations are + * serialized through the `searchMutex` by the facade). + */ + +import { readOwnField, ColumnarStore, type PreparedColumnarRow } from '../filtering/columnar-store'; +import { packDataset, checkMemoryBudget } from '../dataset-packing'; +import { normalizeText } from '../text-normalization'; +import { transitionEngineState, initialEngineState } from '../engine-state'; +import { WebGPUEngine } from '../webgpu-engine'; +import { GpuDevicePool } from '../gpu-device-pool'; +import { normalizeSearchHooks } from '../hooks'; +import { nowMs } from '../guard'; +import { DuplicateIdError, DocumentNotFoundError } from '../text-profile'; +import type { + AddOptions, + DocumentId, + DocumentIndexOptions, + FilterFieldDefinition, + FilterFieldType, + MutationBatch, + MutationResult, +} from '../types'; +import type { DocumentIndexState, InternalField } from './state'; + +/** Own-property write that is safe for arbitrary user field names: a plain + * `obj[key] = v` with key `"__proto__"` would replace the prototype instead + * of creating an entry. (Shared with query-side highlight/result writers.) */ +export function setOwnKey(obj: Record, key: string, value: V): void { + Object.defineProperty(obj, key, { value, enumerable: true, configurable: true, writable: true }); +} + +/** + * Derive the document-id accessor from `options.idField` (string property + * name or function). Called during construction; throws `TypeError` on a + * malformed idField exactly as the pre-refactor constructor did. + */ +export function deriveIdAccessor( + options: DocumentIndexOptions +): (doc: TDoc) => DocumentId { + if (typeof options.idField === 'function') { + return options.idField; + } else if (typeof options.idField === 'string') { + const prop = options.idField; + return (doc: any) => doc[prop]; + } else if (options.idField === undefined) { + return (doc: any) => doc.id; + } else { + throw new TypeError('[webgpu-search] idField must be a string property name or function.'); + } +} + +/** + * Normalize + deduplicate search-field definitions and stratify them + * (sorted descending by weight, ties broken by original declaration order). + * Fail-closed: empty names, invalid weights, duplicate names throw. + */ +export function normalizeIndexFields( + options: DocumentIndexOptions +): InternalField[] { + const seenNames = new Set(); + const normalized: InternalField[] = options.fields.map((f, originalIndex) => { + let name: string; + let weight = 1.0; + let getter: (doc: TDoc) => string | string[] | undefined | null; + + if (typeof f === 'string') { + if (f.trim().length === 0) { + throw new TypeError('[webgpu-search] Field name string must not be empty.'); + } + name = f; + getter = (doc: any) => readOwnField(doc, name); + } else if (f && typeof f === 'object') { + if (typeof f.name !== 'string' || f.name.length === 0) { + throw new TypeError('[webgpu-search] Field definition requires a non-empty string name.'); + } + name = f.name; + if (f.weight !== undefined) { + if (typeof f.weight !== 'number' || !Number.isFinite(f.weight) || f.weight <= 0) { + throw new RangeError( + `[webgpu-search] Field "${name}" weight must be a positive finite number, got ${f.weight}.` + ); + } + weight = f.weight; + } + if (f.getter !== undefined) { + if (typeof f.getter !== 'function') { + throw new TypeError(`[webgpu-search] Field "${name}" getter must be a function.`); + } + getter = f.getter; + } else { + getter = (doc: any) => readOwnField(doc, name); + } + } else { + throw new TypeError('[webgpu-search] Field must be a string or FieldDefinition object.'); + } + + if (seenNames.has(name)) { + throw new Error(`[webgpu-search] Duplicate field name: "${name}".`); + } + seenNames.add(name); + + return { name, weight, getter, originalIndex }; + }); + + // Field-stratified packing: sort fields descending by weight + return [...normalized].sort((a, b) => { + if (b.weight !== a.weight) return b.weight > a.weight ? 1 : -1; + if (a.originalIndex !== b.originalIndex) return a.originalIndex > b.originalIndex ? 1 : -1; + return 0; + }); +} + +/** + * Normalize + deduplicate filter-field definitions (name, type, getter). + * Fail-closed: empty names, invalid getters, duplicate names throw. + */ +export function normalizeFilterFields( + options: DocumentIndexOptions +): FilterFieldDefinition[] { + const normalizedFilterFields: FilterFieldDefinition[] = []; + if (options.filterFields !== undefined) { + if (!Array.isArray(options.filterFields)) { + throw new TypeError('[webgpu-search] DocumentIndex expects options.filterFields to be an array.'); + } + const seenFilterNames = new Set(); + for (let i = 0; i < options.filterFields.length; i++) { + const ff = options.filterFields[i]; + let name: string; + let type: FilterFieldType | undefined = undefined; + let getter: ((doc: TDoc) => any) | undefined = undefined; + let hasGetter = false; + + if (typeof ff === 'string') { + if (ff.trim().length === 0) { + throw new TypeError('[webgpu-search] Filter field name string must not be empty.'); + } + name = ff; + getter = (doc: any) => readOwnField(doc, name); + hasGetter = false; + } else if (ff && typeof ff === 'object') { + if (typeof ff.name !== 'string' || ff.name.trim().length === 0) { + throw new TypeError('[webgpu-search] Filter field definition requires a non-empty string name.'); + } + name = ff.name; + type = ff.type; + if (ff.getter !== undefined) { + if (typeof ff.getter !== 'function') { + throw new TypeError(`[webgpu-search] Filter field "${name}" getter must be a function.`); + } + getter = ff.getter; + hasGetter = true; + } else { + getter = (doc: any) => readOwnField(doc, name); + hasGetter = false; + } + } else { + throw new TypeError('[webgpu-search] Filter field must be a string or FilterFieldDefinition object.'); + } + + if (seenFilterNames.has(name)) { + throw new Error(`[webgpu-search] Duplicate filter field name: "${name}".`); + } + seenFilterNames.add(name); + normalizedFilterFields.push(hasGetter ? { name, type, getter, hasGetter: true } : { name, type, getter }); + } + } + return normalizedFilterFields; +} + +/** + * Construct the shared mutable-state bundle from validated options. + * Mirrors the pre-refactor constructor field-by-field: hooks are validated + * first, then capacity defaults, then the id accessor, then search fields, + * then filter fields — preserving the original error precedence. + */ +export function createDocumentIndexState( + options: DocumentIndexOptions +): DocumentIndexState { + // fail-closed extension hook validation at construction. + const indexHooks = normalizeSearchHooks(options.hooks); + + const initialCapacity = typeof options.initialCapacity === 'number' && Number.isFinite(options.initialCapacity) && options.initialCapacity > 0 + ? Math.floor(options.initialCapacity) + : 0; + const growthFactor = typeof options.growthFactor === 'number' && Number.isFinite(options.growthFactor) && options.growthFactor >= 1.0 + ? options.growthFactor + : 1.5; + + const normalized = !(options.caseSensitive ?? false); + const preferGpu = options.preferGpu ?? false; + const getId = deriveIdAccessor(options); + const sortedFields = normalizeIndexFields(options); + const filterFieldDefinitions = normalizeFilterFields(options); + + const fieldNameToIndex: Map = new Map(); + for (let i = 0; i < sortedFields.length; i++) { + fieldNameToIndex.set(sortedFields[i]!.name, i); + } + + const columnarStore = new ColumnarStore(filterFieldDefinitions, { + initialCapacity, + growthFactor + }); + + return { + options, + normalized, + profileId: 'unicode-default', + preferGpu, + initialCapacity, + growthFactor, + + records: [], + docIds: [], + idToDocIndex: new Map(), + getId, + + sortedFields, + fieldNameToIndex, + + rowTokens: [], + rowToDocIndex: [], + rowToFieldIndex: [], + docToRowIndices: [], + rawFieldStrings: [], + totalTokens: 0, + candidateCapacity: 8192, + buildTimeMs: 0, + lastMutationTimeMs: undefined, + mutationEpoch: 0, + tombstones: new Set(), + + searchMutex: Promise.resolve(), + generation: 0, + isDestroyed: false, + + columnarStore, + filterFieldDefinitions, + indexHooks, + + engineState: initialEngineState(), + gpuEngine: null, + vramAllocatedBytes: 0, + unsubscribeDeviceLost: undefined + }; +} + +/** + * Initial corpus ingestion (`DocumentIndex.create`). Builds the doc/id maps, + * the field-stratified row packing, the columnar store, and attempts the + * GPU engine when the corpus clears the configured routing thresholds. + */ +export async function initializeDocumentIndex(s: DocumentIndexState, records: TDoc[]): Promise { + const t0 = nowMs(); + const docCount = records.length; + const fieldCount = s.sortedFields.length; + const totalRows = docCount * fieldCount; + + s.records = records.slice(); + s.docIds = new Array(docCount); + s.idToDocIndex = new Map(); + + for (let i = 0; i < docCount; i++) { + const doc = s.records[i]; + if (doc === null || typeof doc !== 'object') { + throw new TypeError(`[webgpu-search] Document at index ${i} must be an object.`); + } + const id = s.getId(doc); + if ( + id === null || + id === undefined || + (typeof id !== 'string' && typeof id !== 'number') || + (typeof id === 'number' && !Number.isFinite(id)) + ) { + throw new TypeError( + `[webgpu-search] Document at index ${i} has invalid id: ${String(id)}. DocumentId must be a non-empty string or finite number.` + ); + } + if (typeof id === 'string' && id.length === 0) { + throw new TypeError( + `[webgpu-search] Document at index ${i} has an empty string id. DocumentId must not be empty.` + ); + } + if (s.idToDocIndex.has(id)) { + throw new DuplicateIdError(id); + } + s.docIds[i] = id; + s.idToDocIndex.set(id, i); + } + + // Dynamic Candidate Pool Scaling (Section 3.1): + // candidateCapacity = min(32768, max(8192, docCount * fieldCount * 0.1)) + const rawCap = s.options.candidateCapacity; + const validUserCap = typeof rawCap === 'number' && Number.isFinite(rawCap) && rawCap > 0 + ? Math.floor(rawCap) + : Math.floor(docCount * fieldCount * 0.1); + s.candidateCapacity = Math.min(32768, Math.max(8192, validUserCap)); + + // Field-Stratified Row Packing: + // Rows 0..N-1: Primary fields (highest weight). + // Rows N..2N-1: Secondary fields. + // Rows 2N..3N-1: Lower-weight fields. + s.rowTokens = new Array(totalRows); + s.rowToDocIndex = new Array(totalRows); + s.rowToFieldIndex = new Array(totalRows); + s.docToRowIndices = Array.from({ length: docCount }, () => new Array(fieldCount)); + s.rawFieldStrings = Array.from({ length: docCount }, () => new Array(fieldCount)); + + let corpusTokens = 0; + let rowIdx = 0; + + for (let f = 0; f < fieldCount; f++) { + const field = s.sortedFields[f]; + for (let d = 0; d < docCount; d++) { + const doc = s.records[d]; + const rawVal = field.getter(doc); + let rawStr = ''; + if (rawVal !== null && rawVal !== undefined) { + if (Array.isArray(rawVal)) { + rawStr = rawVal.filter((x) => x !== null && x !== undefined).join(' '); + } else { + rawStr = typeof rawVal === 'string' ? rawVal : String(rawVal); + } + } + s.rawFieldStrings[d][f] = rawStr; + const norm = normalizeText(rawStr, s.normalized); + s.rowTokens[rowIdx] = norm.tokens; + s.rowToDocIndex[rowIdx] = d; + s.rowToFieldIndex[rowIdx] = f; + s.docToRowIndices[d][f] = rowIdx; + corpusTokens += norm.tokenCount; + rowIdx++; + } + } + + s.totalTokens = corpusTokens; + s.columnarStore.init(s.records); + + if (docCount === 0) { + s.engineState = transitionEngineState( + s.engineState, + 'cpu', + s.options.preferGpu === false ? 'prefer-cpu' : 'below-threshold', + ); + s.buildTimeMs = nowMs() - t0; + return; + } + + const threshold = s.options.threshold ?? 30_000; + const shouldAttemptGpu = s.options.preferGpu === false + ? false + : s.preferGpu || totalRows >= threshold; + + if (s.options.preferGpu === false) { + s.engineState = transitionEngineState(s.engineState, 'cpu', 'prefer-cpu'); + } else if (!s.preferGpu && totalRows < threshold) { + s.engineState = transitionEngineState(s.engineState, 'cpu', 'below-threshold'); + } + + if (shouldAttemptGpu) { + const measuredAvgBytes = totalRows > 0 ? (corpusTokens * 4) / totalRows : 0; + const budget = checkMemoryBudget(totalRows, measuredAvgBytes, s.options.device, s.candidateCapacity); + if (!budget.allowed) { + console.warn(`[webgpu-search] ${budget.reason} Falling back to CPU.`); + s.engineState = transitionEngineState(s.engineState, 'cpu', 'memory-budget-exceeded'); + } else { + let gpu: WebGPUEngine | null = null; + try { + gpu = new WebGPUEngine(); + const initialized = await gpu.init( + s.options.device !== undefined || s.options.powerPreference !== undefined + ? { device: s.options.device, powerPreference: s.options.powerPreference } + : undefined + ); + + if (initialized && gpu.isReady) { + gpu.ensureCandidateCapacity(s.candidateCapacity); + const packed = packDataset(s.rowTokens, { + normalized: s.normalized, + totalTokens: corpusTokens + }); + const effectiveCap = s.initialCapacity > 0 ? s.initialCapacity : Math.max(16, docCount); + const rowCap = effectiveCap * fieldCount; + const avgTokens = totalRows > 0 ? corpusTokens / totalRows : 16; + const tokenCap = Math.max(Math.floor(rowCap * avgTokens), 64); + await gpu.loadDataset(packed, { + rowCapacity: rowCap, + tokenCapacity: tokenCap, + growthFactor: s.initialCapacity > 0 ? 1.0 : s.growthFactor + }); + + s.gpuEngine = gpu; + s.engineState = transitionEngineState(s.engineState, 'webgpu'); + s.vramAllocatedBytes = gpu.vramAllocatedBytes; + gpu = null; // Ownership transferred + + subscribeDeviceLost(s); + } else { + s.engineState = transitionEngineState( + s.engineState, + 'cpu', + typeof navigator === 'undefined' || !('gpu' in navigator) || !navigator.gpu + ? 'webgpu-unsupported' + : 'device-request-failed', + ); + } + } catch (gpuErr) { + console.warn('[webgpu-search] WebGPU initialization failed, falling back to CPU:', gpuErr); + s.engineState = transitionEngineState( + s.engineState, + 'cpu', + typeof navigator === 'undefined' || !('gpu' in navigator) || !navigator.gpu + ? 'webgpu-unsupported' + : 'device-request-failed', + ); + } finally { + if (gpu) { + gpu.destroy(); + } + } + } + } else { + s.engineState = transitionEngineState( + s.engineState, + 'cpu', + s.engineState.fallbackReason, + ); + } + + s.buildTimeMs = nowMs() - t0; +} + +/** + * Fail-closed document-id validation shared by batch mutation paths. + */ +export function validateDocumentId(id: unknown, context: string): void { + if ( + id === null || + id === undefined || + (typeof id !== 'string' && typeof id !== 'number') || + (typeof id === 'number' && !Number.isFinite(id)) + ) { + throw new TypeError( + `[webgpu-search] Invalid document ID in ${context}: ${String(id)}. DocumentId must be a non-empty string or finite number.` + ); + } + if (typeof id === 'string' && id.length === 0) { + throw new TypeError( + `[webgpu-search] Document ID in ${context} must not be an empty string.` + ); + } +} + +/** Field-token preparation result for a single document (pre-columnar). */ +export interface PreparedDocFields { + id: DocumentId; + doc: TDoc; + rawStrings: string[]; + tokens: Uint32Array[]; + tokenCount: number; +} + +/** A prepared document with its columnar pre-filter row attached. */ +export interface PreparedDocMutation extends PreparedDocFields { + columnar: PreparedColumnarRow; +} + +/** + * Extract raw strings + normalized tokens for every field of one document + * (mirrors the ingestion row-packing rule for arrays/scalars). + */ +export function prepareDocFields( + s: DocumentIndexState, + doc: TDoc, + id: DocumentId +): PreparedDocFields { + const fieldCount = s.sortedFields.length; + const rawStrings = new Array(fieldCount); + const tokens = new Array(fieldCount); + let tokenCount = 0; + + for (let f = 0; f < fieldCount; f++) { + const field = s.sortedFields[f]; + const rawVal = field.getter(doc); + let rawStr = ''; + if (rawVal !== null && rawVal !== undefined) { + if (Array.isArray(rawVal)) { + rawStr = rawVal.filter((x) => x !== null && x !== undefined).join(' '); + } else { + rawStr = typeof rawVal === 'string' ? rawVal : String(rawVal); + } + } + rawStrings[f] = rawStr; + const norm = normalizeText(rawStr, s.normalized); + tokens[f] = norm.tokens; + tokenCount += norm.tokenCount; + } + + return { id, doc, rawStrings, tokens, tokenCount }; +} + +/** + * Serialize a mutation through the index-wide search mutex. The returned + * promise resolves with the mutation result; the stored mutex tail never + * rejects (failures surface only through the returned promise). + */ +export function enqueue(s: DocumentIndexState, fn: () => Promise): Promise { + const p = s.searchMutex.then(fn, fn); + s.searchMutex = p.catch(() => {}); + return p; +} + +/** + * Two-phase atomic batch mutation (remove → update → add) executed with the + * mutation mutex already held. Phase 1 validates the whole batch and + * prepares all field/columnar work before touching state, so a throwing + * getter (or any validation failure) leaves the index untouched. + */ +export async function applyBatchOps( + s: DocumentIndexState, + batch: MutationBatch, + options?: AddOptions +): Promise { + if (s.isDestroyed) { + throw new Error('[webgpu-search] DocumentIndex has been destroyed.'); + } + const tStart = nowMs(); + const upsert = options?.upsert ?? false; + + // PHASE 1: ATOMIC TWO-PHASE VALIDATION (State remains untouched on error) + const removeIds: DocumentId[] = []; + const willRemoveSet = new Set(); + if (batch.remove !== undefined) { + if (!Array.isArray(batch.remove)) { + throw new TypeError('[webgpu-search] batch.remove must be an array of document IDs.'); + } + for (let i = 0; i < batch.remove.length; i++) { + const id = batch.remove[i]; + validateDocumentId(id, `remove at index ${i}`); + removeIds.push(id); + if (s.idToDocIndex.has(id)) { + willRemoveSet.add(id); + } + } + } + + const preparedUpdates: PreparedDocMutation[] = []; + const updatedIdsInBatch = new Set(); + if (batch.update !== undefined) { + if (!Array.isArray(batch.update)) { + throw new TypeError('[webgpu-search] batch.update must be an array of documents.'); + } + for (let i = 0; i < batch.update.length; i++) { + const doc = batch.update[i]; + if (doc === null || typeof doc !== 'object') { + throw new TypeError(`[webgpu-search] batch.update document at index ${i} must be an object.`); + } + const id = s.getId(doc); + validateDocumentId(id, `update at index ${i}`); + if (!s.idToDocIndex.has(id) || willRemoveSet.has(id)) { + throw new DocumentNotFoundError(id); + } + if (updatedIdsInBatch.has(id)) { + throw new DuplicateIdError(id); + } + updatedIdsInBatch.add(id); + // Filter getters run here (before any mutation) so a throwing + // getter rejects the whole batch with the index untouched. + const prepared = prepareDocFields(s, doc, id); + preparedUpdates.push({ ...prepared, columnar: s.columnarStore.prepare(doc) }); + } + } + + const preparedAdds: PreparedDocMutation[] = []; + const addedIdsInBatch = new Set(); + + if (batch.add !== undefined) { + if (!Array.isArray(batch.add)) { + throw new TypeError('[webgpu-search] batch.add must be an array of documents.'); + } + for (let i = 0; i < batch.add.length; i++) { + const doc = batch.add[i]; + if (doc === null || typeof doc !== 'object') { + throw new TypeError(`[webgpu-search] batch.add document at index ${i} must be an object.`); + } + const id = s.getId(doc); + validateDocumentId(id, `add at index ${i}`); + + const existsInIndex = s.idToDocIndex.has(id) && !willRemoveSet.has(id); + const existsInBatch = addedIdsInBatch.has(id); + + if (existsInIndex || existsInBatch) { + if (!upsert) { + throw new DuplicateIdError(id); + } + } + addedIdsInBatch.add(id); + + const prepared = prepareDocFields(s, doc, id); + preparedAdds.push({ ...prepared, columnar: s.columnarStore.prepare(doc) }); + } + } + + // PHASE 2: DETERMINISTIC EXECUTION (remove -> update -> add) + s.generation++; + let removedCount = 0; + let updatedCount = 0; + let addedCount = 0; + + // A. Remove + for (let i = 0; i < removeIds.length; i++) { + const id = removeIds[i]; + const d = s.idToDocIndex.get(id); + if (d !== undefined) { + const rows = s.docToRowIndices[d]; + if (rows) { + for (let f = 0; f < rows.length; f++) { + s.tombstones.add(rows[f]); + } + } + s.idToDocIndex.delete(id); + s.records[d] = null as any; + s.docIds[d] = null as any; + s.rawFieldStrings[d] = null as any; + s.docToRowIndices[d] = null as any; + s.columnarStore.remove(d); + removedCount++; + } + } + + // B. Update + const newRowTokensToAppend: Uint32Array[] = []; + const newRowDocIndicesToAppend: number[] = []; + const newRowFieldIndicesToAppend: number[] = []; + let newTokensToAppendCount = 0; + + for (let i = 0; i < preparedUpdates.length; i++) { + const upd = preparedUpdates[i]; + const oldD = s.idToDocIndex.get(upd.id)!; + const oldRows = s.docToRowIndices[oldD]; + if (oldRows) { + for (let f = 0; f < oldRows.length; f++) { + s.tombstones.add(oldRows[f]); + } + } + s.records[oldD] = null as any; + s.docIds[oldD] = null as any; + s.rawFieldStrings[oldD] = null as any; + s.docToRowIndices[oldD] = null as any; + s.columnarStore.remove(oldD); + + const newD = s.records.length; + s.records.push(upd.doc); + s.docIds.push(upd.id); + s.idToDocIndex.set(upd.id, newD); + s.rawFieldStrings.push(upd.rawStrings); + s.columnarStore.addPrepared(newD, upd.columnar); + + const docRows: number[] = new Array(s.sortedFields.length); + for (let f = 0; f < s.sortedFields.length; f++) { + const rowIdx = s.rowTokens.length + newRowTokensToAppend.length; + newRowTokensToAppend.push(upd.tokens[f]); + newRowDocIndicesToAppend.push(newD); + newRowFieldIndicesToAppend.push(f); + docRows[f] = rowIdx; + newTokensToAppendCount += upd.tokens[f].length; + } + s.docToRowIndices.push(docRows); + updatedCount++; + } + + // C. Add + for (let i = 0; i < preparedAdds.length; i++) { + const ad = preparedAdds[i]; + if (s.idToDocIndex.has(ad.id)) { + const oldD = s.idToDocIndex.get(ad.id)!; + const oldRows = s.docToRowIndices[oldD]; + if (oldRows) { + for (let f = 0; f < oldRows.length; f++) { + s.tombstones.add(oldRows[f]); + } + } + s.records[oldD] = null as any; + s.docIds[oldD] = null as any; + s.rawFieldStrings[oldD] = null as any; + s.docToRowIndices[oldD] = null as any; + s.columnarStore.remove(oldD); + updatedCount++; + } else { + addedCount++; + } + + const newD = s.records.length; + s.records.push(ad.doc); + s.docIds.push(ad.id); + s.idToDocIndex.set(ad.id, newD); + s.rawFieldStrings.push(ad.rawStrings); + s.columnarStore.addPrepared(newD, ad.columnar); + + const docRows: number[] = new Array(s.sortedFields.length); + for (let f = 0; f < s.sortedFields.length; f++) { + const rowIdx = s.rowTokens.length + newRowTokensToAppend.length; + newRowTokensToAppend.push(ad.tokens[f]); + newRowDocIndicesToAppend.push(newD); + newRowFieldIndicesToAppend.push(f); + docRows[f] = rowIdx; + newTokensToAppendCount += ad.tokens[f].length; + } + s.docToRowIndices.push(docRows); + } + + // Append newly prepared rows to rowTokens, rowToDocIndex, rowToFieldIndex + if (newRowTokensToAppend.length > 0) { + for (let r = 0; r < newRowTokensToAppend.length; r++) { + s.rowTokens.push(newRowTokensToAppend[r]); + s.rowToDocIndex.push(newRowDocIndicesToAppend[r]); + s.rowToFieldIndex.push(newRowFieldIndicesToAppend[r]); + } + s.totalTokens += newTokensToAppendCount; + } + + // D. Check Compaction / GPU Headroom + const totalRows = s.rowTokens.length; + const tombstoneCount = s.tombstones.size; + const tombstoneRatio = totalRows > 0 ? tombstoneCount / totalRows : 0; + + let compacted = false; + const isGpu = s.engineState.engine === 'webgpu' && s.gpuEngine !== null && s.gpuEngine.isReady; + const gpuHeadroomExhausted = isGpu && !s.gpuEngine!.canFitHeadroom(totalRows, s.totalTokens); + + const activeDocCount = s.idToDocIndex.size; + const newCandidateCap = Math.min(32768, Math.max(8192, Math.floor(activeDocCount * s.sortedFields.length * 0.1))); + if (newCandidateCap > s.candidateCapacity) { + s.candidateCapacity = newCandidateCap; + s.gpuEngine?.ensureCandidateCapacity(newCandidateCap); + } + + try { + if (tombstoneRatio >= 0.25 || gpuHeadroomExhausted) { + await compactIndex(s); + compacted = true; + } else if (isGpu && newRowTokensToAppend.length > 0) { + await syncAppendedRowsToGpu(s, newRowTokensToAppend, totalRows, newTokensToAppendCount); + } else if (!isGpu && (s.preferGpu || totalRows >= (s.options.threshold ?? 30_000)) && s.options.preferGpu !== false && activeDocCount > 0) { + await tryInitializeGpuEngine(s); + } + } catch (err) { + console.warn('[webgpu-search] GPU sync failed during mutation; falling back to CPU:', err); + if (s.gpuEngine) { + try { s.gpuEngine.destroy(); } catch {} + s.gpuEngine = null; + } + s.engineState = transitionEngineState(s.engineState, 'cpu', 'gpu-execution-error'); + s.vramAllocatedBytes = 0; + } + + if (s.engineState.engine === 'cpu' && !s.engineState.fallbackReason) { + s.engineState = transitionEngineState( + s.engineState, + 'cpu', + s.options.preferGpu === false ? 'prefer-cpu' : 'below-threshold', + ); + } + + s.mutationEpoch++; + const durationMs = nowMs() - tStart; + s.lastMutationTimeMs = durationMs; + + return { + added: addedCount, + updated: updatedCount, + removed: removedCount, + mutationEpoch: s.mutationEpoch, + compacted, + durationMs + }; +} + +/** + * Incremental GPU sync of appended mutation rows (offsets continue at the + * pre-append corpus token count). VRAM accounting is refreshed afterwards. + */ +export async function syncAppendedRowsToGpu( + s: DocumentIndexState, + newRows: Uint32Array[], + totalRows: number, + newTokensCount: number +): Promise { + if (!s.gpuEngine || !s.gpuEngine.isReady) return; + const numNewRows = newRows.length; + const oldTotalTokens = s.totalTokens - newTokensCount; + + const newOffsets = new Uint32Array(numNewRows); + let runningOffset = oldTotalTokens; + for (let i = 0; i < numNewRows; i++) { + runningOffset += newRows[i].length; + newOffsets[i] = runningOffset; + } + + const newTokens = new Uint32Array(newTokensCount); + let pos = 0; + for (let i = 0; i < numNewRows; i++) { + newTokens.set(newRows[i], pos); + pos += newRows[i].length; + } + + await s.gpuEngine.appendRows(newTokens, newOffsets, totalRows); + s.vramAllocatedBytes = s.gpuEngine.vramAllocatedBytes; +} + +/** + * CPU-side compaction: drop tombstoned slots, re-pack documents/ids/rows + * into the field-stratified layout, and remap columnar store indices. + */ +export function compactCpu(s: DocumentIndexState): void { + const activeDocCount = s.idToDocIndex.size; + const fieldCount = s.sortedFields.length; + const totalRows = activeDocCount * fieldCount; + + const activeOldDocIndices: number[] = []; + for (let d = 0; d < s.records.length; d++) { + const id = s.docIds[d]; + if (id !== null && id !== undefined && s.idToDocIndex.get(id) === d) { + activeOldDocIndices.push(d); + } + } + + const newRecords: TDoc[] = new Array(activeDocCount); + const newDocIds: DocumentId[] = new Array(activeDocCount); + const newIdToDocIndex = new Map(); + const newRawFieldStrings: string[][] = Array.from({ length: activeDocCount }, () => new Array(fieldCount)); + + for (let newD = 0; newD < activeDocCount; newD++) { + const oldD = activeOldDocIndices[newD]; + const doc = s.records[oldD]; + const id = s.docIds[oldD]; + newRecords[newD] = doc as TDoc; + newDocIds[newD] = id as DocumentId; + newIdToDocIndex.set(id as DocumentId, newD); + for (let f = 0; f < fieldCount; f++) { + newRawFieldStrings[newD][f] = s.rawFieldStrings[oldD][f]; + } + } + + const newRowTokens: Uint32Array[] = new Array(totalRows); + const newRowToDocIndex: number[] = new Array(totalRows); + const newRowToFieldIndex: number[] = new Array(totalRows); + const newDocToRowIndices: number[][] = Array.from({ length: activeDocCount }, () => new Array(fieldCount)); + + let corpusTokens = 0; + let rowIdx = 0; + + for (let f = 0; f < fieldCount; f++) { + for (let newD = 0; newD < activeDocCount; newD++) { + const oldD = activeOldDocIndices[newD]; + const oldRowIdx = s.docToRowIndices[oldD][f]; + const tokens = s.rowTokens[oldRowIdx]; + newRowTokens[rowIdx] = tokens; + newRowToDocIndex[rowIdx] = newD; + newRowToFieldIndex[rowIdx] = f; + newDocToRowIndices[newD][f] = rowIdx; + corpusTokens += tokens.length; + rowIdx++; + } + } + + s.records = newRecords; + s.docIds = newDocIds; + s.idToDocIndex = newIdToDocIndex; + s.rawFieldStrings = newRawFieldStrings; + s.rowTokens = newRowTokens; + s.rowToDocIndex = newRowToDocIndex; + s.rowToFieldIndex = newRowToFieldIndex; + s.docToRowIndices = newDocToRowIndices; + s.totalTokens = corpusTokens; + const oldToNewDocIndexMap = new Map(); + for (let newD = 0; newD < activeDocCount; newD++) { + oldToNewDocIndexMap.set(activeOldDocIndices[newD], newD); + } + s.columnarStore.compact(oldToNewDocIndexMap, activeDocCount); + s.tombstones.clear(); +} + +/** True when the row layout is exactly field-stratified (compaction no-op). */ +function isStratified(s: DocumentIndexState): boolean { + const docCount = s.records.length; + if (docCount === 0) return true; + for (let r = 0; r < s.rowTokens.length; r++) { + if ( + s.rowToFieldIndex[r] !== Math.floor(r / docCount) || + s.rowToDocIndex[r] !== (r % docCount) + ) { + return false; + } + } + return true; +} + +/** + * Public synchronous compaction: compacts CPU state when tombstones exist or + * the row layout is no longer stratified, discarding the GPU engine (which + * cannot represent the compacted layout without a full re-upload). + */ +export function compactSync(s: DocumentIndexState): void { + if (s.tombstones.size > 0 || !isStratified(s)) { + compactCpu(s); + if (s.gpuEngine) { + s.gpuEngine.destroy(); + s.gpuEngine = null; + s.engineState = transitionEngineState( + s.engineState, + 'cpu', + s.options.preferGpu === false ? 'prefer-cpu' : 'below-threshold', + ); + } + } +} + +/** + * Full compaction with GPU reload: compact CPU state first, then re-upload + * the compacted corpus into the GPU engine (or reset it for an empty index). + */ +export async function compactIndex(s: DocumentIndexState): Promise { + compactCpu(s); + + if (s.engineState.engine === 'webgpu' && s.gpuEngine && s.gpuEngine.isReady) { + const activeDocCount = s.idToDocIndex.size; + const totalRows = s.rowTokens.length; + if (activeDocCount === 0) { + const packed = packDataset([], { normalized: s.normalized }); + await s.gpuEngine.loadDataset(packed); + s.vramAllocatedBytes = s.gpuEngine.vramAllocatedBytes; + } else { + const packed = packDataset(s.rowTokens, { + normalized: s.normalized, + totalTokens: s.totalTokens + }); + const rowCap = Math.max(totalRows, 16); + const tokenCap = Math.max(s.totalTokens, 64); + await s.gpuEngine.loadDataset(packed, { + rowCapacity: rowCap, + tokenCapacity: tokenCap, + growthFactor: s.growthFactor + }); + s.vramAllocatedBytes = s.gpuEngine.vramAllocatedBytes; + } + } +} + +/** + * Single device-loss subscription (re-subscribes replace the old one so + * rebuild cycles never leak listeners). On loss the GPU engine is torn + * down, VRAM accounting is zeroed, and the state moves to CPU with + * `fallbackReason: 'device-lost'` — CPU corpora (`rowTokens`, + * `records`) are retained so `rebuildGpu()` can re-upload identically. + */ +export function subscribeDeviceLost(s: DocumentIndexState): void { + if (s.unsubscribeDeviceLost) { + s.unsubscribeDeviceLost(); + s.unsubscribeDeviceLost = undefined; + } + s.unsubscribeDeviceLost = GpuDevicePool.onDeviceLost(() => { + console.warn('[webgpu-search] GPU device lost, falling back to CPU.'); + if (s.gpuEngine) { + s.gpuEngine.destroy(); + s.gpuEngine = null; + } + s.vramAllocatedBytes = 0; + s.engineState = transitionEngineState(s.engineState, 'cpu', 'device-lost'); + }); +} + +/** + * Explicit device-loss rebuild path (Phase 2 reliability proof). + * + * After a `fallbackReason: 'device-lost'` (or `gpu-execution-error`) + * transition, attempts to re-acquire a GPU device and re-upload the + * retained in-memory corpus. On success the engine returns to `'webgpu'` + * with identical match semantics; on failure the index stays on CPU with + * an explicit `fallbackReason` and returns `false` (never throws for + * GPU-unavailable). + * + * - `preferGpu: false` indexes are CPU-by-design: rebuild is a no-op + * returning `false`. + * - Pass a fresh injected `device` after loss (e.g. a new mock device + * in headless tests); otherwise the creation-time `options.device` is + * reused, else the shared pool is tried. A supplied device is persisted + * into `options.device` for subsequent rebuilds. + * - Throws only when the index itself is destroyed. + */ +export async function rebuildGpuEngine( + s: DocumentIndexState, + options?: { device?: GPUDevice; powerPreference?: GPUPowerPreference } +): Promise { + if (s.isDestroyed) { + throw new Error('[webgpu-search] DocumentIndex has been destroyed.'); + } + if (s.gpuEngine !== null && s.gpuEngine.isReady && s.engineState.engine === 'webgpu') { + return true; + } + if (s.options.preferGpu === false) { + return false; + } + if (s.idToDocIndex.size === 0) { + return false; + } + if (options?.device !== undefined) { + (s.options as { device?: GPUDevice }).device = options.device; + } + if (options?.powerPreference !== undefined) { + (s.options as { powerPreference?: GPUPowerPreference }).powerPreference = options.powerPreference; + } + if (s.gpuEngine) { + try { s.gpuEngine.destroy(); } catch {} + s.gpuEngine = null; + } + s.vramAllocatedBytes = 0; + await tryInitializeGpuEngine(s); + return s.engineState.engine === 'webgpu'; +} + +/** + * Attempt (re)initialization of the GPU engine from retained CPU state, + * honoring preferGpu/threshold/memory-budget routing. Never throws for + * GPU-unavailable environments; engine failures degrade to CPU fallback. + */ +export async function tryInitializeGpuEngine(s: DocumentIndexState): Promise { + const totalRows = s.rowTokens.length; + const threshold = s.options.threshold ?? 30_000; + const shouldAttempt = s.options.preferGpu === false + ? false + : s.preferGpu || totalRows >= threshold; + + if (!shouldAttempt) { + s.engineState = transitionEngineState( + s.engineState, + 'cpu', + s.options.preferGpu === false ? 'prefer-cpu' : 'below-threshold', + ); + return; + } + + const measuredAvgBytes = totalRows > 0 ? (s.totalTokens * 4) / totalRows : 0; + const budget = checkMemoryBudget(totalRows, measuredAvgBytes, s.options.device, s.candidateCapacity); + if (!budget.allowed) { + s.engineState = transitionEngineState(s.engineState, 'cpu', 'memory-budget-exceeded'); + return; + } + + try { + const gpu = new WebGPUEngine(); + const initialized = await gpu.init( + s.options.device !== undefined || s.options.powerPreference !== undefined + ? { device: s.options.device, powerPreference: s.options.powerPreference } + : undefined + ); + if (initialized && gpu.isReady) { + gpu.ensureCandidateCapacity(s.candidateCapacity); + const packed = packDataset(s.rowTokens, { + normalized: s.normalized, + totalTokens: s.totalTokens + }); + const effectiveCap = s.initialCapacity > 0 ? s.initialCapacity : Math.max(16, s.records.length); + const rowCap = effectiveCap * s.sortedFields.length; + const avgTokens = totalRows > 0 ? s.totalTokens / totalRows : 16; + const tokenCap = Math.max(Math.floor(rowCap * avgTokens), 64); + await gpu.loadDataset(packed, { + rowCapacity: rowCap, + tokenCapacity: tokenCap, + growthFactor: s.growthFactor + }); + s.gpuEngine = gpu; + s.engineState = transitionEngineState(s.engineState, 'webgpu'); + s.vramAllocatedBytes = gpu.vramAllocatedBytes; + subscribeDeviceLost(s); + } else { + s.engineState = transitionEngineState( + s.engineState, + 'cpu', + typeof navigator === 'undefined' || !('gpu' in navigator) || !navigator.gpu + ? 'webgpu-unsupported' + : 'device-request-failed', + ); + } + } catch { + s.engineState = transitionEngineState( + s.engineState, + 'cpu', + typeof navigator === 'undefined' || !('gpu' in navigator) || !navigator.gpu + ? 'webgpu-unsupported' + : 'device-request-failed', + ); + } +} + +/** + * Destroy the index: release the GPU engine + device-loss listener, drop the + * corpus, and reset the engine state so post-destroy calls fail closed. + */ +export function destroyIndex(s: DocumentIndexState): void { + s.isDestroyed = true; + s.generation++; + if (s.unsubscribeDeviceLost) { + s.unsubscribeDeviceLost(); + s.unsubscribeDeviceLost = undefined; + } + if (s.gpuEngine) { + s.gpuEngine.destroy(); + s.gpuEngine = null; + } + s.records = []; + s.docIds = []; + s.idToDocIndex.clear(); + s.rowTokens = []; + s.rowToDocIndex = []; + s.rowToFieldIndex = []; + s.docToRowIndices = []; + s.rawFieldStrings = []; + s.totalTokens = 0; + s.tombstones.clear(); + s.columnarStore = new ColumnarStore(s.filterFieldDefinitions, { + initialCapacity: s.initialCapacity, + growthFactor: s.growthFactor + }); + s.engineState = transitionEngineState( + s.engineState, + 'cpu', + s.engineState.fallbackReason, + ); + s.vramAllocatedBytes = 0; + const p = s.searchMutex.then(() => {}, () => {}); + s.searchMutex = p.catch(() => {}); +} diff --git a/packages/webgpu-search/src/document/filter-facet.ts b/packages/webgpu-search/src/document/filter-facet.ts new file mode 100644 index 0000000..5ad4f39 --- /dev/null +++ b/packages/webgpu-search/src/document/filter-facet.ts @@ -0,0 +1,140 @@ +/** + * Search filter compilation + facet aggregation (extracted verbatim from + * the pre-refactor monolithic `document-index.ts`). The public + * `DocumentIndex` facade delegates here via `document/query.ts`. + * + * Scope of this module: + * - `compileFilterPhase`: compiles `FilterExpression` to a bitset, measures + * `filteringMs`, enforces abort/budget gates, and composes extension + * predicates conjunctively. + * - `buildFacetResults`: aggregates facet buckets over an explicit + * query-matched doc set with disjunctive filter exclusion. + * + * All functions operate on the shared `DocumentIndexState` bundle and are + * read-only with respect to index internals. + */ + +import type { DocumentBitset } from '../filtering/doc-bitset'; +import { compileFilter } from '../filtering/compile-filter'; +import { excludeFieldFromFilter, type FacetEngine, type NormalizedFacet } from '../faceting/facet-engine'; +import { nowMs, throwIfAborted } from '../guard'; +import { throwIfBudgetAborted, assertTimeBudget, type NormalizedCostBudget } from '../diagnostics'; +import type { + DocumentSearchOptions, + FacetResult, + FilterExpression, + SearchHooks, +} from '../types'; +import type { DocumentIndexState } from './state'; + +/** + * Normalized cost budget shape (validated via `normalizeCostBudgetOptions`, + * which may return `undefined` when no budget is configured). Shared by the + * search validation, filter compilation, and engine routing phases. + */ +export type CostBudget = NormalizedCostBudget; +/** + * Compile structured/predicate filters (extracted from `search`). + * Compiles `FilterExpression` to a bitset, measures `filteringMs`, + * enforces abort/budget gates, and composes extension predicates. + * Empty bitsets return early with `filteringMs: 0` so callers can take + * the `noHits` path with identical telemetry. No behavior change. + */ +export function compileFilterPhase( + s: DocumentIndexState, + filter: DocumentSearchOptions['filter'], + effectiveHooks: SearchHooks | undefined, + wantsDiagnostics: boolean, + signal: AbortSignal | undefined, + budget: CostBudget | undefined, + queryStartMs: number +) { + let filterBitset: DocumentBitset | undefined = undefined; + let filterPredicate: ((doc: TDoc) => boolean) | undefined = undefined; + let structuredFilter: FilterExpression | undefined = undefined; + let filteringMs = 0; + + const tFilter0 = wantsDiagnostics ? nowMs() : 0; + if (filter !== undefined) { + if (typeof filter === 'function') { + filterPredicate = filter; + } else if (typeof filter === 'object' && filter !== null) { + structuredFilter = filter as FilterExpression; + filterBitset = compileFilter(structuredFilter, s.columnarStore); + if (filterBitset.isEmpty()) { + return { filterBitset, filterPredicate, structuredFilter, filteringMs }; + } + } else { + throw new TypeError('[webgpu-search] options.filter must be a function or FilterExpression.'); + } + } + if (wantsDiagnostics) filteringMs = nowMs() - tFilter0; + throwIfAborted(signal); + throwIfBudgetAborted(budget); + assertTimeBudget(queryStartMs, budget); + + // extension filter predicates compose conjunctively (AND) with + // `options.filter` functions. Index-level vs per-query hooks already + // merged via `resolveEffectiveHooks` (per-query wins per key). + const extensionPredicate = effectiveHooks?.filterPredicate; + if (extensionPredicate !== undefined) { + const basePredicate = filterPredicate; + if (basePredicate === undefined) { + filterPredicate = extensionPredicate; + } else { + const optFn = basePredicate; + const extFn = extensionPredicate; + filterPredicate = (doc: TDoc) => { + const d = doc as TDoc; + return optFn(d) && extFn(d); + }; + } + } + + return { filterBitset, filterPredicate, structuredFilter, filteringMs }; +} + +/** + * Aggregates facet buckets over an explicit query-matched doc set with + * disjunctive filter exclusion. Each facet ignores structured filter + * clauses on its own field (`excludeFieldFromFilter`) while keeping all + * other clauses; function predicates stay conjunctive. `or`/`not` branches + * referencing the facet field are kept verbatim (conservative, aliased). + * Cost is O(F*n) with F bitset allocs plus an optional extra unfiltered + * parity scan when a structured filter is present. + */ +export function buildFacetResults( + s: DocumentIndexState, + facetEngine: FacetEngine, + facetSpec: NormalizedFacet[], + queryMatched: Iterable, + structuredFilter: FilterExpression | undefined, + filterPredicate: ((doc: TDoc) => boolean) | undefined, + isApproximate: boolean +): Record { + const base: number[] = []; + for (const dIdx of queryMatched) { + const doc = s.records[dIdx]; + if (doc === null || doc === undefined) continue; + if (filterPredicate && !filterPredicate(doc)) continue; + base.push(dIdx); + } + const out: Record = {}; + for (const facet of facetSpec) { + let candidates: number[] = base; + if (structuredFilter !== undefined) { + const excluded = excludeFieldFromFilter(structuredFilter, facet.request.field); + if (excluded !== undefined) { + const mask = compileFilter(excluded, s.columnarStore); + candidates = base.filter((d) => mask.has(d)); + } + } + const result = facetEngine.aggregateOne(facet, candidates, isApproximate); + if (facet.name === '__proto__') { + Object.defineProperty(out, facet.name, { value: result, enumerable: true, configurable: true, writable: true }); + } else { + out[facet.name] = result; + } + } + return out; +} diff --git a/packages/webgpu-search/src/document/query.ts b/packages/webgpu-search/src/document/query.ts new file mode 100644 index 0000000..5a6f70b --- /dev/null +++ b/packages/webgpu-search/src/document/query.ts @@ -0,0 +1,1513 @@ +/** + * Document-index query implementation (extracted verbatim from the + * pre-refactor monolithic `document-index.ts). The public `DocumentIndex` + * facade delegates `search` / `autocomplete` here. + * + * Scope of this module: + * - `validateSearchOptions`: fail-closed option validation (query shape, + * mode/scorer/profile checks, token/prefix/typo normalization, hooks, + * budget/diagnostics, query-length gates, facet, ranking/autocomplete, + * field scoping, tokenizer expansion). + * - `runSearch`: GPU/CPU routed execution with facets, inline + * autocomplete, diagnostics, and abort/budget gates. + * - `runAutocomplete` + `computeSuggestions`: standalone and inline + * suggestion enumeration. + * - `enrichResults` / `enrichHighlights` / `resortResultsAfterScoring`: + * Top-K highlight + post-match hook pipeline. + * - `routeEngine`: candidate-budget + broad-query + field-restriction + * routing decisions. + * - `aggregateDocMatches`: per-document best-field aggregation delegator. + * + * All functions operate on the shared `DocumentIndexState` bundle; + * mutations are serialized through the `searchMutex` by the facade. + */ + +import { CPUEngine } from '../cpu-engine'; +import { + scoreExactMatchesMultiField, + scoreFuzzyTokens, + scoreSubstringTypoTokens, + type CpuModeOptions +} from '../exact-scorer'; +import { + normalizeTypoTolerance, + type NormalizedTypoOptions +} from '../search/typo-tolerance'; +import { + normalizeTokenMatchOptions, + type NormalizedTokenMatchOptions +} from '../search/token-search'; +import { + normalizePrefixOptions, + assertPrefixLengthForQuery, + scorePrefixTokens, + type NormalizedPrefixOptions +} from '../search/prefix-search'; +import { + compareRanked, + isExactTokenMatch, + normalizeTieBreakers, + type RankableCandidate +} from '../ranking'; +import { + aggregateDocMatches as aggregateDocMatchesPure, + accumulateDocMatch, +} from '../document-aggregate'; +import { + normalizeAutocompleteOptions, + type NormalizedAutocompleteOptions +} from '../autocomplete'; +import { + resolveEffectiveHooks, + getTokenTermsForQuery, + applyScoringHook, + applyPostProcess, +} from '../hooks'; +import { + clampLimit, + throwIfAborted, + abortError, + nowMs, + assertValidMode +} from '../guard'; +import { + throwIfBudgetAborted, + assertTimeBudget, + assertCandidateBudget, + computeFilterSelectivity, + isBroadQueryHeuristic, + isBroadSelectivity, + broadQueryRouteWarning, + broadSelectivityWarning, + candidateOverflowWarning, + normalizeCostBudgetOptions +} from '../diagnostics'; +import { + QUERY_TOKENS_MAX, + SCORING_VERSION, + normalizeCpuScorer, + IncompatibleOptionError, + ProfileMismatchError, + QueryTooLongError, +} from '../text-profile'; +import { transitionEngineState } from '../engine-state'; +import type { WebGPUEngine } from '../webgpu-engine'; +import { normalizeText } from '../text-normalization'; +import type { DocumentBitset } from '../filtering/doc-bitset'; +import { + FacetEngine, + normalizeFacetRequests +} from '../faceting/facet-engine'; +import { alignHighlights, renderHighlightedText } from '../highlight'; +import { compileFilterPhase, buildFacetResults, type CostBudget } from './filter-facet'; +import { setOwnKey } from './build'; +import type { DocumentIndexState } from './state'; +import type { + AutocompleteOptions, + DocumentSearchOptions, + DocumentSearchResponse, + DocumentSearchResultItem, + EngineType, + FacetResult, + HighlightRange, + QueryDiagnostics, + QueryDiagnosticsTimings, + SearchHooks, + SearchMode, + SearchTimings, + SuggestionItem, + SuggestResponse, + TieBreakerCriterion, +} from '../types'; + +/** + * Validate search options fail-closed (extracted from `search`). + * Covers query shape, mode/scorer/profile checks, token/prefix/typo + * normalization, hooks, budget/diagnostics, query-length gates, facet, + * ranking/autocomplete, field scoping, and tokenizer expansion. + * No behavior change: code moved verbatim from `search`. + */ +export function validateSearchOptions( + s: DocumentIndexState, + query: string, + options: DocumentSearchOptions +) { + if (typeof query !== 'string') { + throw new TypeError(`[webgpu-search] search expects query: string, got ${typeof query}.`); + } + if ((options as unknown as Record).cpuAlgorithm !== undefined) { + throw new IncompatibleOptionError( + 'cpuScorer', + "[webgpu-search] cpuAlgorithm was removed; use cpuScorer: 'exact' | 'ufuzzy'." + ); + } + if ((options as unknown as Record).extensions !== undefined) { + throw new IncompatibleOptionError( + 'hooks', + "[webgpu-search] extensions was removed; use hooks." + ); + } + if ((options as unknown as Record).suggest !== undefined) { + throw new IncompatibleOptionError( + 'autocomplete', + "[webgpu-search] suggest was removed; use autocomplete." + ); + } + + const { + mode = 'fuzzy', + caseSensitive = false, + signal, + cpuScorer: cpuScorerOpt, + onQueryTooLong = 'throw', + fields: searchFields, + facets, + faceting + } = options; + + const cpuScorer = normalizeCpuScorer(cpuScorerOpt) ?? 'exact'; + if (caseSensitive === s.normalized) { + throw new ProfileMismatchError(!s.normalized, caseSensitive); + } + if (s.preferGpu && cpuScorer === 'ufuzzy') { + throw new IncompatibleOptionError( + 'cpuScorer', + "cpuScorer:'ufuzzy' is CPU-only; use preferGpu:false or cpuScorer:'exact'." + ); + } + assertValidMode(mode); + // fail-closed option validation up front so malformed + // token/prefix/typo shapes throw identically on GPU and CPU paths. + // 'fuzzy' validates typo shape but ignores it (subsequence matching is + // inherently typo-tolerant); 'substring'/'token'/'prefix' honor it. + const tokenOpts: NormalizedTokenMatchOptions = normalizeTokenMatchOptions(options.tokenMatch); + const prefixOpts: NormalizedPrefixOptions = normalizePrefixOptions(options.prefixMatch); + const typo: NormalizedTypoOptions = normalizeTypoTolerance(options.typoTolerance); + if (mode === 'prefix' && options.prefixMatch?.exactCase !== undefined && prefixOpts.exactCase !== caseSensitive) { + throw new ProfileMismatchError(caseSensitive, prefixOpts.exactCase, 'prefixMatch.exactCase'); + } + // Legacy ufuzzy/native scorers only implement fuzzy/substring-exact. + if (cpuScorer === 'ufuzzy' && (mode === 'token' || mode === 'prefix' || typo.enabled)) { + throw new IncompatibleOptionError( + 'cpuScorer', + `cpuScorer:'ufuzzy' supports only exact 'fuzzy'/'substring' modes without typo tolerance (got mode '${mode}'${typo.enabled ? ' with typoTolerance' : ''}). Use cpuScorer:'exact'.` + ); + } + // fail-closed extension hook validation up front so malformed + // hook shapes throw identically on GPU and CPU paths (before early exits). + const effectiveHooks = resolveEffectiveHooks(s.indexHooks, options.hooks); + // Threaded into every exact CPU call below (single normalization). + const cpuModeOptions: CpuModeOptions = { + tokenMatch: { operator: tokenOpts.operator, minMatchCount: tokenOpts.minMatchCount }, + prefixMatch: prefixOpts.prefixLength !== undefined + ? { prefixLength: prefixOpts.prefixLength, exactCase: prefixOpts.exactCase } + : { exactCase: prefixOpts.exactCase }, + typoTolerance: { + enabled: typo.enabled, + maxDistance: typo.maxDistance, + minWordLengthForOneTypo: typo.minWordLengthForOneTypo, + minWordLengthForTwoTypos: typo.minWordLengthForTwoTypos, + prefixExactLength: typo.prefixExactLength + } + }; + // fail-closed cost-budget + diagnostics validation up front so + // malformed budgets throw identically on GPU and CPU paths and on empty + // corpora/queries (before early exits). Enforcement is independent of the + // diagnostics flag: budgets constrain even when telemetry is off. + const budget = normalizeCostBudgetOptions(options.budget); + if (options.diagnostics !== undefined && typeof options.diagnostics !== 'boolean') { + throw new TypeError('[webgpu-search] options.diagnostics must be a boolean.'); + } + const wantsDiagnostics = options.diagnostics === true; + const needsClock = wantsDiagnostics || budget?.maxExecutionTimeMs !== undefined; + const queryStartMs = needsClock ? nowMs() : 0; + const diagWarnings: string[] = []; + + const rawTrimmed = query.trim(); + let forceCpu = false; + if (rawTrimmed.length > QUERY_TOKENS_MAX * 4) { + if (rawTrimmed.length > 1_000_000) { + let cpCount = 0; + for (const _ch of rawTrimmed) cpCount++; + const est = cpCount * 3; + if (est > QUERY_TOKENS_MAX) { + if (onQueryTooLong !== 'cpu-fallback') { + throw new QueryTooLongError(QUERY_TOKENS_MAX, est, s.profileId); + } + forceCpu = true; + } + } + } + + const normalizedQuery = normalizeText(query, s.normalized); + const queryTokenCount = normalizedQuery.tokenCount; + if (queryTokenCount > QUERY_TOKENS_MAX) { + if (onQueryTooLong === 'cpu-fallback') { + forceCpu = true; + } else { + throw new QueryTooLongError(QUERY_TOKENS_MAX, queryTokenCount, s.profileId); + } + } + + const clampedLimit = clampLimit(options.limit ?? options.maxResults ?? 50); + throwIfAborted(signal); + throwIfBudgetAborted(budget); + + // Hoisted prefixLength check: fail-closed even on empty queries/corpora. + // Skipped for empty queries to match the scorer's early noMatch. + if (mode === 'prefix' && normalizedQuery.tokens.length > 0) { + assertPrefixLengthForQuery(prefixOpts, normalizedQuery.tokens.length); + } + + // custom tokenizer terms for 'token' mode (CPU-only). + // Other modes ignore the tokenizer. Empty queries skip hook invocation + // (no-match by design); non-empty token queries expand via the hook so + // scoring and highlights share one term set (score-highlight symmetry). + let customTokenTerms: Uint32Array[] | undefined = undefined; + if (effectiveHooks?.tokenizer !== undefined && mode === 'token' && !normalizedQuery.isEmpty) { + customTokenTerms = getTokenTermsForQuery( + query, + normalizedQuery.tokens, + s.normalized, + effectiveHooks.tokenizer + ); + cpuModeOptions.tokenTermsOverride = customTokenTerms; + } + + // (): validate facet contracts fail-fast, before early exits. + // An empty facet list is equivalent to not requesting facets (absent key). + // `faceting` is ignored unless facets are requested (no throw on absent). + const facetSpec = normalizeFacetRequests(facets); + const wantsFacets = facetSpec !== undefined && facetSpec.length > 0; + let facetingMode: 'auto' | 'force-exact' = 'auto'; + if (wantsFacets && faceting !== undefined) { + if (faceting !== 'auto' && faceting !== 'force-exact') { + throw new TypeError(`[webgpu-search] Invalid faceting option: "${String(faceting)}". Must be 'auto' or 'force-exact'.`); + } + facetingMode = faceting; + } + const facetEngine = wantsFacets ? new FacetEngine(s.columnarStore) : null; + if (wantsFacets && facetEngine && facetSpec) { + // Fail-closed field validation up front: unknown facet fields and + // range-on-non-number must throw even when an empty match set would + // otherwise take an early noHits exit below. + facetEngine.validateRequests(facetSpec); + } + + // Inline autocomplete config: deterministic ranking hierarchy. + // All scope validation (ranking, autocomplete, fields) is fail-closed before + // early exits so malformed shapes throw identically on empty and + // non-empty corpora/queries. + if (options.ranking !== undefined && (typeof options.ranking !== 'object' || options.ranking === null || Array.isArray(options.ranking))) { + throw new TypeError('[webgpu-search] options.ranking must be an object.'); + } + const tieBreakers: TieBreakerCriterion[] = normalizeTieBreakers( + options.ranking?.tieBreakers + ); + let suggestSpec: NormalizedAutocompleteOptions | undefined = undefined; + const suggestRaw = options.autocomplete; + if (suggestRaw !== undefined && suggestRaw !== false) { + suggestSpec = normalizeAutocompleteOptions(suggestRaw as AutocompleteOptions | boolean); + // Inline autocomplete inherits the search ranking hierarchy unless the + // autocomplete object carries its own explicit tieBreakers. + const rawSuggest = suggestRaw as AutocompleteOptions; + if ( + typeof rawSuggest === 'object' && + rawSuggest !== null && + rawSuggest.tieBreakers === undefined && + options.ranking?.tieBreakers !== undefined + ) { + suggestSpec.tieBreakers = [...tieBreakers]; + } + if (suggestSpec.field !== undefined && !s.fieldNameToIndex.has(suggestSpec.field)) { + throw new Error(`[webgpu-search] Unknown autocomplete field: "${suggestSpec.field}".`); + } + } + // Hoisted search-field validation (fail-closed before empty-query, + // empty-corpus, and empty-filter early exits below). + let allowedFieldIndices: Set | undefined = undefined; + if (searchFields !== undefined) { + if (!Array.isArray(searchFields)) { + throw new TypeError('[webgpu-search] options.fields must be an array of string field names.'); + } + if (searchFields.length !== 0) { + allowedFieldIndices = new Set(); + for (const fName of searchFields) { + if (typeof fName !== 'string') { + throw new TypeError('[webgpu-search] options.fields elements must be strings.'); + } + const idx = s.fieldNameToIndex.get(fName); + if (idx === undefined) { + throw new Error(`[webgpu-search] Unknown search field: "${fName}".`); + } + allowedFieldIndices.add(idx); + } + } + } + const hasEmptyFieldList = searchFields !== undefined && searchFields.length === 0; + + return { + mode, + caseSensitive, + signal, + cpuScorer, + tokenOpts, + prefixOpts, + typo, + effectiveHooks, + cpuModeOptions, + budget, + wantsDiagnostics, + queryStartMs, + diagWarnings, + normalizedQuery, + clampedLimit, + forceCpu, + facetSpec, + wantsFacets, + facetingMode, + facetEngine, + tieBreakers, + suggestSpec, + allowedFieldIndices, + hasEmptyFieldList, + customTokenTerms + }; +} + +/** + * Route to WebGPU vs CPU (extracted from `search`). + * Applies candidate-budget enforcement, broad-query CPU routing, and + * field-restriction routing. Mutates `diagWarnings` for the broad-query + * warning exactly as before. No behavior change. + */ +export function routeEngine( + s: DocumentIndexState, + params: { + mode: SearchMode; + typoEnabled: boolean; + cpuScorer: string; + forceCpu: boolean; + allowedFieldIndices: Set | undefined; + normalizedQueryTokenCount: number; + filterBitset: DocumentBitset | undefined; + budget: CostBudget | undefined; + queryStartMs: number; + wantsDiagnostics: boolean; + diagWarnings: string[]; + } +) { + const { + mode, + typoEnabled, + cpuScorer, + forceCpu, + allowedFieldIndices, + normalizedQueryTokenCount, + filterBitset, + budget, + queryStartMs, + wantsDiagnostics, + diagWarnings + } = params; + // token/prefix modes and typo-tolerant queries are CPU-only + // (exact-only WGSL kernels) and skip GPU dispatch with fallbackReason + // 'unsupported-mode' (parity boundary). + const isGpuSupportedMode: boolean = + (mode === 'fuzzy' || mode === 'substring') && !typoEnabled; + + // broad-query pre-dispatch guard + candidate ceiling. + // `filteredCandidateCount` is the exact post-filter population to score; + // exceeding `maxCandidates` throws before any scoring work. Short queries + // over massive corpora route to the CPU streaming scan to avoid GPU + // buffer saturation and driver timeouts (TDR). With function-predicate + // filters the pre-predicate population is checked (conservative — the + // true post-predicate count is unknowable pre-scan). + const activeDocCount = s.idToDocIndex.size; + const filteredCandidateCount = filterBitset !== undefined + ? filterBitset.popcount() + : activeDocCount; + assertCandidateBudget(filteredCandidateCount, budget); + let broadQueryCpuRoute = false; + if (isBroadQueryHeuristic(normalizedQueryTokenCount, activeDocCount)) { + broadQueryCpuRoute = true; + // Suppress the GPU-avoidance warning when already CPU-by-design + // (token/prefix/typo modes, explicit ufuzzy, query-too-long fallback): + // routing outcome is correct but the stated cause would misattribute. + const cpuByDesign = !isGpuSupportedMode || cpuScorer === 'ufuzzy' || forceCpu; + if (wantsDiagnostics && !cpuByDesign) { + diagWarnings.push(broadQueryRouteWarning(activeDocCount, normalizedQueryTokenCount)); + } + } + assertTimeBudget(queryStartMs, budget); + + // When field-restricted, route to CPU to prevent unselected high-priority + // fields from saturating GPU candidate buffer before low-priority allowed fields. + const isFieldRestricted = allowedFieldIndices !== undefined && allowedFieldIndices.size < s.sortedFields.length; + + // 1. WebGPU execution path (exact fuzzy/substring only; token/prefix + // and typo queries skip dispatch via isGpuSupportedMode above). + const gpuHandle = s.gpuEngine; + const useGpu = !forceCpu && + !isFieldRestricted && + !broadQueryCpuRoute && + isGpuSupportedMode && + cpuScorer !== 'ufuzzy' && + s.engineState.engine === 'webgpu' && + gpuHandle !== null && + gpuHandle.isReady; + + return { + isGpuSupportedMode, + broadQueryCpuRoute, + isFieldRestricted, + filteredCandidateCount, + useGpu, + gpuHandle + }; +} + +/** + * Aggregate per-row matches into per-document best-field hits (thin + * delegator over the pure `document-aggregate` helper so GPU/CPU paths + * share one ranking implementation; unit-tested without a device). + * No behavior change. + */ +export function aggregateDocMatches( + s: DocumentIndexState, + docMatches: Map }>, + normalizedQueryTokens: Uint32Array, + tieBreakers: TieBreakerCriterion[] +) { + return aggregateDocMatchesPure(docMatches, { + sortedFields: s.sortedFields, + docToRowIndices: s.docToRowIndices, + rowTokens: s.rowTokens, + docIds: s.docIds, + records: s.records, + normalizedQueryTokens, + tieBreakers + }); +} + +/** + * deterministic re-sort after `scoringHook` boosts. + * Reconstructs rank keys from index state (field weight, exactness via + * `isExactTokenMatch` on the winning row, matched length, id, docIndex) + * with updated scores, then sorts via `compareRanked` + `tieBreakers`. + * Preserves the total-order contract (deterministic across engines). + */ +export function resortResultsAfterScoring( + s: DocumentIndexState, + results: DocumentSearchResultItem[], + queryTokens: Uint32Array, + tieBreakers: TieBreakerCriterion[] +): void { + if (results.length <= 1) return; + const keys: RankableCandidate[] = new Array(results.length); + for (let i = 0; i < results.length; i++) { + const item = results[i] as DocumentSearchResultItem; + const dIdx = s.idToDocIndex.get(item.id); + const fIdx = s.fieldNameToIndex.get(item.matchedField); + if (dIdx !== undefined && fIdx !== undefined) { + const rowIdx = s.docToRowIndices[dIdx]?.[fIdx]; + const rowTok = rowIdx !== undefined ? s.rowTokens[rowIdx] : undefined; + const weight = s.sortedFields[fIdx]?.weight ?? 1.0; + keys[i] = { + score: item.score, + fieldWeight: weight, + isExactMatch: rowTok !== undefined ? isExactTokenMatch(rowTok, queryTokens) : false, + matchedLength: rowTok !== undefined ? rowTok.length : 0, + id: item.id, + docIndex: dIdx + }; + } else { + // Defensive fallback (unreachable in normal flow: resort runs before + // postProcess, so IDs/fields resolve). Positional docIndex preserves + // the total-order guarantee via unique positions. + keys[i] = { + score: item.score, + fieldWeight: 1.0, + isExactMatch: false, + matchedLength: 0, + id: item.id, + docIndex: i + }; + } + } + const order = results.map((_, i) => i); + order.sort((ia, ib) => compareRanked(keys[ia] as RankableCandidate, keys[ib] as RankableCandidate, tieBreakers)); + const sorted = order.map((i) => results[i] as DocumentSearchResultItem); + for (let i = 0; i < results.length; i++) { + results[i] = sorted[i] as DocumentSearchResultItem; + } +} + +export function enrichHighlights( + s: DocumentIndexState, + results: DocumentSearchResultItem[], + query: string, + mode: SearchMode, + options: DocumentSearchOptions, + customTokenTerms?: Uint32Array[], + precomputedTokens?: Uint32Array +): void { + const shouldHighlight = options.highlightOptions?.highlight ?? options.highlight ?? true; + if (!shouldHighlight) return; + + const tag = options.highlightOptions?.tag ?? options.tag; + const highlightFieldsRule = options.highlightOptions?.fields ?? options.highlightFields ?? 'all-matched'; + const escapeHtml = options.highlightOptions?.escapeHtml ?? options.escapeHtml ?? false; + + if ( + typeof highlightFieldsRule === 'string' && + highlightFieldsRule !== 'all-matched' && + highlightFieldsRule !== 'matched-field' && + highlightFieldsRule !== 'all-fields' + ) { + throw new TypeError( + `[webgpu-search] Invalid highlightFields option: "${highlightFieldsRule}". Must be 'all-matched', 'matched-field', 'all-fields', or an array of field names.` + ); + } + + const queryTokens = precomputedTokens ?? normalizeText(query, s.normalized).tokens; + // thread token/prefix/typo options so highlight ranges track + // the serving scorer (score-highlight symmetry across all modes). + // custom tokenizer terms replace the default split when supplied. + const alignOpts = { + mode, + normalized: s.normalized, + queryTokens, + tokenMatch: options.tokenMatch, + prefixMatch: options.prefixMatch, + typoTolerance: options.typoTolerance, + ...(customTokenTerms !== undefined ? { tokenTermsOverride: customTokenTerms } : {}) + }; + + for (let i = 0; i < results.length; i++) { + const item = results[i]; + const dIdx = s.idToDocIndex.get(item.id); + if (dIdx === undefined) continue; + + const highlights: Record = {}; + const highlightedText: Record = {}; + + if (highlightFieldsRule === 'all-fields') { + for (let f = 0; f < s.sortedFields.length; f++) { + const fieldDef = s.sortedFields[f]; + const rawStr = s.rawFieldStrings[dIdx][f] ?? ''; + const ranges = alignHighlights(rawStr, query, alignOpts); + if (ranges.length > 0) { + setOwnKey(highlights, fieldDef.name, ranges); + if (tag) { + setOwnKey(highlightedText, fieldDef.name, renderHighlightedText(rawStr, ranges, tag, escapeHtml)); + } + } + } + } else if (Array.isArray(highlightFieldsRule)) { + for (let k = 0; k < highlightFieldsRule.length; k++) { + const fName = highlightFieldsRule[k]; + const fIdx = s.fieldNameToIndex.get(fName); + if (fIdx !== undefined) { + const rawStr = s.rawFieldStrings[dIdx][fIdx] ?? ''; + const ranges = alignHighlights(rawStr, query, alignOpts); + if (ranges.length > 0) { + setOwnKey(highlights, fName, ranges); + if (tag) { + setOwnKey(highlightedText, fName, renderHighlightedText(rawStr, ranges, tag, escapeHtml)); + } + } + } + } + } else { + // 'matched-field' or 'all-matched' (default) + const primaryFIdx = s.fieldNameToIndex.get(item.matchedField); + if (primaryFIdx !== undefined) { + const primaryRaw = s.rawFieldStrings[dIdx][primaryFIdx] ?? ''; + const ranges = alignHighlights(primaryRaw, query, alignOpts); + setOwnKey(highlights, item.matchedField, ranges); + if (tag) { + setOwnKey(highlightedText, item.matchedField, renderHighlightedText(primaryRaw, ranges, tag, escapeHtml)); + } + } + + if (highlightFieldsRule === 'all-matched' && item.matches) { + for (let m = 0; m < item.matches.length; m++) { + const aux = item.matches[m]; + const auxFIdx = s.fieldNameToIndex.get(aux.field); + if (auxFIdx !== undefined) { + const auxRaw = s.rawFieldStrings[dIdx][auxFIdx] ?? ''; + const ranges = alignHighlights(auxRaw, query, alignOpts); + aux.highlights = ranges; + setOwnKey(highlights, aux.field, ranges); + if (tag) { + setOwnKey(highlightedText, aux.field, renderHighlightedText(auxRaw, ranges, tag, escapeHtml)); + } + } + } + } + } + + // Synchronize auxiliary match highlights across all modes if highlights exist + if (item.matches) { + for (let m = 0; m < item.matches.length; m++) { + const aux = item.matches[m]; + if (Object.prototype.hasOwnProperty.call(highlights, aux.field)) { + aux.highlights = highlights[aux.field]; + } + } + } + + item.highlights = highlights; + if (tag) { + item.highlightedText = highlightedText; + } + } +} + +/** + * Enrich Top-K results with highlights + post-match hooks (extracted). + * Runs highlight enrichment, scoring-hook boost + deterministic re-sort, + * and postProcess transform. Returns updated results with phase timings. + * No behavior change. + */ +export function enrichResults( + s: DocumentIndexState, + results: DocumentSearchResultItem[], + query: string, + mode: SearchMode, + options: DocumentSearchOptions, + customTokenTerms: Uint32Array[] | undefined, + normalizedQueryTokens: Uint32Array, + effectiveHooks: SearchHooks | undefined, + tieBreakers: TieBreakerCriterion[], + wantsDiagnostics: boolean +) { + const tHl0 = wantsDiagnostics ? nowMs() : 0; + enrichHighlights(s, results, query, mode, options, customTokenTerms, normalizedQueryTokens); + const highlightMs = wantsDiagnostics ? nowMs() - tHl0 : 0; + + // post-match extension pipeline (Top-K only). Scoring boosts + // apply identically on GPU and CPU paths (parity preserved), followed + // by a deterministic re-sort; postProcess is the final transform. + // hook time joins the scoring bucket on all paths for comparability. + const tHook0 = wantsDiagnostics ? nowMs() : 0; + if (effectiveHooks?.scoringHook !== undefined) { + applyScoringHook(results, effectiveHooks.scoringHook, query); + resortResultsAfterScoring(s, results, normalizedQueryTokens, tieBreakers); + } + let out = results; + if (effectiveHooks?.postProcess !== undefined) { + out = applyPostProcess(results, effectiveHooks.postProcess); + } + const hookMs = wantsDiagnostics ? nowMs() - tHook0 : 0; + return { results: out, highlightMs, hookMs }; +} + +/** + * Synchronous suggestion enumeration shared by `autocomplete()` and inline + * `search({ autocomplete })`. Callers pass pre-normalized query tokens and a + * validated spec; `searchFields` scopes the scan when the autocomplete spec + * carries no explicit field. + */ +export function computeSuggestions( + s: DocumentIndexState, + query: string, + queryTokens: Uint32Array, + spec: NormalizedAutocompleteOptions, + searchFields?: ReadonlySet +): SuggestionItem[] { + let fieldIndices: number[]; + if (spec.field !== undefined) { + fieldIndices = [s.fieldNameToIndex.get(spec.field) as number]; + } else if (searchFields !== undefined) { + fieldIndices = [...searchFields]; + } else { + fieldIndices = s.sortedFields.map((_, i) => i); + } + + const typoForSuggest = normalizeTypoTolerance( + spec.fuzzyDistance > 0 + ? { enabled: true, maxDistance: spec.fuzzyDistance as 1 | 2 } + : undefined + ); + const prefixOpts = normalizePrefixOptions(undefined); + + interface SuggestHit { + keys: RankableCandidate; + text: string; + dIdx: number; + fIdx: number; + rawStr: string; + } + const candidates: SuggestHit[] = []; + + for (let d = 0; d < s.records.length; d++) { + const doc = s.records[d]; + if (doc === null || doc === undefined) continue; + const id = s.docIds[d]; + if (id === null || id === undefined) continue; + const rows = s.docToRowIndices[d]; + if (!rows) continue; + for (let k = 0; k < fieldIndices.length; k++) { + const fIdx = fieldIndices[k] as number; + const rowIdx = rows[fIdx]; + if (rowIdx === undefined) continue; + if (s.tombstones.has(rowIdx)) continue; + const rowTok = s.rowTokens[rowIdx]; + if (!rowTok || rowTok.length === 0) continue; + const rawStr = s.rawFieldStrings[d]?.[fIdx] ?? ''; + if (rawStr.length === 0) continue; + + let matched = false; + let rawScore = 0; + if (spec.mode === 'prefix') { + const r = scorePrefixTokens(rowTok, queryTokens, prefixOpts, typoForSuggest); + matched = r.matched; + rawScore = r.score; + } else if (spec.fuzzyDistance > 0) { + const r = scoreSubstringTypoTokens(rowTok, queryTokens, typoForSuggest); + matched = r.matched; + rawScore = r.score; + } else { + const r = scoreFuzzyTokens(rowTok, queryTokens); + matched = r.matched; + rawScore = r.score; + } + if (!matched) continue; + const weight = s.sortedFields[fIdx].weight; + candidates.push({ + keys: { + score: Math.round(rawScore * weight), + fieldWeight: weight, + isExactMatch: isExactTokenMatch(rowTok, queryTokens), + matchedLength: rowTok.length, + id, + docIndex: d + }, + text: rawStr, + dIdx: d, + fIdx, + rawStr + }); + } + } + + candidates.sort((a, b) => compareRanked(a.keys, b.keys, spec.tieBreakers)); + const top = candidates.length > spec.limit ? candidates.slice(0, spec.limit) : candidates; + const out: SuggestionItem[] = []; + for (let i = 0; i < top.length; i++) { + const c = top[i] as SuggestHit; + const highlightMode = spec.mode === 'prefix' ? 'prefix' : (spec.fuzzyDistance > 0 ? 'substring' : 'fuzzy'); + let matchedRanges: HighlightRange[] = []; + try { + matchedRanges = alignHighlights(c.rawStr, query, { + mode: highlightMode, + normalized: s.normalized, + queryTokens, + typoTolerance: spec.fuzzyDistance > 0 + ? { enabled: true, maxDistance: spec.fuzzyDistance as 1 | 2 } + : undefined + }); + } catch (err) { + // Only tolerate expected highlight-alignment failures for a single + // candidate (scorer/highlight length gates, polarity mismatch); + // programmer errors and I/O faults must propagate. + if ( + err instanceof IncompatibleOptionError || + err instanceof ProfileMismatchError || + err instanceof RangeError + ) { + matchedRanges = []; + } else { + throw err; + } + } + out.push({ + text: c.text, + score: c.keys.score, + type: spec.mode === 'prefix' ? 'completion' : 'did-you-mean', + matchedRanges, + docId: c.keys.id, + doc: s.records[c.dIdx] as TDoc + }); + } + return out; +} + +/** + * First-party autocomplete / did-you-mean primitive (extracted verbatim + * from `DocumentIndex.autocomplete`). + */ +export async function runAutocomplete( + s: DocumentIndexState, + query: string, + options: AutocompleteOptions = {} +): Promise> { + if (typeof query !== 'string') { + throw new TypeError(`[webgpu-search] autocomplete expects query: string, got ${typeof query}.`); + } + const spec = normalizeAutocompleteOptions(options); + if (spec.field !== undefined && !s.fieldNameToIndex.has(spec.field)) { + throw new Error(`[webgpu-search] Unknown autocomplete field: "${spec.field}".`); + } + const t0 = nowMs(); + const normalized = normalizeText(query, s.normalized); + if (normalized.tokenCount > QUERY_TOKENS_MAX) { + throw new QueryTooLongError(QUERY_TOKENS_MAX, normalized.tokenCount, s.profileId); + } + if (normalized.isEmpty || s.idToDocIndex.size === 0) { + return { suggestions: [], queryDurationMs: nowMs() - t0 }; + } + const allowed = spec.field !== undefined + ? new Set([s.fieldNameToIndex.get(spec.field) as number]) + : undefined; + const suggestions = computeSuggestions(s, query, normalized.tokens, spec, allowed); + return { suggestions, queryDurationMs: nowMs() - t0 }; +} + +/** + * Full search execution (extracted verbatim from + * `DocumentIndex.search`, minus the facade-level destroyed/mutex/generation + * prologue which stays in `document-index.ts`). + * + * @param gen generation snapshot taken after the mutex wait; mid-flight + * mutations bump `generation` and abort in-flight searches. + */ +export async function runSearch( + s: DocumentIndexState, + query: string, + options: DocumentSearchOptions, + gen: number +): Promise> { + const validated = validateSearchOptions(s, query, options); + const { + mode, + caseSensitive, + signal, + cpuScorer, + typo, + effectiveHooks, + cpuModeOptions, + budget, + wantsDiagnostics, + queryStartMs, + diagWarnings, + normalizedQuery, + clampedLimit, + forceCpu, + facetSpec, + wantsFacets, + facetingMode, + facetEngine, + tieBreakers, + suggestSpec, + allowedFieldIndices, + hasEmptyFieldList, + customTokenTerms + } = validated; + + const filterState = compileFilterPhase( + s, + options.filter, + effectiveHooks, + wantsDiagnostics, + signal, + budget, + queryStartMs + ); + const { filterBitset, filterPredicate, structuredFilter, filteringMs } = filterState; + /** + * assemble response diagnostics (undefined unless requested). + * `filteringMs` is read at call time so early exits after compilation + * report measured filter cost; `totalMs` is wall-clock from query entry + * (end-to-end including inline autocomplete; see `autocompleteMs` for the slice). + * Scoring includes post-match scoring-hook time on all paths. + */ + const buildDiagnostics = ( + routedEngine: EngineType, + scoringMs: number, + highlightMs: number, + facetingMs: number | undefined, + hasOverflow: boolean, + autocompleteMs?: number + ): QueryDiagnostics | undefined => { + if (!wantsDiagnostics) return undefined; + const scanned = s.idToDocIndex.size; + const selectivity = filterBitset !== undefined + ? computeFilterSelectivity(filterBitset.popcount(), scanned) + : 1.0; + const totalMs = nowMs() - queryStartMs; + const timings: QueryDiagnosticsTimings = { + filteringMs, + scoringMs, + highlightMs, + ...(facetingMs !== undefined ? { facetingMs } : {}), + ...(autocompleteMs !== undefined ? { autocompleteMs } : {}), + totalMs + }; + const diag: QueryDiagnostics = { + scannedCandidates: scanned, + filterSelectivity: selectivity, + routedEngine, + hasOverflow, + timings + }; + if (diagWarnings.length > 0) diag.warnings = [...diagWarnings]; + return diag; + }; + /** + * post-hoc broad-query + overflow warnings (non-fatal, gated on + * `diagnostics:true` — routing decisions stay ungated for safety). + * `rawTotalMatches` carries the pre-filter pool count on the GPU path so + * the overflow sentence stays accurate; `facetsExact` suppresses the + * stale "approximate" claim after a force-exact rescan. + */ + const pushPostHocWarnings = ( + totalMatches: number, + hasOverflow: boolean, + opts?: { rawTotalMatches?: number; facetsRequested?: boolean; facetsExact?: boolean } + ): void => { + if (!wantsDiagnostics) return; + const scanned = s.idToDocIndex.size; + const selectivity = scanned > 0 ? totalMatches / scanned : 0; + if (isBroadSelectivity(selectivity, scanned)) { + diagWarnings.push(broadSelectivityWarning(selectivity, scanned)); + } + if (hasOverflow) { + diagWarnings.push( + candidateOverflowWarning(totalMatches, s.candidateCapacity, { + facetsRequested: opts?.facetsRequested ?? false, + facetsExact: opts?.facetsExact ?? false, + rawTotalMatches: opts?.rawTotalMatches, + }) + ); + } + }; + + const noHits = (q: string): DocumentSearchResponse => { + // even trivial exits honor caller aborts and time budgets. + throwIfAborted(signal); + throwIfBudgetAborted(budget); + assertTimeBudget(queryStartMs, budget); + const diag = buildDiagnostics('cpu', 0, 0, wantsFacets ? 0 : undefined, false); + // Public totalMs is scorer wall-clock (0 on no-hit: nothing scored); + // diagnostics.totalMs carries the wall-clock including validation. + return { + query: q, + mode, + engine: 'cpu', + totalMatches: 0, + candidateCount: 0, + hasOverflow: false, + results: [], + timings: { + queryUploadMs: 0, + encodeSubmitMs: 0, + gpuExecutionMs: null, + readbackMs: 0, + totalMs: 0, + gpuDispatchMs: 0 + }, + profileId: s.profileId, + scoringVersion: SCORING_VERSION, + cpuScorer, + fallbackReason: forceCpu ? 'query-too-long' : s.engineState.fallbackReason, + ...(wantsFacets && facetEngine && facetSpec + ? { facets: facetEngine.emptyResults(facetSpec) } + : {}), + ...(diag ? { diagnostics: diag } : {}) + }; + }; + + if (filterBitset !== undefined && filterBitset.isEmpty()) { + return noHits(query); + } + + if (normalizedQuery.isEmpty) { + return noHits(''); + } + if (s.idToDocIndex.size === 0) { + return noHits(query); + } + if (hasEmptyFieldList) { + return noHits(query); + } + + const route = routeEngine(s, { + mode, + typoEnabled: typo.enabled, + cpuScorer, + forceCpu, + allowedFieldIndices, + normalizedQueryTokenCount: normalizedQuery.tokens.length, + filterBitset, + budget, + queryStartMs, + wantsDiagnostics, + diagWarnings + }); + const { broadQueryCpuRoute, isGpuSupportedMode, useGpu, gpuHandle } = route; + + if (useGpu && gpuHandle !== null) { + // narrow GPU try to dispatch only so extension hook + // errors (filterPredicate / scoringHook / postProcess) propagate + // directly without destroying the engine or double-invoking via CPU + // fallback. Only dispatch failures fall through to CPU. + let gpuResult: Awaited> | undefined = undefined; + try { + gpuResult = await gpuHandle.search(query, { + ...options, + mode, + limit: s.candidateCapacity, + caseSensitive, + signal + }); + + if (s.isDestroyed || s.generation !== gen) { + throw abortError(); + } + throwIfAborted(signal); + // GPU dispatch is async wall-clock work — enforce the caller + // deadline and budget abort before touching the readback. + throwIfBudgetAborted(budget); + assertTimeBudget(queryStartMs, budget); + } catch (err: any) { + if (err.name === 'AbortError') { + throw err; + } + // Mirror search-index: mode-routed rejections (token/prefix/typo via + // direct or raced calls) are not engine failures — keep a healthy GPU + // up and fall through to exact CPU. + const modeRouted: boolean = + err instanceof IncompatibleOptionError && (err.option === 'mode' || err.option === 'typoTolerance'); + if (!modeRouted) { + console.warn('[webgpu-search] GPU document search failed, falling back to CPU:', err); + s.engineState = transitionEngineState(s.engineState, 'cpu', 'gpu-execution-error'); + if (s.gpuEngine) { + try { s.gpuEngine.destroy(); } catch {} + s.gpuEngine = null; + } + } + gpuResult = undefined; + } + + if (gpuResult !== undefined) { + const gpuResultOk = gpuResult; + + // WebGPU candidate readback & fixed-point score enrichment + const docMatches = new Map; + }>(); + // unfiltered query match set for facet aggregation (filter applied below). + const queryMatchedAll = wantsFacets ? new Set() : null; + + // hook errors propagate (no fallback / no double-invoke). + // Only GPU dispatch failures fall back; readback + hook pipeline runs + // outside the dispatch try above. + for (let i = 0; i < gpuResultOk.results.length; i++) { + const item = gpuResultOk.results[i]; + const r = item.index; + if (s.tombstones.has(r)) continue; + const fIdx = s.rowToFieldIndex[r]; + if (allowedFieldIndices && !allowedFieldIndices.has(fIdx)) continue; + const dIdx = s.rowToDocIndex[r]; + if (queryMatchedAll !== null && s.records[dIdx]) { + queryMatchedAll.add(dIdx); + } + if (filterBitset && !filterBitset.has(dIdx)) continue; + const doc = s.records[dIdx]; + if (!doc) continue; + if (filterPredicate && !filterPredicate(doc)) continue; + + const rawScore = item.score; + const fDef = s.sortedFields[fIdx]; + const weightedScore = Math.round(rawScore * fDef.weight); + + accumulateDocMatch(docMatches, dIdx, fIdx, weightedScore); + } + + const hits = aggregateDocMatches(s, docMatches, normalizedQuery.tokens, tieBreakers); + + const totalMatches = hits.length; + const candidateCount = Math.min(totalMatches, s.candidateCapacity); + let results = hits.slice(0, clampedLimit).map((h) => h.item); + const gpuEnriched = enrichResults(s, results, query, mode, options, customTokenTerms, normalizedQuery.tokens, effectiveHooks, tieBreakers, wantsDiagnostics); + results = gpuEnriched.results; + const gpuHighlightMs = gpuEnriched.highlightMs; + const gpuHookMs = gpuEnriched.hookMs; + + // facet aggregation. Exact when the GPU pool covered all matches; + // approximate over the top pool on overflow unless force-exact rescan. + let gpuFacets: Record | undefined = undefined; + let gpuFacetMs: number | undefined = undefined; + let gpuFacetExact = true; + if (wantsFacets && facetEngine && facetSpec && queryMatchedAll !== null) { + const facetT0 = nowMs(); + let facetCandidates: Iterable = queryMatchedAll; + let facetIsApproximate = gpuResultOk.hasOverflow; + if (gpuResultOk.hasOverflow && facetingMode === 'force-exact') { + const exact = scoreExactMatchesMultiField( + s.records.length, + s.sortedFields, + s.rowTokens, + s.rowToDocIndex, + s.rowToFieldIndex, + normalizedQuery.tokens, + mode, + clampedLimit, + s.candidateCapacity, + allowedFieldIndices, + s.tombstones, + undefined, + true, + cpuModeOptions + ); + facetCandidates = exact.allMatchedDocIndices ?? exact.results.map((r) => r.docIndex); + facetIsApproximate = false; + } + gpuFacets = buildFacetResults( + s, facetEngine, facetSpec, facetCandidates, structuredFilter, filterPredicate, facetIsApproximate + ); + const facetMs = nowMs() - facetT0; + gpuFacetMs = facetMs; + gpuFacetExact = !facetIsApproximate; + } + + // deadline + post-hoc warnings before returning GPU results. + // Inline autocomplete runs before the diagnostics snapshot so + // diagnostics.totalMs covers end-to-end latency (autocomplete bucket). + const tGpuSuggest0 = wantsDiagnostics && suggestSpec ? nowMs() : 0; + const gpuSuggestions = suggestSpec + ? computeSuggestions(s, query, normalizedQuery.tokens, suggestSpec, allowedFieldIndices) + : undefined; + const gpuSuggestMs = wantsDiagnostics && suggestSpec ? nowMs() - tGpuSuggest0 : undefined; + throwIfAborted(signal); + throwIfBudgetAborted(budget); + assertTimeBudget(queryStartMs, budget); + pushPostHocWarnings(totalMatches, gpuResultOk.hasOverflow, { + rawTotalMatches: gpuResultOk.totalMatches, + facetsRequested: wantsFacets, + facetsExact: gpuFacetExact, + }); + // Scoring = engine scan + post-match hooks (facet time excluded — + // no double-counting, no mutation of the engine timings object). + const gpuEngineTotal = gpuResultOk.timings.totalMs; + const gpuDiag = buildDiagnostics( + 'webgpu', + gpuEngineTotal + gpuHookMs, + gpuHighlightMs, + gpuFacetMs, + gpuResultOk.hasOverflow, + gpuSuggestMs + ); + + return { + query: gpuResultOk.query, + mode: gpuResultOk.mode, + engine: 'webgpu', + totalMatches, + candidateCount, + hasOverflow: gpuResultOk.hasOverflow, + results, + timings: { + ...gpuResultOk.timings, + totalMs: gpuEngineTotal + (gpuFacetMs ?? 0) + gpuHighlightMs + gpuHookMs, + }, + profileId: s.profileId, + scoringVersion: SCORING_VERSION, + cpuScorer, + ...(gpuFacets ? { facets: gpuFacets } : {}), + ...(gpuDiag ? { diagnostics: gpuDiag } : {}), + ...(gpuSuggestions ? { suggestions: gpuSuggestions } : {}) + }; + } + } + + // 2. CPU execution path + if (s.isDestroyed || s.generation !== gen) { + throw abortError(); + } + throwIfAborted(signal); + throwIfBudgetAborted(budget); + assertTimeBudget(queryStartMs, budget); + const t0 = nowMs(); + + if (cpuScorer === 'ufuzzy') { + const cpuEngine = new CPUEngine(); + const docMatches = new Map; + }>(); + // unfiltered query match set for exact facet aggregation. + const legacyMatchedAll = wantsFacets ? new Set() : null; + + for (let fIdx = 0; fIdx < s.sortedFields.length; fIdx++) { + if (allowedFieldIndices && !allowedFieldIndices.has(fIdx)) continue; + const fDef = s.sortedFields[fIdx]; + const fieldStrings = s.records.map((doc, dIdx) => (doc && s.rawFieldStrings[dIdx] ? s.rawFieldStrings[dIdx][fIdx] : '')); + const legacyResult = mode === 'fuzzy' + ? cpuEngine.searchWithUFuzzy(fieldStrings, query, s.records.length, caseSensitive) + : cpuEngine.searchNaiveScan(fieldStrings, query, s.records.length, caseSensitive); + + for (const hit of legacyResult.results) { + const dIdx = hit.index; + // Explicit tombstone guard for symmetry with GPU/exact paths: + // removed docs have null records (columnar presence also cleared). + if (!s.records[dIdx] || s.docIds[dIdx] === null || s.docIds[dIdx] === undefined) continue; + if (legacyMatchedAll !== null) { + legacyMatchedAll.add(dIdx); + } + if (filterBitset && !filterBitset.has(dIdx)) continue; + const doc = s.records[dIdx]; + if (!doc) continue; + if (filterPredicate && !filterPredicate(doc)) continue; + const weightedScore = Math.round(hit.score * fDef.weight); + accumulateDocMatch(docMatches, dIdx, fIdx, weightedScore); + } + } + + const hits = aggregateDocMatches(s, docMatches, normalizedQuery.tokens, tieBreakers); + + const totalMatches = hits.length; + const candidateCount = Math.min(totalMatches, s.candidateCapacity); + let results = hits.slice(0, clampedLimit).map((h) => h.item); + const legacyEnriched = enrichResults(s, results, query, mode, options, customTokenTerms, normalizedQuery.tokens, effectiveHooks, tieBreakers, wantsDiagnostics); + results = legacyEnriched.results; + const legacyHighlightMs = legacyEnriched.highlightMs; + const preFacetMs = nowMs() - t0; + // scoring bucket covers scan + ranking + post-match hooks; + // highlight enrichment is metered separately above. + const legacyScoringMs = Math.max(0, preFacetMs - legacyHighlightMs); + + // CPU evaluates the full match set, so legacy facets are exact + // w.r.t. the serving (ufuzzy/native) match set; see types for caveat. + let legacyFacets: Record | undefined = undefined; + let legacyFacetMs = 0; + if (wantsFacets && facetEngine && facetSpec && legacyMatchedAll !== null) { + const fT0 = nowMs(); + legacyFacets = buildFacetResults( + s, facetEngine, facetSpec, legacyMatchedAll, structuredFilter, filterPredicate, false + ); + legacyFacetMs = nowMs() - fT0; + } + + // inline autocomplete precedes the diagnostics snapshot so + // diagnostics.totalMs covers end-to-end latency (autocompleteMs bucket). + const tLegacySuggest0 = wantsDiagnostics && suggestSpec ? nowMs() : 0; + const legacySuggestions = suggestSpec + ? computeSuggestions(s, query, normalizedQuery.tokens, suggestSpec, allowedFieldIndices) + : undefined; + const legacySuggestMs = wantsDiagnostics && suggestSpec ? nowMs() - tLegacySuggest0 : undefined; + throwIfAborted(signal); + throwIfBudgetAborted(budget); + assertTimeBudget(queryStartMs, budget); + pushPostHocWarnings(totalMatches, totalMatches > s.candidateCapacity, { + facetsRequested: wantsFacets, + facetsExact: true, + }); + const legacyDiag = buildDiagnostics( + 'cpu', + legacyScoringMs, + legacyHighlightMs, + wantsFacets ? legacyFacetMs : undefined, + totalMatches > s.candidateCapacity, + legacySuggestMs + ); + + return { + query, + mode, + engine: 'cpu', + totalMatches, + candidateCount, + hasOverflow: totalMatches > s.candidateCapacity, + results, + timings: { + queryUploadMs: 0, + encodeSubmitMs: 0, + gpuExecutionMs: null, + readbackMs: 0, + totalMs: preFacetMs + legacyFacetMs, + gpuDispatchMs: 0 + }, + profileId: s.profileId, + scoringVersion: SCORING_VERSION, + cpuScorer, + fallbackReason: 'cpu-algorithm-requested', + ...(legacyFacets ? { facets: legacyFacets } : {}), + ...(legacyDiag ? { diagnostics: legacyDiag } : {}), + ...(legacySuggestions ? { suggestions: legacySuggestions } : {}) + }; + } + + // Default exact CPU algorithm (collect full match set only when faceting). + // passes the deterministic hierarchy + doc IDs so pre-truncation order + // already reflects (score, weight, exact, length, id). + const parityResult = scoreExactMatchesMultiField( + s.records.length, + s.sortedFields, + s.rowTokens, + s.rowToDocIndex, + s.rowToFieldIndex, + normalizedQuery.tokens, + mode, + clampedLimit, + s.candidateCapacity, + allowedFieldIndices, + s.tombstones, + (filterBitset || filterPredicate) ? (dIdx) => { + if (filterBitset && !filterBitset.has(dIdx)) return false; + if (filterPredicate) { + const doc = s.records[dIdx]; + return doc !== null && doc !== undefined && filterPredicate(doc); + } + return true; + } : undefined, + wantsFacets, + cpuModeOptions, + { tieBreakers, docIds: s.docIds as unknown as ReadonlyArray } + ); + + if (s.isDestroyed || s.generation !== gen) { + throw abortError(); + } + throwIfAborted(signal); + throwIfBudgetAborted(budget); + assertTimeBudget(queryStartMs, budget); + + const idProp = typeof s.options.idField === 'string' ? s.options.idField : 'id'; + let enrichedResults: DocumentSearchResultItem[] = []; + for (let i = 0; i < parityResult.results.length; i++) { + const hit = parityResult.results[i]; + const id = s.docIds[hit.docIndex]; + if (id === null || id === undefined) continue; + const doc = s.records[hit.docIndex] ?? ({ [idProp]: id } as any); + enrichedResults.push({ + id, + doc, + score: hit.score, + matchedField: hit.matchedField, + matches: hit.matches + }); + } + const parityEnriched = enrichResults(s, enrichedResults, query, mode, options, customTokenTerms, normalizedQuery.tokens, effectiveHooks, tieBreakers, wantsDiagnostics); + enrichedResults = parityEnriched.results; + const parityHighlightMs = parityEnriched.highlightMs; + const parityHookMs = parityEnriched.hookMs; + const parityScoringMs = parityResult.durationMs + parityHookMs; + + // exact facet aggregation. Disjunctive facets need the unfiltered + // query match set, so a structured filter triggers one extra unfiltered + // parity scan (results above stay filtered and correctly ranked; O(F*n) + // per-facet re-evaluation + F bitset allocs by design for disjunctive + // semantics). Wall-clock facet work is included in totalMs; future + // attribute facet work to diagnostics.timings.facetingMs. + let parityFacets: Record | undefined = undefined; + const facetT0 = nowMs(); + let facetScanMs = 0; + if (wantsFacets && facetEngine && facetSpec) { + let facetCandidates: Iterable = + parityResult.allMatchedDocIndices ?? parityResult.results.map((r) => r.docIndex); + if (structuredFilter !== undefined) { + const unfiltered = scoreExactMatchesMultiField( + s.records.length, + s.sortedFields, + s.rowTokens, + s.rowToDocIndex, + s.rowToFieldIndex, + normalizedQuery.tokens, + mode, + clampedLimit, + s.candidateCapacity, + allowedFieldIndices, + s.tombstones, + undefined, + true, + cpuModeOptions + ); + facetCandidates = unfiltered.allMatchedDocIndices ?? unfiltered.results.map((r) => r.docIndex); + facetScanMs = unfiltered.durationMs; + } + parityFacets = buildFacetResults( + s, facetEngine, facetSpec, facetCandidates, structuredFilter, filterPredicate, false + ); + } + const facetMs = wantsFacets ? Math.max(0, nowMs() - facetT0 - facetScanMs) : 0; + // the disjunctive unfiltered rescan serves facets, so it joins the + // faceting bucket (scoring stays scan + post-match hooks). + const parityFacetingMs = wantsFacets ? facetScanMs + facetMs : undefined; + + // Public totalMs is scorer wall-clock: scan + hooks + highlight + facets + // (excludes inline autocomplete; diagnostics.totalMs is end-to-end). + const timings: SearchTimings = { + queryUploadMs: 0, + encodeSubmitMs: 0, + gpuExecutionMs: null, + readbackMs: 0, + totalMs: parityScoringMs + parityHighlightMs + facetScanMs + facetMs, + gpuDispatchMs: 0 + }; + + let effectiveFallbackReason = s.engineState.fallbackReason; + if (forceCpu) { + effectiveFallbackReason = 'query-too-long'; + } else if (!isGpuSupportedMode) { + // token/prefix/typo queries are CPU-by-design (exact-only + // WGSL kernels) — recorded per the exact boundary. + effectiveFallbackReason = 'unsupported-mode'; + } else if (useGpu && gpuHandle !== null) { + effectiveFallbackReason = 'gpu-execution-error'; + } else if (broadQueryCpuRoute) { + // broad-query CPU routing is a routing decision (like + // field-restriction below), not a scorer request — leave the reason + // as-is and surface the decision in diagnostics.warnings instead. + } + // Note: field-restricted routing (isFieldRestricted) intentionally leaves + // fallbackReason as-is — it is a routing decision, not a scorer request, + // so reusing 'cpu-algorithm-requested' would mislead telemetry. + + // inline autocomplete precedes the diagnostics snapshot so + // diagnostics.totalMs covers end-to-end latency (autocompleteMs bucket). + const tParitySuggest0 = wantsDiagnostics && suggestSpec ? nowMs() : 0; + const paritySuggestions = suggestSpec + ? computeSuggestions(s, query, normalizedQuery.tokens, suggestSpec, allowedFieldIndices) + : undefined; + const paritySuggestMs = wantsDiagnostics && suggestSpec ? nowMs() - tParitySuggest0 : undefined; + throwIfAborted(signal); + throwIfBudgetAborted(budget); + assertTimeBudget(queryStartMs, budget); + pushPostHocWarnings(parityResult.totalMatches, parityResult.hasOverflow, { + facetsRequested: wantsFacets, + facetsExact: true, + }); + const parityDiag = buildDiagnostics( + 'cpu', + parityScoringMs, + parityHighlightMs, + parityFacetingMs, + parityResult.hasOverflow, + paritySuggestMs + ); + + return { + query, + mode, + engine: 'cpu', + totalMatches: parityResult.totalMatches, + candidateCount: parityResult.candidateCount, + hasOverflow: parityResult.hasOverflow, + results: enrichedResults, + timings, + profileId: s.profileId, + scoringVersion: SCORING_VERSION, + cpuScorer, + fallbackReason: effectiveFallbackReason, + ...(parityFacets ? { facets: parityFacets } : {}), + ...(parityDiag ? { diagnostics: parityDiag } : {}), + ...(paritySuggestions ? { suggestions: paritySuggestions } : {}) + }; +} diff --git a/packages/webgpu-search/src/document/snapshot.ts b/packages/webgpu-search/src/document/snapshot.ts new file mode 100644 index 0000000..4dba716 --- /dev/null +++ b/packages/webgpu-search/src/document/snapshot.ts @@ -0,0 +1,522 @@ +/** + * Snapshot restore/apply implementation for `DocumentIndex` (extracted + * verbatim from the pre-refactor monolithic `document-index.ts`). + * + * Scope of this module: + * - `assertSnapshotFilterGettersSatisfied`: fail-closed custom filter-getter + * restore guard (shared by static `fromSnapshotData` and instance restores; + * re-exported from the public `document-index.ts` facade). + * - `mergeSnapshotRestoreOptions`: pure merge of a reset snapshot + user + * restore options into construction options (includes the profile/getter/ + * hook fail-closed guards for the static restore path). + * - `configureFromSchema` / `applySnapshotData`: rebuild live index state + * from a decoded snapshot, including the CPU re-platform and the + * GPU bootstrap attempt. + * + * This module must never import the `DocumentIndex` facade at runtime + * (the facade imports from here); static restore constructs the instance + * in `document-index.ts` before calling `applySnapshotData`. + */ + +import { readOwnField, ColumnarStore } from '../filtering/columnar-store'; +import { packDataset, checkMemoryBudget } from '../dataset-packing'; +import { transitionEngineState } from '../engine-state'; +import { WebGPUEngine } from '../webgpu-engine'; +import { assertValidPowerPreference } from '../gpu-device-pool'; +import { normalizeSearchHooks, assertHooksSatisfied } from '../hooks'; +import { IncompatibleIndexError, IncompatibleOptionError, ProfileMismatchError } from '../text-profile'; +import type { RestoredDocumentSnapshot } from '../snapshot-codec'; +import type { + DocumentId, + DocumentIndexOptions, + DocumentIndexSchema, + FilterFieldDefinition, + FilterFieldType, + RestoreDocumentIndexOptions, + SearchHooks, +} from '../types'; +import { subscribeDeviceLost } from './build'; +import type { DocumentIndexState, InternalField } from './state'; + +/** + * fail-closed custom filter-getter restore guard (mirrors the + * hookIds idiom). Snapshots recording `hasGetter: true` for a filter field + * require the caller to supply a matching getter override via + * `options.options.filterFields`; otherwise columnar rebuild via default + * `doc[name]` would silently change filter semantics. + */ +export function assertSnapshotFilterGettersSatisfied( + snapshotFilterFields: Array<{ name: string; hasGetter?: boolean }> | undefined, + userFilterFields: Array | undefined +): void { + if (!snapshotFilterFields || snapshotFilterFields.length === 0) return; + const needed = snapshotFilterFields.filter((ff) => ff && ff.hasGetter === true); + if (needed.length === 0) return; + const userGetterNames = new Set(); + if (Array.isArray(userFilterFields)) { + for (const uf of userFilterFields) { + if (uf && typeof uf === 'object' && typeof (uf as { name?: unknown }).name === 'string') { + const u = uf as { name: string; getter?: unknown }; + if (typeof u.getter === 'function') userGetterNames.add(u.name); + } + } + } + for (const ff of needed) { + if (!userGetterNames.has(ff.name)) { + throw new IncompatibleIndexError( + `filter-getter:${ff.name}`, + 'missing custom getter override for snapshot filter field (supply options.options.filterFields with getter)' + ); + } + } +} + +/** + * Merge a decoded snapshot with user restore options into construction + * options (verbatim from the pre-refactor `fromSnapshotData` data-merge + * block; instance construction stays in the facade). Includes the fail-closed + * profile guards, the filter-getter guard, and the hook restore guard. + */ +export function mergeSnapshotRestoreOptions>( + snapshot: RestoredDocumentSnapshot, + options?: RestoreDocumentIndexOptions +): DocumentIndexOptions { + // Fail-closed profile guards (never read under the wrong profile). + // `textProfile` only admits 'unicode-default'; any other value throws + // `ProfileMismatchError`. An explicit `caseSensitive` override must agree + // with the snapshot polarity (`!header.normalized`); otherwise the + // retained token stream would be queried under the wrong normalization. + const rawTextProfile = (options?.options as { textProfile?: unknown } | undefined)?.textProfile; + if (rawTextProfile !== undefined && rawTextProfile !== 'unicode-default') { + throw new ProfileMismatchError('unicode-default', rawTextProfile, 'textProfile'); + } + const overrideCase = options?.options?.caseSensitive; + if (overrideCase !== undefined && overrideCase === snapshot.header.normalized) { + throw new ProfileMismatchError(!snapshot.header.normalized, overrideCase, 'caseSensitive'); + } + // fail-closed custom filter-getter guard. Snapshot `hasGetter` + // entries require a matching getter override; otherwise restore would + // silently rebuild columnar via default `doc[name]` (presence cleared). + assertSnapshotFilterGettersSatisfied( + snapshot.schema.filterFields as Array<{ name: string; hasGetter?: boolean }> | undefined, + options?.options?.filterFields as Array | undefined + ); + const userFields = options?.options?.fields; + const userFieldMap = new Map(); + if (Array.isArray(userFields)) { + for (const uf of userFields) { + if (typeof uf === 'string') { + userFieldMap.set(uf, uf); + } else if (uf && typeof uf === 'object' && typeof uf.name === 'string') { + userFieldMap.set(uf.name, uf); + } + } + } + + const resolvedFields = snapshot.schema.fields.map((sf) => { + const uf = userFieldMap.get(sf.name); + const getter = typeof uf === 'object' && uf !== null && typeof uf.getter === 'function' + ? uf.getter + : undefined; + return { + name: sf.name, + weight: sf.weight, + getter + }; + }); + + const mergedOptions: DocumentIndexOptions = { + ...(options?.options as any), + fields: resolvedFields, + idField: options?.options?.idField ?? snapshot.schema.idField ?? 'id', + caseSensitive: options?.options?.caseSensitive ?? snapshot.schema.caseSensitive ?? !(snapshot.header.normalized), + device: (options?.device ?? options?.options?.device) as GPUDevice | undefined, + preferGpu: options?.options?.preferGpu ?? snapshot.schema.preferGpu, + threshold: options?.options?.threshold ?? snapshot.schema.threshold, + candidateCapacity: options?.options?.candidateCapacity ?? snapshot.schema.candidateCapacity, + initialCapacity: options?.options?.initialCapacity ?? snapshot.schema.initialCapacity, + growthFactor: options?.options?.growthFactor ?? snapshot.schema.growthFactor, + filterFields: options?.options?.filterFields ?? (snapshot.schema.filterFields as any), + hooks: options?.options?.hooks + }; + + // fail-closed hook restore guard. Snapshots recording hookIds + // require matching handlers via `options.options.hooks`; closures + // are never serialized, only declarative IDs. + assertHooksSatisfied( + snapshot.schema.hookIds, + normalizeSearchHooks(mergedOptions.hooks) + ); + + return mergedOptions; +} + +/** + * Rebuild the field/id accessors from a snapshot schema (verbatim from the + * pre-refactor `configureFromSchema`). User-supplied field getters take + * precedence over the default own-property read. + */ +export function configureFromSchema( + s: DocumentIndexState, + schema: DocumentIndexSchema, + options?: RestoreDocumentIndexOptions +): void { + const userFields = options?.options?.fields; + const userFieldMap = new Map(); + if (Array.isArray(userFields)) { + for (const uf of userFields) { + if (typeof uf === 'string') { + userFieldMap.set(uf, uf); + } else if (uf && typeof uf === 'object' && typeof uf.name === 'string') { + userFieldMap.set(uf.name, uf); + } + } + } + + const normalized: InternalField[] = schema.fields.map( + (f: { name: string; weight: number }, originalIndex: number) => { + const name = f.name; + if (typeof name !== 'string' || name.length === 0) { + throw new TypeError('[webgpu-search] Restored schema field requires a non-empty string name.'); + } + const weight = f.weight ?? 1.0; + if (typeof weight !== 'number' || !Number.isFinite(weight) || weight <= 0) { + throw new RangeError( + `[webgpu-search] Restored field "${name}" weight must be a positive finite number, got ${String(weight)}.` + ); + } + const uf = userFieldMap.get(name); + let getter: (doc: TDoc) => string | string[] | undefined | null; + if (uf && typeof uf === 'object' && typeof uf.getter === 'function') { + getter = uf.getter; + } else { + getter = (doc: any) => readOwnField(doc, name); + } + return { name, weight, getter, originalIndex }; + }); + + s.sortedFields = [...normalized].sort((a, b) => { + if (b.weight !== a.weight) return b.weight > a.weight ? 1 : -1; + if (a.originalIndex !== b.originalIndex) return a.originalIndex > b.originalIndex ? 1 : -1; + return 0; + }); + + s.fieldNameToIndex.clear(); + for (let i = 0; i < s.sortedFields.length; i++) { + s.fieldNameToIndex.set(s.sortedFields[i].name, i); + } + + if (typeof options?.options?.idField === 'function') { + s.getId = options.options.idField; + } else if (typeof options?.options?.idField === 'string') { + const prop = options.options.idField; + s.getId = (doc: any) => doc[prop]; + } else if (typeof schema.idField === 'string') { + const prop = schema.idField; + s.getId = (doc: any) => doc[prop]; + } else if (options?.options?.idField === undefined && schema.idField === undefined) { + s.getId = (doc: any) => doc.id; + } +} + +/** + * Apply decoded snapshot data to live index state (verbatim from the + * pre-refactor `applySnapshotData`). Assumes the mutation mutex is held and + * the constructor-time profile guards have already run (instance restores in + * the facade, static restores in `fromSnapshotData`). Ends by bumping + * `generation` so in-flight searches observe the replace. + */ +export async function applySnapshotData( + s: DocumentIndexState, + snapshot: RestoredDocumentSnapshot, + options?: RestoreDocumentIndexOptions +): Promise { + if ((options?.options as unknown as Record | undefined)?.extensions !== undefined) { + throw new IncompatibleOptionError( + 'hooks', + "[webgpu-search] extensions was removed; use hooks." + ); + } + // Fail-closed profile guards (never read under the wrong profile). + // The live index normalization is fixed at construction; a snapshot + // built under the opposite polarity must be rejected, not reinterpreted. + // An explicit `caseSensitive`/`textProfile` override must agree with both + // the live index and the snapshot. + const rawTextProfile = (options?.options as { textProfile?: unknown } | undefined)?.textProfile; + if (rawTextProfile !== undefined && rawTextProfile !== 'unicode-default') { + throw new ProfileMismatchError('unicode-default', rawTextProfile, 'textProfile'); + } + const snapshotCaseSensitive = !snapshot.header.normalized; + const liveCaseSensitive = !s.normalized; + if (snapshotCaseSensitive !== liveCaseSensitive) { + throw new ProfileMismatchError(liveCaseSensitive, snapshotCaseSensitive, 'caseSensitive'); + } + const overrideCaseApply = options?.options?.caseSensitive; + if (overrideCaseApply !== undefined && overrideCaseApply !== snapshotCaseSensitive) { + throw new ProfileMismatchError(snapshotCaseSensitive, overrideCaseApply, 'caseSensitive'); + } + // fail-closed hook restore guard for instance `restore()`. + // Restore-supplied handlers (if any) replace the live index hooks; + // otherwise the live hooks must satisfy the snapshot. + const restoreHooks = normalizeSearchHooks(options?.options?.hooks); + if (restoreHooks !== undefined) { + s.indexHooks = restoreHooks; + (s.options as { hooks?: SearchHooks }).hooks = restoreHooks; + } + assertHooksSatisfied(snapshot.schema.hookIds, s.indexHooks); + + // fail-closed custom filter-getter guard for instance restore. + // Live filter defs (plus any restore override) must supply getters for + // snapshot `hasGetter` fields; otherwise columnar rebuild would silently + // change semantics. When the override supplies the getters, adopt it so + // `columnarStore.init` uses the correct accessors. + { + const overrideFF = options?.options?.filterFields as + | Array any }> + | undefined; + const liveCustom = new Set(); + for (const ff of s.filterFieldDefinitions) { + if ((ff as { hasGetter?: boolean }).hasGetter === true) liveCustom.add(ff.name as string); + } + const overrideGetterNames = new Set(); + if (Array.isArray(overrideFF)) { + for (const uf of overrideFF) { + if (uf && typeof uf === 'object' && typeof (uf as { getter?: unknown }).getter === 'function') { + overrideGetterNames.add((uf as { name: string }).name); + } + } + } + const needed = ((snapshot.schema.filterFields as Array<{ name: string; hasGetter?: boolean }> | undefined) ?? []).filter( + (ff) => ff && ff.hasGetter === true + ); + for (const ff of needed) { + if (!liveCustom.has(ff.name) && !overrideGetterNames.has(ff.name)) { + throw new IncompatibleIndexError( + `filter-getter:${ff.name}`, + 'missing custom getter override for snapshot filter field (supply options.options.filterFields with getter)' + ); + } + } + if (Array.isArray(overrideFF) && overrideFF.length > 0) { + const normalized: FilterFieldDefinition[] = []; + const seen = new Set(); + for (const uf of overrideFF) { + if (typeof uf === 'string') { + if (seen.has(uf)) continue; + seen.add(uf); + normalized.push({ name: uf as never, getter: ((doc: any) => (doc as any)[uf]) as never }); + } else if (uf && typeof uf === 'object' && typeof (uf as { name?: unknown }).name === 'string') { + const o = uf as { name: string; type?: FilterFieldType; getter?: (doc: TDoc) => any }; + if (seen.has(o.name)) continue; + seen.add(o.name); + const hasCustom = typeof o.getter === 'function'; + normalized.push( + hasCustom + ? { name: o.name as never, type: o.type, getter: o.getter as never, hasGetter: true } + : { name: o.name as never, type: o.type, getter: ((doc: any) => (doc as any)[o.name]) as never } + ); + } + } + if (normalized.length > 0) { + s.filterFieldDefinitions = normalized; + s.columnarStore = new ColumnarStore(normalized, { + initialCapacity: s.initialCapacity, + growthFactor: s.growthFactor + }); + } + } + } + + configureFromSchema(s, snapshot.schema, options); + + if (s.unsubscribeDeviceLost) { + s.unsubscribeDeviceLost(); + s.unsubscribeDeviceLost = undefined; + } + if (s.gpuEngine) { + s.gpuEngine.destroy(); + s.gpuEngine = null; + } + s.tombstones.clear(); + s.mutationEpoch = + typeof snapshot.schema.mutationEpoch === 'number' && + Number.isFinite(snapshot.schema.mutationEpoch) && + snapshot.schema.mutationEpoch >= 0 + ? Math.floor(snapshot.schema.mutationEpoch) + : 0; + + const docCount = snapshot.header.docCount; + const fieldCount = s.sortedFields.length; + const totalRows = snapshot.header.rowCount; + + s.records = snapshot.records.slice(); + s.docIds = snapshot.docIds.slice(); + s.idToDocIndex.clear(); + + if (s.records.length > 0 && s.docIds.length === 0) { + s.docIds = new Array(docCount); + for (let i = 0; i < docCount; i++) { + const id = s.getId(s.records[i]); + s.docIds[i] = id; + s.idToDocIndex.set(id, i); + } + } else { + for (let i = 0; i < s.docIds.length; i++) { + s.idToDocIndex.set(s.docIds[i], i); + } + } + + // If both records and docIds are present, ensure exact ID correspondence alignment + if (s.records.length === docCount && s.docIds.length === docCount) { + const recMap = new Map(); + for (const rec of s.records) { + recMap.set(s.getId(rec), rec); + } + if (recMap.size === docCount) { + const aligned: TDoc[] = new Array(docCount); + let allFound = true; + for (let i = 0; i < docCount; i++) { + const rec = recMap.get(s.docIds[i]); + if (rec === undefined) { + allFound = false; + break; + } + aligned[i] = rec; + } + if (allFound) { + s.records = aligned; + } + } + } + + if (s.records.length === 0 && docCount > 0) { + const idProp = typeof s.options.idField === 'string' ? s.options.idField : 'id'; + s.records = new Array(docCount); + for (let i = 0; i < docCount; i++) { + s.records[i] = { [idProp]: s.docIds[i] } as any; + } + } + + s.columnarStore.init(s.records); + + const rawCap = s.options.candidateCapacity; + const validUserCap = typeof rawCap === 'number' && Number.isFinite(rawCap) && rawCap > 0 + ? Math.floor(rawCap) + : Math.floor(docCount * fieldCount * 0.1); + s.candidateCapacity = Math.min(32768, Math.max(8192, validUserCap)); + + s.rowTokens = new Array(totalRows); + s.rowToDocIndex = new Array(totalRows); + s.rowToFieldIndex = new Array(totalRows); + s.docToRowIndices = Array.from({ length: docCount }, () => new Array(fieldCount)); + s.rawFieldStrings = Array.from({ length: docCount }, () => new Array(fieldCount)); + + for (let r = 0; r < totalRows; r++) { + const start = snapshot.offsets[r]; + const end = snapshot.offsets[r + 1]; + s.rowTokens[r] = snapshot.tokens.subarray(start, end); + + const f = fieldCount > 0 && docCount > 0 ? Math.floor(r / docCount) : 0; + const d = docCount > 0 ? r % docCount : 0; + s.rowToDocIndex[r] = d; + s.rowToFieldIndex[r] = f; + if (s.docToRowIndices[d]) { + s.docToRowIndices[d][f] = r; + } + } + + s.totalTokens = snapshot.header.tokenCount; + + if (s.records.length === docCount && docCount > 0) { + for (let d = 0; d < docCount; d++) { + const doc = s.records[d]; + for (let f = 0; f < fieldCount; f++) { + const field = s.sortedFields[f]; + const rawVal = field.getter(doc); + let rawStr = ''; + if (rawVal !== null && rawVal !== undefined) { + if (Array.isArray(rawVal)) { + rawStr = rawVal.filter((x) => x !== null && x !== undefined).join(' '); + } else { + rawStr = typeof rawVal === 'string' ? rawVal : String(rawVal); + } + } + s.rawFieldStrings[d][f] = rawStr; + } + } + } + + s.engineState = transitionEngineState( + s.engineState, + 'cpu', + s.engineState.fallbackReason, + ); + s.vramAllocatedBytes = 0; + + if (docCount > 0) { + const threshold = s.options.threshold ?? 30_000; + const shouldAttemptGpu = s.options.preferGpu === false + ? false + : s.preferGpu || totalRows >= threshold; + + if (s.options.preferGpu === false) { + s.engineState = transitionEngineState(s.engineState, 'cpu', 'prefer-cpu'); + } else if (!s.preferGpu && totalRows < threshold) { + s.engineState = transitionEngineState(s.engineState, 'cpu', 'below-threshold'); + } + + if (shouldAttemptGpu) { + const measuredAvgBytes = totalRows > 0 ? (s.totalTokens * 4) / totalRows : 0; + const deviceToUse = options?.device ?? s.options.device; + const powerPreference = options?.options?.powerPreference ?? s.options.powerPreference; + assertValidPowerPreference(powerPreference); + const budget = checkMemoryBudget(totalRows, measuredAvgBytes, deviceToUse, s.candidateCapacity); + if (!budget.allowed) { + s.engineState = transitionEngineState(s.engineState, 'cpu', 'memory-budget-exceeded'); + } else { + try { + const gpu = new WebGPUEngine(); + const initialized = await gpu.init( + deviceToUse !== undefined || powerPreference !== undefined + ? { device: deviceToUse ?? undefined, powerPreference } + : undefined + ); + if (initialized && gpu.isReady) { + gpu.ensureCandidateCapacity(s.candidateCapacity); + const packed = packDataset(s.rowTokens, { + normalized: s.normalized, + totalTokens: s.totalTokens + }); + const effectiveCap = s.initialCapacity > 0 ? s.initialCapacity : Math.max(16, docCount); + const rowCap = effectiveCap * fieldCount; + const avgTokens = totalRows > 0 ? s.totalTokens / totalRows : 16; + const tokenCap = Math.max(Math.floor(rowCap * avgTokens), 64); + await gpu.loadDataset(packed, { + rowCapacity: rowCap, + tokenCapacity: tokenCap, + growthFactor: s.initialCapacity > 0 ? 1.0 : s.growthFactor + }); + + s.gpuEngine = gpu; + s.engineState = transitionEngineState(s.engineState, 'webgpu'); + s.vramAllocatedBytes = gpu.vramAllocatedBytes; + + subscribeDeviceLost(s); + } else { + s.engineState = transitionEngineState(s.engineState, 'cpu', 'device-request-failed'); + } + } catch { + s.engineState = transitionEngineState( + s.engineState, + 'cpu', + typeof navigator === 'undefined' || !('gpu' in navigator) || !navigator.gpu + ? 'webgpu-unsupported' + : 'device-request-failed', + ); + } + } + } + } + + s.generation++; +} diff --git a/packages/webgpu-search/src/document/state.ts b/packages/webgpu-search/src/document/state.ts new file mode 100644 index 0000000..8ff71f0 --- /dev/null +++ b/packages/webgpu-search/src/document/state.ts @@ -0,0 +1,106 @@ +/** + * Internal shared-state contract for the decomposed document-index modules. + * + * The public `DocumentIndex` facade in `document-index.ts` owns lifecycle + * concerns (construction-time fail-closed validation, mutation mutex + * scheduling, destroy) and delegates the implementation to the sibling + * modules in `document/`: + * - `build.ts` : ingestion/index building, batch mutations, GPU engine + * bootstrap, compaction (the imperative mutation shell). + * - `query.ts` : `search`, `autocomplete`, inline suggestions, and + * highlight enrichment. + * - `filter-facet.ts`: search filter compilation + facet aggregation. + * - `snapshot.ts` : snapshot application/restoration. + * + * `DocumentIndexState` is an internal implementation detail only: it is + * never exported from the package entrypoint and never exposed as public + * API. Modules mutate its fields directly, exactly as the pre-refactor + * private class members did — including the `null`-sentinel slots for + * tombstoned documents. Keeping the mutable fields in one explicit + * interface (instead of prototype tricks) keeps the cross-module contract + * type-checked and cohesive. + */ + +import type { EngineState } from '../engine-state'; +import type { WebGPUEngine } from '../webgpu-engine'; +import type { ColumnarStore } from '../filtering/columnar-store'; +import type { FieldScoreDefinition } from '../exact-scorer'; +import type { TextProfileId } from '../text-profile'; +import type { + DocumentId, + DocumentIndexOptions, + FilterFieldDefinition, + SearchHooks, +} from '../types'; + +/** + * Normalized search-field definition (weight-stratified). Reserved for + * internal modules; the public `getSortedFields()` accessor returns live + * instances of this shape. + */ +export interface InternalField extends FieldScoreDefinition { + name: string; + weight: number; + getter: (doc: TDoc) => string | string[] | undefined | null; + originalIndex: number; +} + +/** + * All mutable document-index internals, grouped into one explicit, + * type-checked bundle owned by the facade and passed to the `document/` + * modules by reference. + */ +export interface DocumentIndexState { + /** User options (kept by reference; `restore` may write hooks/device back into it). */ + options: DocumentIndexOptions; + + /** + * Frozen normalization polarity: token streams are case-folded when + * `normalized` is `true` (i.e. `options.caseSensitive` is falsy). + */ + normalized: boolean; + profileId: TextProfileId; + preferGpu: boolean; + initialCapacity: number; + growthFactor: number; + + // Active documents state + records: TDoc[]; + docIds: DocumentId[]; + idToDocIndex: Map; + getId: (doc: TDoc) => DocumentId; + + // Fields and weights (stratified: sorted by weight desc) + sortedFields: InternalField[]; + fieldNameToIndex: Map; + + // Tokens and row mapping + rowTokens: Uint32Array[]; + rowToDocIndex: number[]; + rowToFieldIndex: number[]; + docToRowIndices: number[][]; // [docIndex][fieldIndex] -> row index + rawFieldStrings: string[][]; // [docIndex][fieldIndex] + totalTokens: number; + candidateCapacity: number; + buildTimeMs: number; + restoreTimeMs?: number; + lastMutationTimeMs?: number; + mutationEpoch: number; + tombstones: Set; // tombstone candidate filtering + + // Concurrency + lifecycle (owned by the facade, read by modules) + searchMutex: Promise; + generation: number; + isDestroyed: boolean; + + // Filter columnar + hooks + columnarStore: ColumnarStore; + filterFieldDefinitions: FilterFieldDefinition[]; + indexHooks: SearchHooks | undefined; + + // Engine lifecycle + engineState: EngineState; + gpuEngine: WebGPUEngine | null; + vramAllocatedBytes: number; + unsubscribeDeviceLost?: () => void; +} diff --git a/packages/webgpu-search/src/gpu/buffers.ts b/packages/webgpu-search/src/gpu/buffers.ts new file mode 100644 index 0000000..5b7d714 --- /dev/null +++ b/packages/webgpu-search/src/gpu/buffers.ts @@ -0,0 +1,276 @@ +/** + * GPU buffer lifecycle for the WebGPU engine (verbatim move of the buffer + * half of the facade): fixed uniform/query allocations, timestamp query + * resources, output/staging (re)allocation, dataset buffer swap with a safe + * swap, append tail writes, and the three teardown paths (re-init = + * unmap-each; explicit destroy = unmap staging pair; device loss = null + * refs, allocated sizes kept). + */ + +import { QUERY_TOKENS_MAX } from '../text-profile'; +import type { PackedDataset } from '../dataset-packing'; +import { BufferUsage, type GpuEngineState } from './state'; + +export class GpuBufferManager { + constructor(private readonly state: GpuEngineState) {} + + /** Uniform 32 B; query QUERY_TOKENS_MAX*4 (min 16 B), destroyed under mutex. */ + createFixedBuffers(): void { + const device = this.state.device!; + + this.state.uniformBuffer = device.createBuffer({ + label: 'Uniform Buffer', + size: 32, + usage: BufferUsage.UNIFORM | BufferUsage.COPY_DST + }); + + this.state.queryBuffer = device.createBuffer({ + label: 'Query Buffer', + size: Math.max(QUERY_TOKENS_MAX * 4, 16), + usage: BufferUsage.STORAGE | BufferUsage.COPY_DST + }); + } + + /** Timestamp query + resolve/staging, degraded to null on alloc failure (warn). */ + createTimestampResources(): void { + const device = this.state.device!; + if (device.features && device.features.has('timestamp-query')) { + try { + this.state.querySet = device.createQuerySet({ + label: 'ts-query', + type: 'timestamp', + count: 2 + }); + this.state.queryResolveBuffer = device.createBuffer({ + label: 'ts-resolve', + size: 16, + usage: BufferUsage.QUERY_RESOLVE | BufferUsage.COPY_SRC + }); + this.state.queryStagingBuffer = device.createBuffer({ + label: 'ts-staging', + size: 16, + usage: BufferUsage.MAP_READ | BufferUsage.COPY_DST + }); + } catch (qErr) { + console.warn('timestamp alloc failed:', qErr); + this.state.querySet = null; + this.state.queryResolveBuffer = null; + this.state.queryStagingBuffer = null; + } + } + } + + ensureCandidateCapacity(capacity: number): void { + const state = this.state; + const safeCap = typeof capacity === 'number' && Number.isFinite(capacity) && capacity > 0 + ? Math.floor(capacity) + : 8192; + const clamped = Math.min(32768, Math.max(8192, safeCap)); + if (clamped !== state.candidateCapacity || !state.outputBuffer) { + state.candidateCapacity = clamped; + if (state.device) { + this.allocateOutputBuffers(clamped); + } + } + } + + allocateOutputBuffers(candidateCapacity: number = 8192): void { + const state = this.state; + const device = state.device; + if (!device) return; + const safeCap = typeof candidateCapacity === 'number' && Number.isFinite(candidateCapacity) && candidateCapacity > 0 + ? Math.floor(candidateCapacity) + : 8192; + const clampedCap = Math.min(32768, Math.max(safeCap, 8192)); + if (state.outputBuffer && clampedCap <= state.candidateCapacity && state.outputByteLength >= 8 + clampedCap * 8) return; + state.candidateCapacity = clampedCap; + state.outputByteLength = 8 + state.candidateCapacity * 8; + + // Advance generation to discard any in-flight reads against old staging buffers + state.generation++; + + // Unmap-before-destroy (symmetric with destroy()): destroying a mapped + // buffer throws / leaves torn state on some implementations. + for (const b of [state.outputBuffer, state.stagingBuffer]) { + try { (b as GPUBuffer | null)?.unmap?.(); } catch {} + } + if (state.outputBuffer) state.outputBuffer.destroy(); + if (state.stagingBuffer) state.stagingBuffer.destroy(); + + state.outputBuffer = device.createBuffer({ + label: 'Output Buffer', + size: state.outputByteLength, + usage: BufferUsage.STORAGE | BufferUsage.COPY_SRC | BufferUsage.COPY_DST + }); + + state.stagingBuffer = device.createBuffer({ + label: 'Staging Buffer', + size: state.outputByteLength, + usage: BufferUsage.MAP_READ | BufferUsage.COPY_DST + }); + } + + /** True when every buffer/pipeline the search pass needs is live. */ + isSearchReady(): boolean { + const s = this.state; + return ( + s.device !== null && + s.recordsBuffer !== null && + s.offsetsBuffer !== null && + s.uniformBuffer !== null && + s.queryBuffer !== null && + s.outputBuffer !== null && + s.stagingBuffer !== null && + s.substringPipeline !== null && + s.fuzzyPipeline !== null + ); + } + + /** Headroom check: dataset bindings must fit `requiredRows` offsets + `requiredTokens` records. */ + hasDatasetBufferCapacity(requiredRows: number, requiredTokens: number): boolean { + const needOffsets = (requiredRows + 1) * 4; + const needRecords = requiredTokens * 4; + const s = this.state; + return ( + s.offsetsBuffer !== null && + s.recordsBuffer !== null && + needOffsets <= s.allocatedOffsetsByteLength && + needRecords <= s.allocatedRecordsByteLength + ); + } + + /** + * Create + upload the packed dataset's storage buffers, then swap them for + * the previous pair (old buffers unmap/destroy only after new allocations + * succeed). Destroys partials and rethrows on failure — the facade restores + * CPU metadata (same split as the original inlined block). + */ + swapDatasetBuffers( + count: number, + packed: PackedDataset, + allocOffsets: number, + allocRecords: number + ): void { + let newOffsetsBuffer: GPUBuffer | null = null; + let newRecordsBuffer: GPUBuffer | null = null; + const device = this.state.device!; + const state = this.state; + + try { + newOffsetsBuffer = device.createBuffer({ + label: `Offsets (${count})`, + size: allocOffsets, + usage: BufferUsage.STORAGE | BufferUsage.COPY_DST + }); + + newRecordsBuffer = device.createBuffer({ + label: `Records (${count})`, + size: allocRecords, + usage: BufferUsage.STORAGE | BufferUsage.COPY_DST + }); + + device.queue.writeBuffer(newOffsetsBuffer, 0, packed.offsetsBufferData); + if (packed.recordsByteLength > 0) { + device.queue.writeBuffer(newRecordsBuffer, 0, packed.recordsBufferData); + } + + // Safe swap: only unmap/destroy previous buffers after new allocations succeed + for (const b of [state.offsetsBuffer, state.recordsBuffer]) { + try { (b as GPUBuffer | null)?.unmap?.(); } catch {} + try { b?.destroy(); } catch {} + } + state.offsetsBuffer = newOffsetsBuffer; + state.recordsBuffer = newRecordsBuffer; + state.allocatedOffsetsByteLength = allocOffsets; + state.allocatedRecordsByteLength = allocRecords; + } catch (allocErr) { + try { newOffsetsBuffer?.destroy(); } catch {} + try { newRecordsBuffer?.destroy(); } catch {} + throw allocErr; + } + } + + /** + * Guards + generation bump + queue writes for appended rows (GPU half of + * appendRows; the facade does the CPU merges). Counts are pre-append. + */ + writeAppendedRows( + prevRows: number, + prevTokens: number, + newTokens: Uint32Array, + newOffsets: Uint32Array + ): void { + const state = this.state; + if (!state.device || !state.offsetsBuffer || !state.recordsBuffer) { + throw new Error('[webgpu-search] Cannot append rows: GPU buffers not initialized.'); + } + const needOffsetsBytes = (prevRows + 1) * 4 + newOffsets.byteLength; + const needRecordsBytes = prevTokens * 4 + newTokens.byteLength; + if (needOffsetsBytes > state.allocatedOffsetsByteLength || needRecordsBytes > state.allocatedRecordsByteLength) { + throw new Error('[webgpu-search] appendRows exceeds allocated buffer headroom.'); + } + state.generation++; + + // Write new offsets: row offsets starting at (prevRows + 1) * 4 + state.device.queue.writeBuffer(state.offsetsBuffer, (prevRows + 1) * 4, newOffsets); + // Write new tokens at prevTokens * 4 + if (newTokens.byteLength > 0) { + state.device.queue.writeBuffer(state.recordsBuffer, prevTokens * 4, newTokens); + } + } + + /** init() re-entry: unmap-each + destroy all buffers, zero allocated sizes. */ + disposeForReinit(): void { + for (const b of [ + this.state.offsetsBuffer, this.state.recordsBuffer, this.state.uniformBuffer, + this.state.queryBuffer, this.state.outputBuffer, this.state.stagingBuffer, + this.state.queryResolveBuffer, this.state.queryStagingBuffer + ]) { + try { (b as GPUBuffer | null)?.unmap?.(); } catch {} + try { (b as GPUBuffer | null)?.destroy(); } catch {} + } + this.state.offsetsBuffer = this.state.recordsBuffer = this.state.uniformBuffer = this.state.queryBuffer = null; + this.state.outputBuffer = this.state.stagingBuffer = this.state.queryResolveBuffer = this.state.queryStagingBuffer = null; + if (this.state.querySet) { + try { this.state.querySet.destroy(); } catch {} + this.state.querySet = null; + } + this.state.allocatedOffsetsByteLength = 0; + this.state.allocatedRecordsByteLength = 0; + } + + /** + * Explicit destroy(): unmap the (possibly mapped) staging pair first, then + * destroy all buffers + querySet and zero allocated sizes. Does not unmap + * dataset/uniform buffers (original did too); pipeline handles + CPU + * metadata parked by the facade. + */ + destroyForTeardown(): void { + try { this.state.stagingBuffer?.unmap(); } catch {} + try { this.state.queryStagingBuffer?.unmap(); } catch {} + const bufs = [ + this.state.offsetsBuffer, this.state.recordsBuffer, this.state.uniformBuffer, + this.state.queryBuffer, this.state.outputBuffer, this.state.stagingBuffer, + this.state.queryResolveBuffer, this.state.queryStagingBuffer + ]; + for (const b of bufs) { try { b?.destroy(); } catch {} } + this.state.offsetsBuffer = this.state.recordsBuffer = this.state.uniformBuffer = this.state.queryBuffer = null; + this.state.outputBuffer = this.state.stagingBuffer = this.state.queryResolveBuffer = this.state.queryStagingBuffer = null; + if (this.state.querySet) { + try { this.state.querySet.destroy(); } catch {} + this.state.querySet = null; + } + this.state.allocatedOffsetsByteLength = 0; + this.state.allocatedRecordsByteLength = 0; + } + + /** + * Device loss: null refs without destroying (device is already lost). + * Keeps allocated sizes (stale `vramAllocatedBytes` is pre-existing). + */ + dropDeviceRefs(): void { + this.state.offsetsBuffer = this.state.recordsBuffer = this.state.uniformBuffer = this.state.queryBuffer = null; + this.state.outputBuffer = this.state.stagingBuffer = this.state.queryResolveBuffer = this.state.queryStagingBuffer = null; + this.state.querySet = null; + } +} diff --git a/packages/webgpu-search/src/gpu/dispatch.ts b/packages/webgpu-search/src/gpu/dispatch.ts new file mode 100644 index 0000000..ac47a0b --- /dev/null +++ b/packages/webgpu-search/src/gpu/dispatch.ts @@ -0,0 +1,226 @@ +/** + * Search submission + readback for the WebGPU engine (verbatim move of the + * facade's dispatch/readback helpers): queue uploads, bind group creation, + * single-pass vs chunked dispatch with counter-reset fallbacks, and the + * generation-epoch-guarded mapAsync readbacks. Sync blocks use a device + * snapshot (no awaits inside); readback paths keep live epoch reads. + */ + +import { abortError } from '../guard'; +import type { SearchResultItem } from '../types'; +import { MapMode, type GpuEngineState } from './state'; + +/** Compute workgroup size the WGSL kernels assume (dispatch math: ceil(rows/128)). */ +export const WORKGROUP_SIZE = 128; + +/** Decoded staging buffer contents for one search. */ +export interface MappedSearchOutput { + totalMatches: number; + candidateCount: number; + hasOverflow: boolean; + candidates: SearchResultItem[]; +} + +export class GpuDispatcher { + constructor(private readonly state: GpuEngineState) {} + + /** Queue uploads for one search: query tokens, then the 32 B uniform header. */ + uploadSearchRequest(tokens: Uint32Array, uniform: ArrayBuffer): void { + const device = this.state.device!; + device.queue.writeBuffer(this.state.queryBuffer!, 0, tokens); + device.queue.writeBuffer(this.state.uniformBuffer!, 0, uniform); + } + + /** Bind group over the fixed search bindings (auto-layout pipeline). */ + createSearchBindGroup(pipeline: GPUComputePipeline): GPUBindGroup { + const s = this.state; + return s.device!.createBindGroup({ + label: 'search-bg', + layout: pipeline.getBindGroupLayout(0), + entries: [ + { binding: 0, resource: { buffer: s.uniformBuffer! } }, + { binding: 1, resource: { buffer: s.offsetsBuffer! } }, + { binding: 2, resource: { buffer: s.recordsBuffer! } }, + { binding: 3, resource: { buffer: s.queryBuffer! } }, + { binding: 4, resource: { buffer: s.outputBuffer! } } + ] + }); + } + + /** + * Zero the atomic counter fallback for Safari without `clearBuffer`. + * 8-byte zero write (count + _pad0), queue-ordered before dispatch. + * Never skip the reset (stale count corruption). No scoring change. + */ + private writeZeroCounterFallback(): void { + this.state.device!.queue.writeBuffer(this.state.outputBuffer!, 0, new Uint8Array(8)); + } + + /** + * Single-pass dispatch: counter reset + compute pass + timestamp resolve + + * readback copy in one command buffer. No scoring change. + */ + dispatchSinglePass( + pipeline: GPUComputePipeline, + bg: GPUBindGroup, + totalWorkgroups: number + ): void { + const state = this.state; + const device = state.device!; + const passDesc: GPUComputePassDescriptor = { + label: 'search-pass' + }; + if (state.querySet) { + (passDesc as any).timestampWrites = { + querySet: state.querySet, + beginningOfPassWriteIndex: 0, + endOfPassWriteIndex: 1 + }; + } + const enc = device.createCommandEncoder({ label: 'search' }); + + if (typeof (enc as any).clearBuffer === 'function') { + enc.clearBuffer(state.outputBuffer!, 0, state.outputByteLength); + } else { + // No clearBuffer: encode the zero via copy path — submit the + // queue write before the compute submit (serial ordering). + this.writeZeroCounterFallback(); + } + + const pass = enc.beginComputePass(passDesc); + pass.setPipeline(pipeline); + pass.setBindGroup(0, bg); + pass.dispatchWorkgroups(totalWorkgroups); + pass.end(); + + if (state.querySet && state.queryResolveBuffer && state.queryStagingBuffer) { + enc.resolveQuerySet(state.querySet, 0, 2, state.queryResolveBuffer, 0); + enc.copyBufferToBuffer(state.queryResolveBuffer, 0, state.queryStagingBuffer, 0, 16); + } + + enc.copyBufferToBuffer(state.outputBuffer!, 0, state.stagingBuffer!, 0, state.outputByteLength); + device.queue.submit([enc.finish()]); + } + + /** + * Chunked dispatch: when ceil(rows/128) exceeds the device dimension + * limit, each chunk re-writes the uniform (_pad.x = base row) and submits + * its own pass; the shared atomic counter accumulates. WebGPU's serial + * queue guarantees each dispatch samples only prior uniform writes. + * Timestamps skipped (null). No scoring change. + */ + dispatchChunked( + pipeline: GPUComputePipeline, + bg: GPUBindGroup, + uniform: ArrayBuffer, + U32: Uint32Array, + maxDim: number, + totalRows: number + ): void { + const state = this.state; + const device = state.device!; + const step = maxDim * WORKGROUP_SIZE; + const clr = device.createCommandEncoder({ label: 'clear' }); + if (typeof (clr as any).clearBuffer === 'function') { + (clr as any).clearBuffer(state.outputBuffer!, 0, state.outputByteLength); + device.queue.submit([clr.finish()]); + } else { + this.writeZeroCounterFallback(); + } + let base = 0; + let left = totalRows; + while (left > 0) { + const n = left > step ? step : left; + const wgs = Math.ceil(n / WORKGROUP_SIZE); + U32[4] = base; + device.queue.writeBuffer(state.uniformBuffer!, 0, uniform); + const chunkEnc = device.createCommandEncoder({ label: 'chunk' }); + const chunkPass = chunkEnc.beginComputePass({ label: 'chunk-pass' }); + chunkPass.setPipeline(pipeline); + chunkPass.setBindGroup(0, bg); + chunkPass.dispatchWorkgroups(wgs); + chunkPass.end(); + device.queue.submit([chunkEnc.finish()]); + base += n; + left -= n; + } + const cp = device.createCommandEncoder({ label: 'copy' }); + cp.copyBufferToBuffer(state.outputBuffer!, 0, state.stagingBuffer!, 0, state.outputByteLength); + device.queue.submit([cp.finish()]); + } + + /** + * Map the timestamp staging buffer and extract the GPU execution delta. + * Epoch-only guard: a teardown during the await discards the stale sample + * (null); warn-suppressed once the epoch moved too. + */ + async readTimestampGpuMs(tsStaging: GPUBuffer, epoch: number): Promise { + let gpuExecutionMs: number | null = null; + try { + await tsStaging.mapAsync(MapMode.READ); + if (epoch === this.state.generation) { + const tu = new BigUint64Array(tsStaging.getMappedRange()); + if (tu[1] >= (tu[0]) && (tu[0]) > 0n) { + gpuExecutionMs = Number((tu[1]) - (tu[0])) / 1_000_000; + } + } + } catch (tsErr) { + if (epoch === this.state.generation) console.warn('timestamp read failed:', tsErr); + } finally { + try { tsStaging.unmap(); } catch {} + } + return gpuExecutionMs; + } + + /** + * Map the output staging buffer. An epoch that moved during the await + * (destroy/reallocation raced the map) surfaces as AbortError — never a + * fallback-eligible GPU failure; genuine map failures rethrow untouched. + */ + async mapStagedOutput(stagingBuffer: GPUBuffer, epoch: number): Promise { + try { + await stagingBuffer.mapAsync(MapMode.READ); + } catch (mapErr) { + if (epoch !== this.state.generation) { + try { stagingBuffer.unmap(); } catch {} + throw abortError(); + } + throw mapErr; + } + } + + /** Re-check epoch + device + staging identity after an await; drift → AbortError. */ + assertStagingLive(stagingBuffer: GPUBuffer, epoch: number): void { + const s = this.state; + if (epoch !== s.generation || !s.device || s.stagingBuffer !== stagingBuffer) { + throw abortError(); + } + } + + /** + * Decode one mapped staging buffer: u32[0] = total matches, then + * per-candidate pairs (u32 index, i32 score). `resolveText` maps rows to + * display text (token-only datasets resolve `''`; see DatasetLike). + */ + readMappedOutput( + stagingBuffer: GPUBuffer, + candidateCapacity: number, + resolveText: (index: number) => string + ): MappedSearchOutput { + const arrayBuffer = stagingBuffer.getMappedRange(); + const u32Read = new Uint32Array(arrayBuffer); + const i32Read = new Int32Array(arrayBuffer); + + const totalMatches = u32Read[0]; + const candidateCount = Math.min(totalMatches, candidateCapacity); + const hasOverflow = totalMatches > candidateCapacity; + + const candidates: SearchResultItem[] = new Array(candidateCount); + for (let i = 0; i < candidateCount; i++) { + const index = u32Read[2 + i * 2]; + const score = i32Read[3 + i * 2]; + candidates[i] = { index, score, text: resolveText(index) }; + } + return { totalMatches, candidateCount, hasOverflow, candidates }; + } +} diff --git a/packages/webgpu-search/src/gpu/pipelines.ts b/packages/webgpu-search/src/gpu/pipelines.ts new file mode 100644 index 0000000..037aa0d --- /dev/null +++ b/packages/webgpu-search/src/gpu/pipelines.ts @@ -0,0 +1,62 @@ +/** + * Shader module + compute pipeline setup for the WebGPU engine (verbatim + * move of the pipeline half of the facade's `setupPipelinesAndBuffers`). + * Device reads stay live (`state.device`): a device lost mid-setup still + * throws (fail-closed teardown), never completes init on a stale device. + */ + +import SUBSTRING_WGSL from '../shaders/substring.wgsl'; +import FUZZY_WGSL from '../shaders/fuzzy.wgsl'; +import type { GpuEngineState } from './state'; + +export class GpuPipelineManager { + constructor(private readonly state: GpuEngineState) {} + + /** + * Create substring/fuzzy modules and compute pipelines. Verbatim call + * order (module → module → pipeline → pipeline): mock/spy compat. + * `createComputePipelineAsync` when available, sync fallback otherwise. + */ + async setup(): Promise { + const substringModule = this.state.device!.createShaderModule({ + label: 'substring', + code: SUBSTRING_WGSL + }); + + const fuzzyModule = this.state.device!.createShaderModule({ + label: 'fuzzy', + code: FUZZY_WGSL + }); + + const createPipeline = async (desc: GPUComputePipelineDescriptor) => { + if (typeof (this.state.device as any).createComputePipelineAsync === 'function') { + return await this.state.device!.createComputePipelineAsync(desc); + } + return (this.state.device as any).createComputePipeline(desc); + }; + + this.state.substringPipeline = await createPipeline({ + label: 'substring', + layout: 'auto', + compute: { + module: substringModule, + entryPoint: 'main' + } + }); + + this.state.fuzzyPipeline = await createPipeline({ + label: 'fuzzy', + layout: 'auto', + compute: { + module: fuzzyModule, + entryPoint: 'main' + } + }); + } + + /** Null the pipeline handles (destroy / device loss / re-init). */ + drop(): void { + this.state.substringPipeline = null; + this.state.fuzzyPipeline = null; + } +} diff --git a/packages/webgpu-search/src/gpu/state.ts b/packages/webgpu-search/src/gpu/state.ts new file mode 100644 index 0000000..a33515b --- /dev/null +++ b/packages/webgpu-search/src/gpu/state.ts @@ -0,0 +1,74 @@ +/** + * Shared GPU state + portable GPU constant tables for the extracted engine + * subsystems (`gpu/buffers.ts`, `gpu/pipelines.ts`, `gpu/dispatch.ts`). + * One `GpuEngineState` per `WebGPUEngine`; managers mutate it by reference. + * Dataset CPU metadata stays on the facade. Portable: no DOM refs. + */ + +/** GPUBufferUsage flags: global table under WebGPU runtimes, constants under Node/Bun. */ +export const BufferUsage = (typeof globalThis !== 'undefined' && 'GPUBufferUsage' in globalThis ? (globalThis as any).GPUBufferUsage : { + MAP_READ: 0x0001, + MAP_WRITE: 0x0002, + COPY_SRC: 0x0004, + COPY_DST: 0x0008, + INDEX: 0x0010, + VERTEX: 0x0020, + UNIFORM: 0x0040, + STORAGE: 0x0080, + INDIRECT: 0x0100, + QUERY_RESOLVE: 0x0200 +}); + +/** GPUMapMode flags: same portability contract as BufferUsage. */ +export const MapMode = (typeof globalThis !== 'undefined' && 'GPUMapMode' in globalThis ? (globalThis as any).GPUMapMode : { + READ: 0x0001, + WRITE: 0x0002 +}); + +/** All mutable GPU-side engine state (buffers/pipelines/dispatch/dispose + epoch). */ +export interface GpuEngineState { + device: GPUDevice | null; // null until init / after device loss / teardown + substringPipeline: GPUComputePipeline | null; + fuzzyPipeline: GPUComputePipeline | null; + offsetsBuffer: GPUBuffer | null; + recordsBuffer: GPUBuffer | null; + uniformBuffer: GPUBuffer | null; // 32 B uniform header, written per query + queryBuffer: GPUBuffer | null; // QUERY_TOKENS_MAX*4 (min 16 B), written per query + outputBuffer: GPUBuffer | null; // atomic counter + candidates: 8 + cap*8 bytes + stagingBuffer: GPUBuffer | null; // MAP_READ twin of outputBuffer, remapped per search + querySet: GPUQuerySet | null; // timestamp resources (null when unsupported/failure) + queryResolveBuffer: GPUBuffer | null; + queryStagingBuffer: GPUBuffer | null; + candidateCapacity: number; // output/staging sizing, clamped 8192..32768 + outputByteLength: number; // 8 + candidateCapacity * 8 + /** Epoch bumped on every teardown/realloc/load/append; in-flight readbacks discard. */ + generation: number; + /** + * Reserved dataset binding sizes. NOT zeroed on device loss (stale + * `vramAllocatedBytes` is pre-existing); explicit destroy()/dispose zero them. + */ + allocatedOffsetsByteLength: number; + allocatedRecordsByteLength: number; +} + +export function createGpuEngineState(): GpuEngineState { + return { + device: null, + substringPipeline: null, + fuzzyPipeline: null, + offsetsBuffer: null, + recordsBuffer: null, + uniformBuffer: null, + queryBuffer: null, + outputBuffer: null, + stagingBuffer: null, + querySet: null, + queryResolveBuffer: null, + queryStagingBuffer: null, + candidateCapacity: 8192, + outputByteLength: 0, + generation: 0, + allocatedOffsetsByteLength: 0, + allocatedRecordsByteLength: 0, + }; +} diff --git a/packages/webgpu-search/src/webgpu-engine.ts b/packages/webgpu-search/src/webgpu-engine.ts index fbf9565..bf90f78 100644 --- a/packages/webgpu-search/src/webgpu-engine.ts +++ b/packages/webgpu-search/src/webgpu-engine.ts @@ -1,5 +1,3 @@ -import SUBSTRING_WGSL from './shaders/substring.wgsl'; -import FUZZY_WGSL from './shaders/fuzzy.wgsl'; import { GpuDevicePool, assertValidPowerPreference } from './gpu-device-pool'; import { checkMemoryBudget, @@ -12,7 +10,7 @@ import { } from './dataset-packing'; import { normalizeText } from './text-normalization'; import { compareExactResults } from './exact-scorer'; -import { clampLimit, nowMs, throwIfAborted, abortError, assertValidMode } from './guard'; +import { clampLimit, nowMs, throwIfAborted, assertValidMode } from './guard'; import { PROFILE_TO_ENUM, QUERY_TOKENS_MAX, @@ -28,24 +26,10 @@ import { import { normalizeTypoTolerance } from './search/typo-tolerance'; import { WebGPUSearchError } from './errors'; import type { AdapterInfo, SearchOptions, SearchResultItem, SearchTimings, SearchMode } from './types'; - -const BufferUsage = (typeof globalThis !== 'undefined' && 'GPUBufferUsage' in globalThis ? (globalThis as any).GPUBufferUsage : { - MAP_READ: 0x0001, - MAP_WRITE: 0x0002, - COPY_SRC: 0x0004, - COPY_DST: 0x0008, - INDEX: 0x0010, - VERTEX: 0x0020, - UNIFORM: 0x0040, - STORAGE: 0x0080, - INDIRECT: 0x0100, - QUERY_RESOLVE: 0x0200 -}); - -const MapMode = (typeof globalThis !== 'undefined' && 'GPUMapMode' in globalThis ? (globalThis as any).GPUMapMode : { - READ: 0x0001, - WRITE: 0x0002 -}); +import { createGpuEngineState } from './gpu/state'; +import { GpuBufferManager } from './gpu/buffers'; +import { GpuPipelineManager } from './gpu/pipelines'; +import { GpuDispatcher, WORKGROUP_SIZE } from './gpu/dispatch'; export interface DatasetLike { // `strings?` optional: search-index enriches `text` from its own `items`, so @@ -88,24 +72,17 @@ export interface ColdSearchResult extends WebGPUSearchResult { coldTotalMs: number; } +/** Search-engine facade over the gpu/ managers: device pool, CPU dataset metadata, mutex, fail-closed gates. */ export class WebGPUEngine { - private device: GPUDevice | null = null; + private readonly state = createGpuEngineState(); + private readonly pipelines = new GpuPipelineManager(this.state); + private readonly buffers = new GpuBufferManager(this.state); + private readonly dispatch = new GpuDispatcher(this.state); + private isSharedDevice: boolean = false; private deviceReleased: boolean = false; - private substringPipeline: GPUComputePipeline | null = null; - private fuzzyPipeline: GPUComputePipeline | null = null; - - private offsetsBuffer: GPUBuffer | null = null; - private recordsBuffer: GPUBuffer | null = null; - private uniformBuffer: GPUBuffer | null = null; - private queryBuffer: GPUBuffer | null = null; - private outputBuffer: GPUBuffer | null = null; - private stagingBuffer: GPUBuffer | null = null; - - private querySet: GPUQuerySet | null = null; - private queryResolveBuffer: GPUBuffer | null = null; - private queryStagingBuffer: GPUBuffer | null = null; + // Dataset CPU metadata (engine-owned; buffer writes flow through managers). private currentDatasetSize: number = 0; private currentStrings: string[] | null = null; private currentTokens: Uint32Array | null = null; @@ -114,45 +91,32 @@ export class WebGPUEngine { private profileId: string = 'unicode-default'; private unicodeVersion: string = UNICODE_VERSION; private scoringVersion: string = SCORING_VERSION; - private candidateCapacity: number = 8192; - private outputByteLength: number = 0; private searchMutex: Promise = Promise.resolve(); - private generation: number = 0; public adapterInfo: AdapterInfo | null = null; get isReady(): boolean { - return this.device !== null; + return this.state.device !== null; } get currentSize(): number { return this.currentDatasetSize; } - private allocatedOffsetsByteLength: number = 0; - private allocatedRecordsByteLength: number = 0; - get allocatedOffsetsBytes(): number { - return this.allocatedOffsetsByteLength; + return this.state.allocatedOffsetsByteLength; } get allocatedRecordsBytes(): number { - return this.allocatedRecordsByteLength; + return this.state.allocatedRecordsByteLength; } get vramAllocatedBytes(): number { - return this.allocatedOffsetsByteLength + this.allocatedRecordsByteLength; + return this.state.allocatedOffsetsByteLength + this.state.allocatedRecordsByteLength; } canFitHeadroom(requiredRows: number, requiredTokens: number): boolean { - const needOffsets = (requiredRows + 1) * 4; - const needRecords = requiredTokens * 4; - return ( - this.offsetsBuffer !== null && - this.recordsBuffer !== null && - needOffsets <= this.allocatedOffsetsByteLength && - needRecords <= this.allocatedRecordsByteLength - ); + return this.buffers.hasDatasetBufferCapacity(requiredRows, requiredTokens); } isHeadroomExhausted(requiredRows: number, requiredTokens: number): boolean { @@ -165,9 +129,10 @@ export class WebGPUEngine { // Guard re-entry: dispose existing GPU buffers AND release the previous // device reference before re-acquiring, so init() twice leaks neither the // first buffer set nor a pool refcount (benchmark/power users). - if (this.device) { - this.generation++; - this.disposeGpuBuffers(); + if (this.state.device) { + this.state.generation++; + this.buffers.disposeForReinit(); + this.pipelines.drop(); this.releaseCurrentDevice(); } // Normalize overloads: bare GPUDevice (legacy) or options bag. @@ -191,7 +156,7 @@ export class WebGPUEngine { } if (injected) { assertValidPowerPreference(powerPreference); - this.device = injected; + this.state.device = injected; this.isSharedDevice = false; this.deviceReleased = false; this.adapterInfo = { @@ -217,7 +182,7 @@ export class WebGPUEngine { return false; } - this.device = acquired.device; + this.state.device = acquired.device; this.adapterInfo = acquired.adapterInfo; this.isSharedDevice = acquired.isShared; this.deviceReleased = false; @@ -241,130 +206,42 @@ export class WebGPUEngine { /** Release the held device reference exactly once (deviceReleased guard). */ private releaseCurrentDevice(): void { - if (this.device && !this.deviceReleased) { + if (this.state.device && !this.deviceReleased) { this.deviceReleased = true; try { - GpuDevicePool.releaseDevice(this.device, this.isSharedDevice); + GpuDevicePool.releaseDevice(this.state.device, this.isSharedDevice); } catch {} } - this.device = null; + this.state.device = null; } private attachLostHandler(): void { - const dev = this.device as unknown as { lost?: Promise } | null; + const dev = this.state.device as unknown as { lost?: Promise } | null; if (!dev?.lost?.then) return; // Capture identity: re-init() attaches a new handler per device; a stale // handler for a previous device must not tear down the live engine. - const mine = this.device; + const mine = this.state.device; dev.lost.then(() => { - if (this.device !== mine) return; + if (this.state.device !== mine) return; // Device loss: null handles, mark unready. Shared refCount decremented // once (deviceReleased guard) so live siblings keep the device. - this.substringPipeline = this.fuzzyPipeline = null; - this.offsetsBuffer = this.recordsBuffer = this.uniformBuffer = this.queryBuffer = null; - this.outputBuffer = this.stagingBuffer = this.queryResolveBuffer = this.queryStagingBuffer = null; - this.querySet = null; + this.pipelines.drop(); + this.buffers.dropDeviceRefs(); this.releaseCurrentDevice(); - this.generation++; + this.state.generation++; }, () => {}); } - private disposeGpuBuffers(): void { - for (const b of [this.offsetsBuffer, this.recordsBuffer, this.uniformBuffer, this.queryBuffer, this.outputBuffer, this.stagingBuffer, this.queryResolveBuffer, this.queryStagingBuffer]) { - try { (b as GPUBuffer | null)?.unmap?.(); } catch {} - try { (b as GPUBuffer | null)?.destroy(); } catch {} - } - this.offsetsBuffer = this.recordsBuffer = this.uniformBuffer = this.queryBuffer = null; - this.outputBuffer = this.stagingBuffer = this.queryResolveBuffer = this.queryStagingBuffer = null; - if (this.querySet) { - try { this.querySet.destroy(); } catch {} - this.querySet = null; - } - this.substringPipeline = this.fuzzyPipeline = null; - this.allocatedOffsetsByteLength = 0; - this.allocatedRecordsByteLength = 0; - } - private async setupPipelinesAndBuffers(): Promise { - if (!this.device) return false; - - const substringModule = this.device.createShaderModule({ - label: 'substring', - code: SUBSTRING_WGSL - }); - - const fuzzyModule = this.device.createShaderModule({ - label: 'fuzzy', - code: FUZZY_WGSL - }); - - const createPipeline = async (desc: GPUComputePipelineDescriptor) => { - if (typeof (this.device as any).createComputePipelineAsync === 'function') { - return await this.device!.createComputePipelineAsync(desc); - } - return (this.device as any).createComputePipeline(desc); - }; - - this.substringPipeline = await createPipeline({ - label: 'substring', - layout: 'auto', - compute: { - module: substringModule, - entryPoint: 'main' - } - }); - - this.fuzzyPipeline = await createPipeline({ - label: 'fuzzy', - layout: 'auto', - compute: { - module: fuzzyModule, - entryPoint: 'main' - } - }); + if (!this.state.device) return false; - // Uniform header: 32 B, 16 B-aligned. Query buffer: 512 B persistent - // storage (QUERY_TOKENS_MAX*4, min 16 B), written per query. - this.uniformBuffer = this.device.createBuffer({ - label: 'Uniform Buffer', - size: 32, - usage: BufferUsage.UNIFORM | BufferUsage.COPY_DST - }); + // Verbatim original order: substring module → fuzzy module → substring pipeline → fuzzy pipeline → uniform → query → timestamps → output buffers (mock/spy compatibility). + await this.pipelines.setup(); - // Persistent storage query buffer: QUERY_TOKENS_MAX * 4 = 512 B, min 16 B, - // written per query via queue.writeBuffer, destroyed under mutex. - this.queryBuffer = this.device.createBuffer({ - label: 'Query Buffer', - size: Math.max(QUERY_TOKENS_MAX * 4, 16), - usage: BufferUsage.STORAGE | BufferUsage.COPY_DST - }); + this.buffers.createFixedBuffers(); + this.buffers.createTimestampResources(); - if (this.device.features && this.device.features.has('timestamp-query')) { - try { - this.querySet = this.device.createQuerySet({ - label: 'ts-query', - type: 'timestamp', - count: 2 - }); - this.queryResolveBuffer = this.device.createBuffer({ - label: 'ts-resolve', - size: 16, - usage: BufferUsage.QUERY_RESOLVE | BufferUsage.COPY_SRC - }); - this.queryStagingBuffer = this.device.createBuffer({ - label: 'ts-staging', - size: 16, - usage: BufferUsage.MAP_READ | BufferUsage.COPY_DST - }); - } catch (qErr) { - console.warn('timestamp alloc failed:', qErr); - this.querySet = null; - this.queryResolveBuffer = null; - this.queryStagingBuffer = null; - } - } - - this.allocateOutputBuffers(this.candidateCapacity); + this.buffers.allocateOutputBuffers(this.state.candidateCapacity); return true; } @@ -372,7 +249,7 @@ export class WebGPUEngine { * Current candidate capacity allocated for the output buffer. */ get currentCandidateCapacity(): number { - return this.candidateCapacity; + return this.state.candidateCapacity; } /** @@ -380,50 +257,11 @@ export class WebGPUEngine { * (clamped between 8,192 and 32,768). */ ensureCandidateCapacity(capacity: number): void { - const safeCap = typeof capacity === 'number' && Number.isFinite(capacity) && capacity > 0 - ? Math.floor(capacity) - : 8192; - const clamped = Math.min(32768, Math.max(8192, safeCap)); - if (clamped !== this.candidateCapacity || !this.outputBuffer) { - this.candidateCapacity = clamped; - if (this.device) { - this.allocateOutputBuffers(clamped); - } - } + this.buffers.ensureCandidateCapacity(capacity); } allocateOutputBuffers(candidateCapacity: number = 8192): void { - if (!this.device) return; - const safeCap = typeof candidateCapacity === 'number' && Number.isFinite(candidateCapacity) && candidateCapacity > 0 - ? Math.floor(candidateCapacity) - : 8192; - const clampedCap = Math.min(32768, Math.max(safeCap, 8192)); - if (this.outputBuffer && clampedCap <= this.candidateCapacity && this.outputByteLength >= 8 + clampedCap * 8) return; - this.candidateCapacity = clampedCap; - this.outputByteLength = 8 + this.candidateCapacity * 8; - - // Advance generation to discard any in-flight reads against old staging buffers - this.generation++; - - // Unmap-before-destroy (symmetric with destroy()): destroying a mapped - // buffer throws / leaves torn state on some implementations. - for (const b of [this.outputBuffer, this.stagingBuffer]) { - try { (b as GPUBuffer | null)?.unmap?.(); } catch {} - } - if (this.outputBuffer) this.outputBuffer.destroy(); - if (this.stagingBuffer) this.stagingBuffer.destroy(); - - this.outputBuffer = this.device.createBuffer({ - label: 'Output Buffer', - size: this.outputByteLength, - usage: BufferUsage.STORAGE | BufferUsage.COPY_SRC | BufferUsage.COPY_DST - }); - - this.stagingBuffer = this.device.createBuffer({ - label: 'Staging Buffer', - size: this.outputByteLength, - usage: BufferUsage.MAP_READ | BufferUsage.COPY_DST - }); + this.buffers.allocateOutputBuffers(candidateCapacity); } private queued(fn: () => Promise): Promise { @@ -509,7 +347,7 @@ export class WebGPUEngine { dataset: EngineDataset, options?: { rowCapacity?: number; tokenCapacity?: number; growthFactor?: number } ): Promise<{ uploadTimeMs: number }> { - this.generation++; + this.state.generation++; const { packed, strings } = this.resolvePacked(dataset); const count = packed.rowCount; @@ -534,7 +372,7 @@ export class WebGPUEngine { this.unicodeVersion = packed.unicodeVersion; this.scoringVersion = packed.scoringVersion; - if (!this.device) { + if (!this.state.device) { Object.assign(this, { currentDatasetSize: prev.size, currentStrings: prev.strings, @@ -554,7 +392,7 @@ export class WebGPUEngine { // to check would churn GPU contexts (rejected: acquire-always), while // allocation-time-only would OOM mid-load leaving half-state (rejected). const avgBytes = count === 0 ? 0 : packed.recordsByteLength / count; - const budget = checkMemoryBudget(count, avgBytes, this.device, this.candidateCapacity); + const budget = checkMemoryBudget(count, avgBytes, this.state.device, this.state.candidateCapacity); if (!budget.allowed) { Object.assign(this, { currentDatasetSize: prev.size, @@ -574,11 +412,11 @@ export class WebGPUEngine { const targetOffsets = computeClampedHeadroomBytes( Math.max((options?.rowCapacity !== undefined ? options.rowCapacity + 1 : count + 1) * 4, reqOffsets), - { growthFactor: options?.growthFactor, device: this.device } + { growthFactor: options?.growthFactor, device: this.state.device } ); const targetRecords = computeClampedHeadroomBytes( Math.max((options?.tokenCapacity !== undefined ? options.tokenCapacity : packed.tokenCount) * 4, reqRecords), - { growthFactor: options?.growthFactor, device: this.device } + { growthFactor: options?.growthFactor, device: this.state.device } ); const allocOffsets = Math.max(reqOffsets, targetOffsets); @@ -588,7 +426,7 @@ export class WebGPUEngine { // actual corpus above the limit can land here. Throw the budget error // (callers fall back to CPU) instead of creating an oversized binding // that fails validation and silently yields 0 results. - const bindingLimit = resolveStorageBufferLimit(this.device); + const bindingLimit = resolveStorageBufferLimit(this.state.device); if (allocOffsets > bindingLimit || allocRecords > bindingLimit) { Object.assign(this, { currentDatasetSize: prev.size, @@ -605,39 +443,9 @@ export class WebGPUEngine { } const t0 = nowMs(); - let newOffsetsBuffer: GPUBuffer | null = null; - let newRecordsBuffer: GPUBuffer | null = null; - try { - newOffsetsBuffer = this.device.createBuffer({ - label: `Offsets (${count})`, - size: allocOffsets, - usage: BufferUsage.STORAGE | BufferUsage.COPY_DST - }); - - newRecordsBuffer = this.device.createBuffer({ - label: `Records (${count})`, - size: allocRecords, - usage: BufferUsage.STORAGE | BufferUsage.COPY_DST - }); - - this.device.queue.writeBuffer(newOffsetsBuffer, 0, packed.offsetsBufferData); - if (packed.recordsByteLength > 0) { - this.device.queue.writeBuffer(newRecordsBuffer, 0, packed.recordsBufferData); - } - - // Safe swap: only unmap/destroy previous buffers after new allocations succeed - for (const b of [this.offsetsBuffer, this.recordsBuffer]) { - try { (b as GPUBuffer | null)?.unmap?.(); } catch {} - try { b?.destroy(); } catch {} - } - this.offsetsBuffer = newOffsetsBuffer; - this.recordsBuffer = newRecordsBuffer; - this.allocatedOffsetsByteLength = allocOffsets; - this.allocatedRecordsByteLength = allocRecords; + this.buffers.swapDatasetBuffers(count, packed, allocOffsets, allocRecords); } catch (allocErr) { - try { newOffsetsBuffer?.destroy(); } catch {} - try { newRecordsBuffer?.destroy(); } catch {} // Restore previous CPU metadata so stats don't lie after OOM. Object.assign(this, { currentDatasetSize: prev.size, @@ -663,24 +471,9 @@ export class WebGPUEngine { newTotalRows: number ): Promise { return this.queued(async () => { - if (!this.device || !this.offsetsBuffer || !this.recordsBuffer) { - throw new Error('[webgpu-search] Cannot append rows: GPU buffers not initialized.'); - } const prevRows = this.currentDatasetSize; const prevTokens = this.currentTokens ? this.currentTokens.length : 0; - const needOffsetsBytes = (prevRows + 1) * 4 + newOffsets.byteLength; - const needRecordsBytes = prevTokens * 4 + newTokens.byteLength; - if (needOffsetsBytes > this.allocatedOffsetsByteLength || needRecordsBytes > this.allocatedRecordsByteLength) { - throw new Error('[webgpu-search] appendRows exceeds allocated buffer headroom.'); - } - this.generation++; - - // Write new offsets: row offsets starting at (prevRows + 1) * 4 - this.device.queue.writeBuffer(this.offsetsBuffer, (prevRows + 1) * 4, newOffsets); - // Write new tokens at prevTokens * 4 - if (newTokens.byteLength > 0) { - this.device.queue.writeBuffer(this.recordsBuffer, prevTokens * 4, newTokens); - } + this.buffers.writeAppendedRows(prevRows, prevTokens, newTokens, newOffsets); const combinedTokens = new Uint32Array(prevTokens + newTokens.length); if (this.currentTokens) combinedTokens.set(this.currentTokens, 0); @@ -700,107 +493,6 @@ export class WebGPUEngine { return this.queued(() => this.searchInternal(query, options)); } - /** - * Zero the atomic counter fallback for Safari without `clearBuffer`. - * 8-byte zero write (count + _pad0), queue-ordered before dispatch. - * Never skip the reset (stale count corruption). No scoring change. - */ - private writeZeroCounterFallback(): void { - this.device!.queue.writeBuffer(this.outputBuffer!, 0, new Uint8Array(8)); - } - - /** - * Single-pass dispatch (extracted). Fast path — keeps timestamps. - * Encodes counter reset + compute pass + timestamp resolve + readback - * copy in one command buffer. No scoring change. - */ - private dispatchSinglePass( - pipeline: GPUComputePipeline, - bg: GPUBindGroup, - totalWorkgroups: number - ): void { - const passDesc: GPUComputePassDescriptor = { - label: 'search-pass' - }; - if (this.querySet) { - (passDesc as any).timestampWrites = { - querySet: this.querySet, - beginningOfPassWriteIndex: 0, - endOfPassWriteIndex: 1 - }; - } - const enc = this.device!.createCommandEncoder({ label: 'search' }); - - if (typeof (enc as any).clearBuffer === 'function') { - enc.clearBuffer(this.outputBuffer!, 0, this.outputByteLength); - } else { - // No clearBuffer: encode the zero via copy path — submit the - // queue write before the compute submit (serial ordering). - this.writeZeroCounterFallback(); - } - - const pass = enc.beginComputePass(passDesc); - pass.setPipeline(pipeline); - pass.setBindGroup(0, bg); - pass.dispatchWorkgroups(totalWorkgroups); - pass.end(); - - if (this.querySet && this.queryResolveBuffer && this.queryStagingBuffer) { - enc.resolveQuerySet(this.querySet, 0, 2, this.queryResolveBuffer, 0); - enc.copyBufferToBuffer(this.queryResolveBuffer, 0, this.queryStagingBuffer, 0, 16); - } - - enc.copyBufferToBuffer(this.outputBuffer!, 0, this.stagingBuffer!, 0, this.outputByteLength); - this.device!.queue.submit([enc.finish()]); - } - - /** - * Chunked dispatch (extracted). When ceil(rows/128) exceeds the device - * dimension limit, each chunk re-writes the uniform (_pad.x = base row) - * and submits its own pass; the shared atomic counter accumulates. - * Queue-ordering dependency: WebGPU guarantees serial queue execution, - * so each dispatch samples only prior uniform writes (chunk counts are - * tiny in practice: step ≈ 8.4M rows/chunk). Timestamps skipped (null). - * No scoring change. - */ - private dispatchChunked( - pipeline: GPUComputePipeline, - bg: GPUBindGroup, - udata: ArrayBuffer, - U32: Uint32Array, - maxDim: number, - workgroupSize: number - ): void { - const step = maxDim * workgroupSize; - const clr = this.device!.createCommandEncoder({ label: 'clear' }); - if (typeof (clr as any).clearBuffer === 'function') { - (clr as any).clearBuffer(this.outputBuffer!, 0, this.outputByteLength); - this.device!.queue.submit([clr.finish()]); - } else { - this.writeZeroCounterFallback(); - } - let base = 0; - let left = this.currentDatasetSize; - while (left > 0) { - const n = left > step ? step : left; - const wgs = Math.ceil(n / workgroupSize); - U32[4] = base; - this.device!.queue.writeBuffer(this.uniformBuffer!, 0, udata); - const chunkEnc = this.device!.createCommandEncoder({ label: 'chunk' }); - const chunkPass = chunkEnc.beginComputePass({ label: 'chunk-pass' }); - chunkPass.setPipeline(pipeline); - chunkPass.setBindGroup(0, bg); - chunkPass.dispatchWorkgroups(wgs); - chunkPass.end(); - this.device!.queue.submit([chunkEnc.finish()]); - base += n; - left -= n; - } - const cp = this.device!.createCommandEncoder({ label: 'copy' }); - cp.copyBufferToBuffer(this.outputBuffer!, 0, this.stagingBuffer!, 0, this.outputByteLength); - this.device!.queue.submit([cp.finish()]); - } - private async searchInternal(query: string, options: SearchOptions): Promise { const rawMode = options.mode ?? 'fuzzy'; // Unified with hybrid/document/exact-scorer: unknown modes throw @@ -862,7 +554,7 @@ export class WebGPUEngine { throw new ProfileMismatchError(!this.normalized, options.caseSensitive); } - if (!this.device || !this.recordsBuffer || !this.offsetsBuffer || !this.uniformBuffer || !this.queryBuffer || !this.outputBuffer || !this.stagingBuffer || !this.substringPipeline || !this.fuzzyPipeline) { + if (!this.buffers.isSearchReady()) { return noHits(query); } @@ -872,19 +564,18 @@ export class WebGPUEngine { return noHits(query); } - const pipeline = mode === 'fuzzy' ? this.fuzzyPipeline! : this.substringPipeline!; + const pipeline = mode === 'fuzzy' ? this.state.fuzzyPipeline! : this.state.substringPipeline!; const caseSensitive = qcs; const queryLen = nq.tokens.length; const totalStart = nowMs(); const tQueryStart = nowMs(); - this.device.queue.writeBuffer(this.queryBuffer, 0, nq.tokens); const udata = new ArrayBuffer(32); const U32 = new Uint32Array(udata); U32[0] = this.currentDatasetSize; U32[1] = queryLen; - U32[2] = this.candidateCapacity; + U32[2] = this.state.candidateCapacity; // flagsAndProfile: caseSensitive:1b + normalized:1b + profile:6b + unicode:8b // + scoring:8b. Reserved for harness/debug — shaders are pure-`==` by // construction and do not read it (host enforces ProfileMismatchError). @@ -896,24 +587,14 @@ export class WebGPUEngine { U32[5] = 0; U32[6] = 0; U32[7] = 0; - this.device.queue.writeBuffer(this.uniformBuffer, 0, udata); + this.dispatch.uploadSearchRequest(nq.tokens, udata); const queryUploadMs = nowMs() - tQueryStart; throwIfAborted(options.signal); const tDispatchStart = nowMs(); - const bg = this.device.createBindGroup({ - label: 'search-bg', - layout: pipeline.getBindGroupLayout(0), - entries: [ - { binding: 0, resource: { buffer: this.uniformBuffer } }, - { binding: 1, resource: { buffer: this.offsetsBuffer } }, - { binding: 2, resource: { buffer: this.recordsBuffer } }, - { binding: 3, resource: { buffer: this.queryBuffer } }, - { binding: 4, resource: { buffer: this.outputBuffer } } - ] - }); + const bg = this.dispatch.createSearchBindGroup(pipeline); // Chunked dispatch: when ceil(rows/128) exceeds the device dimension // limit, each chunk re-writes the uniform (_pad.x = base row) and submits @@ -921,18 +602,18 @@ export class WebGPUEngine { // timestamps. Queue-ordering dependency: WebGPU guarantees serial // queue execution, so each dispatch samples only prior uniform writes // (chunk counts are tiny in practice: step ≈ 8.4M rows/chunk). - const md = (this.device.limits as unknown as Record).maxComputeWorkgroupsPerDimension; + const md = (this.state.device!.limits as unknown as Record).maxComputeWorkgroupsPerDimension; const maxDim = typeof md === 'number' && Number.isFinite(md) && md >= 1 ? Math.floor(md) : 65535; - const workgroupSize = 128; + const workgroupSize = WORKGROUP_SIZE; const totalWorkgroups = Math.ceil(this.currentDatasetSize / workgroupSize); // searchInternal only routes; single vs chunked lives above. // Chunked corpora exceed any latency budget; timestamps skipped (null). let skipTimestamps = false; if (totalWorkgroups <= maxDim) { - this.dispatchSinglePass(pipeline, bg, totalWorkgroups); + this.dispatch.dispatchSinglePass(pipeline, bg, totalWorkgroups); } else { - this.dispatchChunked(pipeline, bg, udata, U32, maxDim, workgroupSize); + this.dispatch.dispatchChunked(pipeline, bg, udata, U32, maxDim, this.currentDatasetSize); skipTimestamps = true; } const encodeSubmitMs = nowMs() - tDispatchStart; @@ -952,72 +633,37 @@ export class WebGPUEngine { // destroy-while-mapped error. mapAsync rejection itself (e.g. destroy // raced it via unmap) also maps to AbortError when the epoch moved, so // hybrid doesn't misclassify it as fallback-eligible GPU failure. - const gen = this.generation; - const stagingBuffer = this.stagingBuffer; - const tsStaging = this.queryStagingBuffer; + const gen = this.state.generation; + const stagingBuffer = this.state.stagingBuffer!; + const tsStaging = this.state.queryStagingBuffer; if (tsStaging && !skipTimestamps) { - try { - await tsStaging.mapAsync(MapMode.READ); - if (gen === this.generation) { - const tu = new BigUint64Array(tsStaging.getMappedRange()); - if (tu[1] >= (tu[0]) && (tu[0]) > 0n) { - gpuExecutionMs = Number((tu[1]) - (tu[0])) / 1_000_000; - } - } - } catch (tsErr) { - if (gen === this.generation) console.warn('timestamp read failed:', tsErr); - } finally { - try { tsStaging.unmap(); } catch {} - } + gpuExecutionMs = await this.dispatch.readTimestampGpuMs(tsStaging, gen); } // Torn down / reallocated during the timestamp await. - if (gen !== this.generation || !this.device || this.stagingBuffer !== stagingBuffer) { - throw abortError(); - } + this.dispatch.assertStagingLive(stagingBuffer, gen); let totalMatches = 0; let candidateCount = 0; let hasOverflow = false; let candidates: SearchResultItem[] = []; - try { - await stagingBuffer.mapAsync(MapMode.READ); - } catch (mapErr) { - if (gen !== this.generation) { - try { stagingBuffer.unmap(); } catch {} - throw abortError(); - } - throw mapErr; - } + await this.dispatch.mapStagedOutput(stagingBuffer, gen); try { // Re-check epoch + liveness after the await: destroy() nulls // device/buffers, reallocation swaps them. - if (gen !== this.generation || !this.device || this.stagingBuffer !== stagingBuffer) { - throw abortError(); - } + this.dispatch.assertStagingLive(stagingBuffer, gen); throwIfAborted(options.signal); - const arrayBuffer = stagingBuffer.getMappedRange(); - const u32Read = new Uint32Array(arrayBuffer); - const i32Read = new Int32Array(arrayBuffer); - - totalMatches = u32Read[0]; - candidateCount = Math.min(totalMatches, this.candidateCapacity); - hasOverflow = totalMatches > this.candidateCapacity; - - candidates = new Array(candidateCount); - for (let i = 0; i < candidateCount; i++) { - const index = u32Read[2 + i * 2]; - const score = i32Read[3 + i * 2]; - // Token-only datasets carry no strings: resolve `text` to `''` - // (documented degradation; see DatasetLike). - candidates[i] = { - index, - score, - text: this.currentStrings ? (this.currentStrings[index] ?? '') : '' - }; - } + const decoded = this.dispatch.readMappedOutput( + stagingBuffer, + this.state.candidateCapacity, + (index: number) => (this.currentStrings ? (this.currentStrings[index] ?? '') : '') + ); + totalMatches = decoded.totalMatches; + candidateCount = decoded.candidateCount; + hasOverflow = decoded.hasOverflow; + candidates = decoded.candidates; } finally { try { stagingBuffer.unmap(); } catch {} } @@ -1067,25 +713,13 @@ export class WebGPUEngine { // readback site; rejected: Promise.race abandons the mapping). Unmap // (never destroy a mapped buffer), destroy, then park a barrier on the // mutex so queued ops observe the torn-down state. - this.generation++; - try { this.stagingBuffer?.unmap(); } catch {} - try { this.queryStagingBuffer?.unmap(); } catch {} - const bufs = [this.offsetsBuffer, this.recordsBuffer, this.uniformBuffer, this.queryBuffer, this.outputBuffer, this.stagingBuffer, this.queryResolveBuffer, this.queryStagingBuffer]; - for (const b of bufs) { try { b?.destroy(); } catch {} } - this.offsetsBuffer = this.recordsBuffer = this.uniformBuffer = this.queryBuffer = null; - this.outputBuffer = this.stagingBuffer = this.queryResolveBuffer = this.queryStagingBuffer = null; - if (this.querySet) { - try { this.querySet.destroy(); } catch {} - this.querySet = null; - } - this.substringPipeline = null; - this.fuzzyPipeline = null; + this.state.generation++; + this.buffers.destroyForTeardown(); + this.pipelines.drop(); this.currentStrings = null; this.currentTokens = null; this.currentOffsets = null; this.currentDatasetSize = 0; - this.allocatedOffsetsByteLength = 0; - this.allocatedRecordsByteLength = 0; this.releaseCurrentDevice(); const p = this.searchMutex.then(() => {}, () => {}); this.searchMutex = p.catch(() => {}); diff --git a/packages/webgpu-search/test/engine-pool-regressions.test.ts b/packages/webgpu-search/test/engine-pool-regressions.test.ts index afedafe..218982b 100644 --- a/packages/webgpu-search/test/engine-pool-regressions.test.ts +++ b/packages/webgpu-search/test/engine-pool-regressions.test.ts @@ -204,34 +204,42 @@ describe('headroom clamp vs binding limit (finding 2)', () => { }); describe('teardown during timestamp readback (finding 3)', () => { - async function readyEngine() { + async function readyEngine(onTsMapAsync: (engine: WebGPUEngine) => void) { + // Inject at the real allocation boundary: the fake `ts-staging` buffer is + // handed out by device.createBuffer for the timestamp staging descriptor, + // so no engine-private field or test-only seam is touched. The vgpu mock + // adapter does not propagate `features` onto the device, so advertise + // timestamp-query explicitly — readTimestampGpuMs only runs when the + // device exposes the feature (as production would). + const dev = await mockDevice(); + dev.features = new Set(['timestamp-query']) as any; + const createBuffer = (dev.createBuffer as any).bind(dev); + (dev as any).createBuffer = (d: GPUBufferDescriptor) => { + if (d.label === 'ts-staging') { + return { + mapAsync: async () => { onTsMapAsync(engine); }, + getMappedRange: () => new ArrayBuffer(16), + unmap() {}, + destroy() {}, + }; + } + return createBuffer(d); + }; const engine = new WebGPUEngine(); - await engine.init(await mockDevice()); + await engine.init(dev); await engine.loadDataset(['hello world', 'help']); return engine; } test('destroy() during timestamp mapAsync rejects with AbortError', async () => { - const engine = await readyEngine(); - (engine as any).queryStagingBuffer = { - mapAsync: async () => { engine.destroy(); }, - getMappedRange: () => new ArrayBuffer(16), - unmap() {}, - destroy() {}, - }; + const engine = await readyEngine((e) => e.destroy()); const err = await engine.search('hel', { mode: 'fuzzy' }).then(() => null, (e) => e); expect(err).not.toBeNull(); expect((err as Error).name).toBe('AbortError'); }); test('output reallocation during timestamp mapAsync rejects with AbortError', async () => { - const engine = await readyEngine(); - (engine as any).queryStagingBuffer = { - mapAsync: async () => { engine.ensureCandidateCapacity(20_000); }, - getMappedRange: () => new ArrayBuffer(16), - unmap() {}, - destroy() {}, - }; + const engine = await readyEngine((e) => e.ensureCandidateCapacity(20_000)); const err = await engine.search('hel', { mode: 'fuzzy' }).then(() => null, (e) => e); expect((err as Error | null)?.name).toBe('AbortError'); engine.destroy(); diff --git a/scripts/check-public-api.ts b/scripts/check-public-api.ts index 797e00d..6df106e 100644 --- a/scripts/check-public-api.ts +++ b/scripts/check-public-api.ts @@ -102,12 +102,48 @@ for (const key of ['profileId', 'scoringVersion', 'cpuScorer', 'fallbackReason', if (!types.includes(key)) fail(`SearchResponse echo key missing in src/types.ts: ${key}`); } -// --- 5. limit/maxResults alias ---------------------------------------------- -const guardFiles = [join(pkgDir, 'src', 'search-index.ts'), join(pkgDir, 'src', 'document-index.ts'), join(pkgDir, 'src', 'webgpu-engine.ts')]; -for (const f of guardFiles) { - const body = read(f); - if (!body.includes('options.limit ?? options.maxResults ?? 50') && !body.includes('limit ?? maxResults ?? 50')) { - fail(`limit/maxResults alias broken in ${f} (expected "limit ?? maxResults ?? 50" with limit precedence).`); +// --- 5. limit/maxResults alias (public behavioral check) --------------------- +// Frozen contract: default 50, maxResults honored, limit honored, and limit +// wins when both are set. Checked at RUNTIME through both public gates +// (SearchIndex + DocumentIndex) instead of grepping source files, so the alias +// expression may move between internal modules (e.g. document/query.ts) with +// the modular refactor without breaking the gate; breaking the behavior does. +type AliasSearchable = { + search(query: string, options?: Record): Promise<{ results: unknown[] }>; + destroy?(): void; +}; +{ + const aliasVecs: Array<[Record, number, string]> = [ + [{}, 50, 'default clamps to 50'], + [{ maxResults: 7 }, 7, 'maxResults alias honored'], + [{ limit: 3 }, 3, 'limit honored'], + [{ limit: 2, maxResults: 9 }, 2, 'limit precedes maxResults'], + ]; + const gates: Array<[string, AliasSearchable]> = []; + const si = await ((mod as Record)['SearchIndex'] as { + create(items: unknown[], options?: Record): Promise; + }).create( + Array.from({ length: 60 }, (_, i) => `aliasdoc-${i}`), + { preferGpu: false } + ); + gates.push(['SearchIndex', si]); + const di = await ((mod as Record)['DocumentIndex'] as { + create(records: T[], options: Record): Promise; + }).create( + Array.from({ length: 60 }, (_, i) => ({ id: i, title: `aliasdoc-${i} note` })), + { fields: [{ name: 'title' }], preferGpu: false } + ); + gates.push(['DocumentIndex', di]); + for (const [gate, searcher] of gates) { + for (const [opts, want, label] of aliasVecs) { + const res = await searcher.search('aliasdoc', { ...opts, mode: 'substring' }); + if (res.results.length !== want) { + fail(`${gate} limit/maxResults alias broken: ${label} (got ${res.results.length}, want ${want})`); + } + } + } + for (const [, searcher] of gates) { + if (typeof searcher.destroy === 'function') searcher.destroy(); } } if (!types.includes('maxResults?: number')) fail('SearchOptions.maxResults alias missing in src/types.ts.'); diff --git a/scripts/test-document-records.ts b/scripts/test-document-records.ts index 4b471a0..e97b81a 100644 --- a/scripts/test-document-records.ts +++ b/scripts/test-document-records.ts @@ -93,8 +93,8 @@ async function runM2Tests() { console.log('2. Testing field-stratified packing priority (descending weights)...'); { const items: ArticleDoc[] = [ - { id: '1', title: 'Alpha', content: 'Beta', tags: ['Gamma'] }, - { id: '2', title: 'Delta', content: 'Epsilon', tags: ['Zeta'] } + { id: '1', title: 'Hi', content: 'Beta Gamma Delta', tags: ['Uno Duo'] }, + { id: '2', title: 'Hi', content: 'Beta Gamma Delta', tags: ['Uno Duo'] } ]; // Give 'content' highest weight (3.0), 'tags' (2.0), 'title' (0.5) @@ -113,20 +113,20 @@ async function runM2Tests() { assert.strictEqual(stats.docCount, 2); assert.strictEqual(stats.rowCount, 6); // 2 docs * 3 fields - // Inspect internal sortedFields: content (3.0), tags (2.0), title (0.5) - const sorted = (index as any).sortedFields; - assert.strictEqual(sorted[0].name, 'content'); - assert.strictEqual(sorted[1].name, 'tags'); - assert.strictEqual(sorted[2].name, 'title'); - - // Inspect internal rowToFieldIndex: rows 0,1 -> field 0 (content); rows 2,3 -> field 1 (tags); rows 4,5 -> field 2 (title) - const rowToField = (index as any).rowToFieldIndex; - assert.strictEqual(rowToField[0], 0); - assert.strictEqual(rowToField[1], 0); - assert.strictEqual(rowToField[2], 1); - assert.strictEqual(rowToField[3], 1); - assert.strictEqual(rowToField[4], 2); - assert.strictEqual(rowToField[5], 2); + // Stratification verified through the public getter surface instead of + // private state: `getSortedFields()` reports the descending weight order + // (content 3.0, tags 2.0, title 0.5). + const sorted = index.getSortedFields().map((f) => f.name); + assert.deepStrictEqual(sorted, ['content', 'tags', 'title']); + + // Row stratification verified via `getRowTokens()`: each row holds one + // token stream per (doc, field), and each field uses a distinct folded + // text length so a row's length identifies its field: content rows + // (weight 3.0) are packed first ('Beta Gamma Delta' = 16), tags rows + // (weight 2.0) next ('Uno Duo' = 7), title rows (weight 0.5) last + // ('Hi' = 2). + const rowLengths = index.getRowTokens().map((row) => row.length); + assert.deepStrictEqual(rowLengths, [16, 16, 7, 7, 2, 2]); console.log(' ✅ Field-stratified packing in descending weight order verified'); } @@ -137,51 +137,98 @@ async function runM2Tests() { console.log('3. Testing dynamic candidate capacity scaling (up to 32,768)...'); { // Formula: min(32768, max(8192, docCount * fieldCount * 0.1)) - // Small corpus: 10 docs * 2 fields * 0.1 = 2 -> min capacity = 8192 - const smallDocs = Array.from({ length: 10 }, (_, i) => ({ + // + // `candidateCapacity` is an internal budget with no public getter, so it + // is verified through the caller-visible overflow contract instead of a + // private field read: + // candidateCount = min(totalMatches, candidateCapacity) + // hasOverflow = totalMatches > candidateCapacity + // + the diagnostics overflow warning literally names the effective + // capacity ("candidate pool capacity N"), so an unclamped 50000, + // an unclamped 100, a leaked NaN, or a missing scaling formula each + // produce a distinguishable observation. + + // Default formula scaling: 100000 docs * 1 field * 0.1 = 10000 + const scaledDocs = Array.from({ length: 100000 }, (_, i) => ({ id: `id-${i}`, - title: `Title ${i}` + f1: 'boost' })); - const smallIndex = await DocumentIndex.create(smallDocs, { - fields: ['title'], + const scaledIndex = await DocumentIndex.create(scaledDocs, { + fields: ['f1'], preferGpu: false }); - assert.strictEqual((smallIndex as any).candidateCapacity, 8192); + try { + const scaledRes = await scaledIndex.search('boost', { mode: 'substring', diagnostics: true, limit: 5 }); + assert.strictEqual(scaledRes.totalMatches, 100000, 'Every scaled-corpus document must match'); + assert.strictEqual(scaledRes.hasOverflow, true, 'Scaled corpus must exceed the computed capacity'); + assert.strictEqual(scaledRes.candidateCount, 10000, 'Effective capacity must scale to 10000 (100000 * 1 * 0.1)'); + assert( + scaledRes.diagnostics?.warnings?.some((w) => w.includes('candidate pool capacity 10000')), + 'Overflow warning must report the computed 10000 capacity' + ); + } finally { + scaledIndex.destroy(); + } - // Explicit candidateCapacity override clamped to 32,768 - const largeCapIndex = await DocumentIndex.create(smallDocs, { - fields: ['title'], + // Explicit candidateCapacity override clamped to 32,768: + // 35000 matches overflow ONLY because the 50000 request was clamped down + // to 32768 (unclamped, 35000 < 50000 would not overflow). + const largeCapDocs = Array.from({ length: 35000 }, (_, i) => ({ + id: `id-${i}`, + f1: 'boost' + })); + const largeCapIndex = await DocumentIndex.create(largeCapDocs, { + fields: ['f1'], candidateCapacity: 50000, preferGpu: false }); - assert.strictEqual((largeCapIndex as any).candidateCapacity, 32768); + try { + const clampRes = await largeCapIndex.search('boost', { mode: 'substring', diagnostics: true, limit: 5 }); + assert.strictEqual(clampRes.hasOverflow, true, '35000 matches must overflow the clamped 32768 capacity'); + assert.strictEqual(clampRes.candidateCount, 32768, 'An override of 50000 must clamp down to 32768'); + assert( + clampRes.diagnostics?.warnings?.some((w) => w.includes('candidate pool capacity 32768')), + 'Overflow warning must report the clamped 32768 capacity' + ); + } finally { + largeCapIndex.destroy(); + } - // Explicit candidateCapacity clamped to minimum 8192 - const minCapIndex = await DocumentIndex.create(smallDocs, { - fields: ['title'], + // Explicit low override clamped up to minimum 8192, and NaN resets to + // 8192: an unclamped user cap of 100 (or a leaked NaN) would name a + // different capacity in the overflow warning. + const minCapDocs = Array.from({ length: 9000 }, (_, i) => ({ + id: `id-${i}`, + f1: 'boost' + })); + const minCapIndex = await DocumentIndex.create(minCapDocs, { + fields: ['f1'], candidateCapacity: 100, preferGpu: false }); - assert.strictEqual((minCapIndex as any).candidateCapacity, 8192); - - // Large simulated corpus scaling - const largeDocs = Array.from({ length: 20000 }, (_, i) => ({ - id: `id-${i}`, - f1: `v1-${i}`, - f2: `v2-${i}` - })); - // 20,000 docs * 2 fields * 0.1 = 4,000 -> still 8192 - // If 100,000 docs * 2 fields * 0.1 = 20,000 -> 20,000 - const mediumDocs = Array.from({ length: 100000 }, (_, i) => ({ - id: `id-${i}`, - f1: `v1-${i}` - })); - // 100,000 docs * 1 field * 0.1 = 10,000 - const scaledIndex = await DocumentIndex.create(mediumDocs, { + const nanCapIndex = await DocumentIndex.create(minCapDocs, { fields: ['f1'], + candidateCapacity: NaN, preferGpu: false }); - assert.strictEqual((scaledIndex as any).candidateCapacity, 10000); + try { + const clampedScenarios = [ + { name: 'candidateCapacity:100 clamps up to 8192', idx: minCapIndex }, + { name: 'candidateCapacity:NaN resets to 8192', idx: nanCapIndex } + ] as const; + for (const scenario of clampedScenarios) { + const res = await scenario.idx.search('boost', { mode: 'substring', diagnostics: true, limit: 5 }); + assert.strictEqual(res.hasOverflow, true, `${scenario.name}: 9000 matches must overflow capacity 8192`); + assert.strictEqual(res.candidateCount, 8192, `${scenario.name}: effective capacity must clamp to 8192`); + assert( + res.diagnostics?.warnings?.some((w) => w.includes('candidate pool capacity 8192')), + `${scenario.name}: warning must report the clamped 8192 capacity` + ); + } + } finally { + nanCapIndex.destroy(); + minCapIndex.destroy(); + } // WebGPUEngine ensureCandidateCapacity verification const gpuEngine = new WebGPUEngine(); @@ -616,15 +663,18 @@ async function runM2Tests() { // ========================================================================= console.log('11. Testing multi-agent review hardening and validation guards...'); { - // 11a. candidateCapacity NaN / negative handling + // 11a. candidateCapacity NaN handling const docs = [{ id: '1', title: 'Test Document', body: 'Content' }]; const nanCapIndex = await DocumentIndex.create(docs, { fields: ['title', 'body'], preferGpu: false, candidateCapacity: NaN }); - // Internal candidate capacity safely clamped to default 8192 - assert.strictEqual((nanCapIndex as any).candidateCapacity, 8192); + // A leaked NaN capacity would surface as candidateCount = min(matches, NaN) = NaN; + // a clamped NaN resets to the 8192 floor, so this 1-match search reports candidateCount 1. + const nanRes = await nanCapIndex.search('Test'); + assert.strictEqual(nanRes.candidateCount, 1, 'NaN candidateCapacity must not leak into candidateCount'); + assert.strictEqual(nanRes.hasOverflow, false, 'A single match must not overflow the clamped capacity'); nanCapIndex.destroy(); // 11b. WebGPUEngine ensureCandidateCapacity NaN handling diff --git a/scripts/test-faceting.ts b/scripts/test-faceting.ts index a4f5619..f57e5a1 100644 --- a/scripts/test-faceting.ts +++ b/scripts/test-faceting.ts @@ -58,6 +58,103 @@ function productIndexOpts(extra?: Partial>): Docum }; } +// =========================================================================== +// GPU-boundary output simulator +// +// The vgpu mock never executes compute (dispatchWorkgroups is a no-op), so a +// real WebGPUEngine readback always decodes an empty candidate pool. To +// exercise overflow routing through the production path, the mock device is +// wrapped so a test can *program* the search output buffer with the kernel's +// readback layout — u32[0] = total match count, then (u32 row index, i32 +// score) pairs. The programmed bytes land when the engine copies the output +// buffer to its MAP_READ staging twin (the post-dispatch moment a real GPU +// would have filled it), and are decoded by the engine's own mapAsync → +// readMappedOutput path. Injection only fires for STORAGE-backed copies large +// enough to carry the payload, so the 16-byte timestamp staging copy is never +// touched. Candidate slots beyond the programmed pairs decode as (row 0, +// score 0), which dedupes into doc row 0 — harmless padding. No production +// seams: the wrapper is a plain WebGPU-device facade over public APIs. +// =========================================================================== + +/** GPUBufferUsage.STORAGE (spec value; Node/Bun has no WebGPU global). */ +const GPU_BUFFER_USAGE_STORAGE = 0x0080; + +interface SimulatedGpuSearchOutput { + /** Decoded `totalMatches`; drives `hasOverflow = total > candidateCapacity`. */ + totalMatches: number; + /** Candidate pool: field-stratified row indices + raw match scores. */ + rows: Array<{ index: number; score: number }>; +} + +function createSearchOutputSimulator(device: GPUDevice): { + device: GPUDevice; + /** Arm (or disarm with `null`) the simulated output for the next searches. */ + program: (output: SimulatedGpuSearchOutput | null) => void; +} { + let pending: SimulatedGpuSearchOutput | null = null; + + const encodeOutput = (out: SimulatedGpuSearchOutput): Uint8Array => { + // Readback layout (GpuDispatcher.readMappedOutput): u32[0] total matches, + // u32[1] padding, then (u32 row index, i32 score) pairs. + const u32 = new Uint32Array(2 + out.rows.length * 2); + const i32 = new Int32Array(u32.buffer); + u32[0] = out.totalMatches; + for (let i = 0; i < out.rows.length; i++) { + u32[2 + i * 2] = out.rows[i].index; + i32[3 + i * 2] = out.rows[i].score; + } + return new Uint8Array(u32.buffer); + }; + + const simulated = new Proxy(device, { + get(target, prop) { + if (prop === 'createCommandEncoder') { + return (desc?: GPUCommandEncoderDescriptor) => { + const enc = target.createCommandEncoder(desc); + return new Proxy(enc, { + get(encTarget, encProp) { + if (encProp === 'copyBufferToBuffer') { + return ( + source: GPUBuffer, + sourceOffset: number, + destination: GPUBuffer, + destinationOffset: number, + size?: number + ) => { + const copyBytes = size ?? source.size - sourceOffset; + if (pending) { + const payload = encodeOutput(pending); + if ( + (source.usage & GPU_BUFFER_USAGE_STORAGE) !== 0 && + copyBytes >= payload.byteLength && + source.size >= payload.byteLength + ) { + // The compute dispatch has "run": the output buffer now + // holds the programmed pool for the engine to copy into + // its staging twin and decode. + target.queue.writeBuffer(source, 0, payload); + } + } + return enc.copyBufferToBuffer(source, sourceOffset, destination, destinationOffset, size); + }; + } + const value = Reflect.get(encTarget, encProp); + return typeof value === 'function' ? (value as (...args: unknown[]) => unknown).bind(encTarget) : value; + }, + }); + }; + } + const value = Reflect.get(target, prop); + return typeof value === 'function' ? (value as (...args: unknown[]) => unknown).bind(target) : value; + }, + }); + + return { + device: simulated, + program: (output) => { pending = output; }, + }; +} + async function runTests() { console.log('=== Running Issue #10 Facet Aggregation Engine ===\n'); @@ -428,13 +525,14 @@ async function runTests() { } // ========================================================================= - // 7. GPU parity + overflow semantics (mock adapter + stubbed engine) + // 7. GPU parity + overflow semantics (mock adapter + programmed output) // // NOTE: the vgpu mock stubs compute dispatch (it never produces match // rows), so the repo-wide convention (test-vgpu-mock.ts, filter tests) // asserts GPU *routing*, not GPU match quality. Overflow/approximate - // routing — which needs a capped candidate pool — is covered with a - // stubbed WebGPUEngine behind the same readback contract. + // routing — which needs a capped candidate pool — programs the search + // output buffer on the same mock device (the kernel's readback format) + // and runs the engine's real copy → mapAsync → decode path. // ========================================================================= console.log('7. Testing GPU facet parity and overflow semantics...'); { @@ -458,40 +556,32 @@ async function runTests() { assert.strictEqual(mockRes.facets!.c.isApproximate, false); assert.deepStrictEqual(mockRes.facets!.c.buckets, []); - // 7b. Stubbed engine: rows are field-stratified, so title rows for the - // 8-doc corpus are 0..7 and body rows are 8..15. - const stubGpu = (poolRows: number[], hasOverflow: boolean) => ({ - isReady: true, - destroy() {}, - search: async (query: string, opts: { mode?: string }) => ({ - query, - mode: opts.mode ?? 'fuzzy', - engine: 'webgpu', - totalMatches: poolRows.length + (hasOverflow ? 5000 : 0), - candidateCount: poolRows.length, - hasOverflow, - results: poolRows.map((index) => ({ index, score: 500, text: '' })), - timings: { - queryUploadMs: 0, encodeSubmitMs: 0, gpuExecutionMs: null, - readbackMs: 0, totalMs: 1, gpuDispatchMs: 0, - }, - profileId: 'unicode-default', - scoringVersion: 'parity-v1', - cpuScorer: 'exact', - }), - }); - - const stubIndex = await DocumentIndex.create(CORPUS, productIndexOpts()); - (stubIndex as unknown as { engineState: { engine: string } }).engineState = { engine: 'webgpu' }; + // 7b. Overflow semantics through the real engine, simulated at the GPU + // boundary. The mock device is extended so each search's output buffer + // carries a programmed candidate pool (the kernel's readback format); + // everything downstream — staging copy, mapAsync decode, pool-driven + // facet aggregation, overflow routing — is the production path. + // Rows are field-stratified: title rows 0..7 map to docs p1..p8, so + // rows 0..2 match docs p1..p3 (all 'tech'). + const { device: simulatedDevice, program } = createSearchOutputSimulator(mockDevice); + const simulatedIndex = await DocumentIndex.create( + CORPUS, + productIndexOpts({ device: simulatedDevice, preferGpu: true }) + ); - // Capped pool (docs p1..p3) with overflow, auto mode -> approximate. - (stubIndex as unknown as { gpuEngine: unknown }).gpuEngine = stubGpu([0, 1, 2], true); - const autoRes = await stubIndex.search('Widget', { + // Capped pool (docs p1..p3) with the match counter one past the 8192 + // candidate capacity -> overflow, auto mode -> approximate facets over + // the pool only. + program({ + totalMatches: 8193, // GPU candidate capacity (8192) + 1 + rows: [0, 1, 2].map((index) => ({ index, score: 500 })), + }); + const autoRes = await simulatedIndex.search('Widget', { facets: { c: { type: 'terms', field: 'category' } }, }); assert.strictEqual(autoRes.engine, 'webgpu'); assert.strictEqual(autoRes.hasOverflow, true); - assert.strictEqual(autoRes.totalMatches, 3); + assert.strictEqual(autoRes.totalMatches, 3); // pool decodes to 3 distinct docs if (autoRes.facets!.c.type !== 'terms') throw new Error('unreachable'); assert.strictEqual(autoRes.facets!.c.isApproximate, true); assert.deepStrictEqual( @@ -499,8 +589,10 @@ async function runTests() { [['tech', 3]] ); - // Same capped pool with force-exact -> full CPU rescan, exact counts. - const exactRes = await stubIndex.search('Widget', { + // Same capped pool with force-exact -> full CPU rescan: hits still come + // from the GPU pool, but facets cover the whole corpus exactly (all four + // categories, matching the CPU reference). + const exactRes = await simulatedIndex.search('Widget', { facets: { c: { type: 'terms', field: 'category' } }, faceting: 'force-exact', }); @@ -512,8 +604,11 @@ async function runTests() { assert.deepStrictEqual(exactRes.facets!.c.buckets, cpuRef.facets!.c.buckets); // Capped pool without overflow flag -> exact over the pool. - (stubIndex as unknown as { gpuEngine: unknown }).gpuEngine = stubGpu([0, 1, 2], false); - const noOverflowRes = await stubIndex.search('Widget', { + program({ + totalMatches: 3, + rows: [0, 1, 2].map((index) => ({ index, score: 500 })), + }); + const noOverflowRes = await simulatedIndex.search('Widget', { facets: { c: { type: 'terms', field: 'category' }, p: { type: 'range', field: 'price', ranges: [{ to: 150 }, { from: 150 }] }, @@ -533,6 +628,9 @@ async function runTests() { [2, 1] // 49.99, 129.50 vs 199.99 ); + // Disarm so later sections never see stale simulated output. + program(null); + // 7c. CPU route stays exact past candidate capacity (9000 docs). const N = 9000; const bulk: Product[] = new Array(N); @@ -557,7 +655,7 @@ async function runTests() { gpuIndex.destroy(); cpuIndex.destroy(); - stubIndex.destroy(); + simulatedIndex.destroy(); cpuBulk.destroy(); console.log(' ✅ GPU parity & overflow semantics verified 100%'); } diff --git a/scripts/test-parity-harness.ts b/scripts/test-parity-harness.ts index 3771288..0e3b286 100644 --- a/scripts/test-parity-harness.ts +++ b/scripts/test-parity-harness.ts @@ -141,6 +141,36 @@ function multisetsEqual(a: Map, b: Map): boolean return true; } +/** + * Capture an engine's output staging buffer through the injected-device + * boundary: window-wrap `device.createBuffer` (same pattern as the + * records-OOM injection below) around `create()` and pick out the + * 'Staging Buffer' label. Mock-boundary injection only -- never reads engine + * private state. `create()` must be the only staging-allocating call; the + * first 'Staging Buffer' created is captured. + */ +async function withCapturedStaging( + device: GPUDevice, + create: () => Promise +): Promise<{ value: T; staging: GPUBuffer | null }> { + const mutable = device as unknown as { + createBuffer: (descriptor: GPUBufferDescriptor) => GPUBuffer; + }; + const realCreate = mutable.createBuffer.bind(device); + let staging: GPUBuffer | null = null; + mutable.createBuffer = (descriptor: GPUBufferDescriptor): GPUBuffer => { + const buf = realCreate(descriptor); + if (descriptor.label === 'Staging Buffer' && staging === null) staging = buf; + return buf; + }; + try { + const value = await create(); + return { value, staging }; + } finally { + mutable.createBuffer = realCreate; + } +} + const CELL_TIMEOUT_MS = 15000; function parseShardEnv(raw: string | undefined, def: number): number { if (raw === undefined || raw === '') return def; @@ -581,19 +611,29 @@ async function main(): Promise { '(m) fast-path timestamp key well-formed', 'gpuExecutionMs' in fast.timings && (fast.timings.gpuExecutionMs === null || Number.isFinite(fast.timings.gpuExecutionMs)) ); - // Force chunked: maxDim=1 -> step=128 rows, 300 rows -> 3 chunks + clear + copy = 5 encoders. - const realDevice = (eng as unknown as { device: GPUDevice }).device; - (eng as unknown as { device: GPUDevice }).device = { - ...realDevice, - limits: { ...(realDevice as unknown as { limits: Record }).limits, maxComputeWorkgroupsPerDimension: 1 }, + // Force chunked through the public device boundary: init a second engine on + // a cloned injected mock device whose dimension limit is 1. The dispatcher + // reads maxComputeWorkgroupsPerDimension from the public device.limits at + // dispatch time; the clone shares the mock's instrumentation object, so + // encoder deltas stay comparable. No engine private state is touched. + const limitedDevice = { + ...mockDevice, + limits: { ...(mockDevice as unknown as { limits: Record }).limits, maxComputeWorkgroupsPerDimension: 1 }, } as unknown as GPUDevice; + const engLimited = new WebGPUEngine(); + const initLim = await engLimited.init(limitedDevice); + ok('(m) limited-dimension engine inits', initLim); + await engLimited.loadDataset(rows); const e1 = calls.createCommandEncoder; - const chunked = await eng.search('chunk-row-1', { mode: 'substring' }); + const chunked = await engLimited.search('chunk-row-1', { mode: 'substring' }); const delta = calls.createCommandEncoder - e1; ok('(m) chunked multi-dispatch (5 encoders)', delta === 5, `delta=${delta}`); ok('(m) chunked skips timestamps (null)', chunked.timings.gpuExecutionMs === null); - (eng as unknown as { device: GPUDevice }).device = realDevice; - ok('(m) grow-only output 65544 B', (eng as unknown as { outputByteLength: number }).outputByteLength === 65544); + // Grow-only output: base candidate pool stays 8192 slots after loads + // (65544 B = 8 B counter + 8192 x 8 B) -- pinned via the public capacity + // getter, which the output/staging allocation derives byte length from. + ok('(m) grow-only output 65544 B', eng.currentCandidateCapacity === 8192 && 8 + eng.currentCandidateCapacity * 8 === 65544 && engLimited.currentCandidateCapacity === 8192); + engLimited.destroy(); await eng.loadDataset([]); const e2 = calls.createCommandEncoder; const zr = await eng.search('chunk-row-1', { mode: 'substring' }); @@ -875,70 +915,93 @@ async function main(): Promise { // mapAsync reject: epoch intact -> raw error (fallback-eligible); epoch moved -> AbortError. { - const eng = new WebGPUEngine(); - await eng.init(mockDevice); - await eng.loadDataset(['apple']); - const sb = (eng as unknown as { stagingBuffer: { mapAsync: () => Promise } }).stagingBuffer; - const orig = sb.mapAsync.bind(sb); - try { - sb.mapAsync = () => Promise.reject(new Error('injected mapAsync failure')); - let raw = false; - try { - await eng.search('apple', { mode: 'substring' }); - } catch (e) { - raw = (e as Error).message === 'injected mapAsync failure'; - } - ok('mapAsync reject propagates when epoch intact', raw); - // Move the epoch *inside* the in-flight await (destroy-race shape): - // gen is captured before mapAsync, so a bump during the await reads - // as moved and must map to AbortError (never fallback-eligible). - sb.mapAsync = () => { - (eng as unknown as { generation: number }).generation += 1; - return Promise.reject(new Error('injected mapAsync failure')); - }; - let aborted = false; - try { - await eng.search('apple', { mode: 'substring' }); - } catch (e) { - aborted = (e as { name?: string }).name === 'AbortError'; - } - ok('mapAsync reject + moved epoch -> AbortError', aborted); - } finally { + // Capture the engine's output staging buffer through the injected-device + // boundary (windowed createBuffer wrap, same pattern as the records-OOM + // injection); map-failure injection then patches only that buffer, and + // the epoch plays out through real public lifecycle operations. + const { value: eng, staging } = await withCapturedStaging(mockDevice, async () => { + const e = new WebGPUEngine(); + await e.init(mockDevice); + await e.loadDataset(['apple']); + return e; + }); + const sb = staging as unknown as { mapAsync: (mode?: number) => Promise } | null; + ok('mapAsync probe captured staging via device boundary', sb !== null); + if (sb) { + const orig = sb.mapAsync.bind(sb); try { - sb.mapAsync = orig; - } catch { - // Engine may have torn down stagingBuffer; restore best-effort. + sb.mapAsync = (): Promise => Promise.reject(new Error('injected mapAsync failure')); + let raw = false; + try { + await eng.search('apple', { mode: 'substring' }); + } catch (e) { + raw = (e as Error).message === 'injected mapAsync failure'; + } + ok('mapAsync reject propagates when epoch intact', raw); + // Move the epoch *inside* the in-flight await (destroy-race shape): + // gen is captured before mapAsync, so a bump during the await reads + // as moved and must map to AbortError (never fallback-eligible). + // The bump is a real public lifecycle operation: + // ensureCandidateCapacity reallocates output/staging and advances + // the epoch (no private generation mutation). + sb.mapAsync = (): Promise => { + eng.ensureCandidateCapacity(10000); + return Promise.reject(new Error('injected mapAsync failure')); + }; + let aborted = false; + try { + await eng.search('apple', { mode: 'substring' }); + } catch (e) { + aborted = (e as { name?: string }).name === 'AbortError'; + } + ok('mapAsync reject + moved epoch -> AbortError', aborted); + } finally { + try { + sb.mapAsync = orig; + } catch { + // Engine may have torn down stagingBuffer; restore best-effort. + } + eng.destroy(); } + } else { eng.destroy(); } } // AbortError through SearchIndex never converts to CPU fallback. { - const idx = await SearchIndex.create(['apple'], { device: mockDevice, preferGpu: true }); - const engHandle = (idx as unknown as { gpuEngine: WebGPUEngine }).gpuEngine; - const sb = (engHandle as unknown as { stagingBuffer: { mapAsync: () => Promise } }).stagingBuffer; - const orig = sb.mapAsync.bind(sb); - try { - // Epoch moves mid-flight (destroy-race shape) -> AbortError, which the - // hybrid must rethrow instead of converting to a CPU fallback. - sb.mapAsync = () => { - (engHandle as unknown as { generation: number }).generation += 1; - return Promise.reject(new Error('injected mid-flight failure')); - }; - let aborted = false; - try { - await idx.search('apple', { mode: 'substring' }); - } catch (e) { - aborted = (e as { name?: string }).name === 'AbortError'; - } - ok('AbortError propagates through SearchIndex (no fallback)', aborted); - } finally { + // Staging captured via the injected-device boundary; the mid-flight + // epoch bump is the real public lifecycle destroy() (exactly the + // destroy-race the abort guard exists for). The hybrid must rethrow + // the resulting AbortError instead of converting to a CPU fallback. + const { value: idx, staging } = await withCapturedStaging(mockDevice, () => + SearchIndex.create(['apple'], { device: mockDevice, preferGpu: true }) + ); + const sb = staging as unknown as { mapAsync: (mode?: number) => Promise } | null; + ok('AbortError probe captured staging via device boundary', sb !== null); + if (sb) { + const orig = sb.mapAsync.bind(sb); try { - sb.mapAsync = orig; - } catch { - // Best-effort restore; index destroy below still runs. + sb.mapAsync = (): Promise => { + idx.destroy(); + return Promise.reject(new Error('injected mid-flight failure')); + }; + let aborted = false; + try { + await idx.search('apple', { mode: 'substring' }); + } catch (e) { + aborted = (e as { name?: string }).name === 'AbortError'; + } + ok('AbortError propagates through SearchIndex (no fallback)', aborted); + } finally { + try { + sb.mapAsync = orig; + } catch { + // Best-effort restore; index destroy below still runs. + } + idx.destroy(); } + } else { idx.destroy(); } } @@ -1061,30 +1124,39 @@ async function main(): Promise { } // Abort mid-mapAsync via controller (deterministic: abort inside mapAsync). { - const eng = new WebGPUEngine(); - await eng.init(mockDevice); - await eng.loadDataset(['apple']); - const sb = (eng as unknown as { stagingBuffer: { mapAsync: () => Promise } }).stagingBuffer; - const orig = sb.mapAsync.bind(sb); - try { - const ac = new AbortController(); - sb.mapAsync = async (): Promise => { - ac.abort(); - return orig(); - }; - let aborted = false; - try { - await eng.search('apple', { mode: 'substring', signal: ac.signal }); - } catch (e) { - aborted = (e as { name?: string }).name === 'AbortError'; - } - ok('abort during mapAsync -> AbortError (epoch discard)', aborted); - } finally { + // Staging captured via the injected-device boundary (windowed createBuffer). + const { value: eng, staging } = await withCapturedStaging(mockDevice, async () => { + const e = new WebGPUEngine(); + await e.init(mockDevice); + await e.loadDataset(['apple']); + return e; + }); + const sb = staging as unknown as { mapAsync: (mode?: number) => Promise } | null; + ok('abort probe captured staging via device boundary', sb !== null); + if (sb) { + const orig = sb.mapAsync.bind(sb); try { - sb.mapAsync = orig; - } catch { - // Best-effort restore. + const ac = new AbortController(); + sb.mapAsync = async (): Promise => { + ac.abort(); + return orig(); + }; + let aborted = false; + try { + await eng.search('apple', { mode: 'substring', signal: ac.signal }); + } catch (e) { + aborted = (e as { name?: string }).name === 'AbortError'; + } + ok('abort during mapAsync -> AbortError (epoch discard)', aborted); + } finally { + try { + sb.mapAsync = orig; + } catch { + // Best-effort restore. + } + eng.destroy(); } + } else { eng.destroy(); } } diff --git a/scripts/test-vgpu-mock.ts b/scripts/test-vgpu-mock.ts index b7532f5..1a2df7f 100644 --- a/scripts/test-vgpu-mock.ts +++ b/scripts/test-vgpu-mock.ts @@ -427,16 +427,27 @@ async function runMockTests() { }); console.log(' ✅ Concurrent searches stay isolated with no cross-talk'); - const gpuEngine = (concurrentIndex as any).gpuEngine; - const originalSearch = gpuEngine.search.bind(gpuEngine); - gpuEngine.search = async () => { throw new Error('injected GPU query failure'); }; + // GPU failure must fall back to CPU results — injected at the REAL + // boundary the engine talks to (the mock device queue submit), the same + // seam an actual driver failure surfaces through (cf. the gpu-execution- + // error fallback contract verified in test-observability.ts). No private + // engine reference is reached into; the submit hook is restored so the + // shared mock device stays healthy for later sections. + const originalSubmit = mockDevice.queue.submit; + (mockDevice.queue as any).submit = () => { + throw new Error('injected GPU query failure'); + }; + let fallbackResult: SearchResponse; try { - const fallbackResult = await concurrentIndex.search('Auth', { mode: 'substring', limit: 20 }); - if (fallbackResult.engine !== 'cpu' || fallbackResult.results.length === 0) { - throw new Error('A GPU query failure must fall back to CPU results'); - } + fallbackResult = await concurrentIndex.search('Auth', { mode: 'substring', limit: 20 }); } finally { - gpuEngine.search = originalSearch; + mockDevice.queue.submit = originalSubmit; + } + if (fallbackResult.engine !== 'cpu' || fallbackResult.results.length === 0) { + throw new Error('A GPU query failure must fall back to CPU results'); + } + if (fallbackResult.fallbackReason !== 'gpu-execution-error') { + throw new Error(`GPU failure fallback must report gpu-execution-error, got: ${fallbackResult.fallbackReason}`); } concurrentIndex.destroy(); console.log(' ✅ GPU query errors fall back to CPU results'); @@ -669,10 +680,14 @@ async function runMockTests() { // Mock has no clearBuffer: search must still succeed via writeBuffer fallback. await abortEngine.search('x', { mode: 'substring' }); // Forged maxDim 0 must not hang (clamped to 65535 → fast path). + // The hostile device is injected through the public init() seam on a + // dedicated engine instead of mutating the private device reference. const zeroDimDevice = { ...mockDevice, limits: { ...(mockDevice as any).limits, maxComputeWorkgroupsPerDimension: 0 } } as unknown as GPUDevice; - (abortEngine as any).device = zeroDimDevice; - await abortEngine.search('x', { mode: 'substring' }); - (abortEngine as any).device = mockDevice; + const zeroDimEngine = new WebGPUEngine(); + await zeroDimEngine.init(zeroDimDevice); + await zeroDimEngine.loadDataset(['x']); + await zeroDimEngine.search('x', { mode: 'substring' }); + zeroDimEngine.destroy(); // Re-init must not leak/throw. await abortEngine.init(mockDevice); await abortEngine.loadDataset(['re', 'init']);