albex 0.1.0 → 0.6.0

package/src/errors.ts

@@ -0,0 +1,76 @@
+/**
+ * Albex error hierarchy.
+ *
+ * All errors thrown by Albex extend `AlbexError`, which itself extends
+ * `Error`. This lets callers `catch` on the base class to handle every Albex
+ * failure, or on a specific subclass when they want to react differently to,
+ * say, "wrong file type" versus "PDF parse error".
+ *
+ * The strings are intentionally short and machine-friendly so they can be
+ * logged or mapped to user-facing messages without parsing the `.message`.
+ */
+
+export class AlbexError extends Error {
+  /** Discriminator for runtime branching when the prototype chain is lost
+   *  (e.g. after structuredClone across a Worker boundary). */
+  readonly kind: string;
+  constructor(kind: string, message: string) {
+    super(message);
+    this.name = 'AlbexError';
+    this.kind = kind;
+  }
+}
+
+/** Thrown when `init()` was not called or the WASM module failed to load. */
+export class AlbexInitError extends AlbexError {
+  constructor(message: string) {
+    super('init', message);
+    this.name = 'AlbexInitError';
+  }
+}
+
+/** Thrown when a file's extension is not in the supported list. */
+export class AlbexUnsupportedFormatError extends AlbexError {
+  /** The lowercase extension that was rejected, without leading dot. */
+  readonly ext: string;
+  constructor(ext: string) {
+    super('unsupported_format', `Unsupported format: .${ext}`);
+    this.name = 'AlbexUnsupportedFormatError';
+    this.ext = ext;
+  }
+}
+
+/** Thrown when a document parser fails (corrupt ZIP, malformed PDF, etc.). */
+export class AlbexParseError extends AlbexError {
+  /** The format that failed (`'pdf'`, `'docx'`, …). */
+  readonly format: string;
+  constructor(format: string, message: string) {
+    super('parse', message);
+    this.name = 'AlbexParseError';
+    this.format = format;
+  }
+}
+
+/**
+ * Thrown when an indexing operation does not fit: either the scratchpad was
+ * too small for a single write, or one of the engine's pools (chunks, text,
+ * documents, names) ran out of room mid-document. Before 0.6.0 the latter was
+ * silent — the corpus was truncated with no signal.
+ *
+ * `limit` names which pool overflowed (or `'scratchpad'`), so callers can
+ * branch — e.g. start a fresh shard, `compact()`, or surface "library full".
+ * When a capacity error is raised during `indexFile`, the engine may hold a
+ * partially-indexed copy of the offending document; treat the index as full
+ * and stop adding.
+ */
+export type AlbexCapacityLimit = 'chunks' | 'text' | 'docs' | 'names' | 'scratchpad';
+
+export class AlbexCapacityError extends AlbexError {
+  /** Which pool overflowed. Undefined for older call sites that didn't set it. */
+  readonly limit?: AlbexCapacityLimit;
+  constructor(message: string, limit?: AlbexCapacityLimit) {
+    super('capacity', message);
+    this.name = 'AlbexCapacityError';
+    if (limit) this.limit = limit;
+  }
+}

package/src/gpu/bloom-runtime.ts

@@ -0,0 +1,229 @@
+/**
+ * WebGPU Bloom pre-filter accelerator.
+ *
+ * For large corpora (>20 k chunks) the per-chunk `(bloom & pattern) == pattern`
+ * check dominates Albex's CPU search time even though each check is one
+ * AND and one compare. Running ~100 k of them in a single GPU dispatch
+ * cuts that cost by an order of magnitude on modern integrated and
+ * discrete GPUs alike.
+ *
+ * Usage flow (driven by `AlbexEngine` when `shouldUseGpu` returns true):
+ *
+ *   1. After indexing settles, call `uploadChunkBlooms(blooms)` once to
+ *      copy the per-chunk Bloom values to a `GPUBuffer`.
+ *   2. Before each search, call `scan(patternBloom)` to dispatch the
+ *      compute shader. The function returns a packed Uint32Array where
+ *      bit `i` indicates that chunk `i` passed the Bloom test.
+ *   3. The CPU then runs the (much smaller) Bitap pass over only the
+ *      candidates, in WASM as before.
+ *
+ * The class is intentionally one-shot and tolerant: if `init()` fails
+ * because WebGPU is unavailable or refuses a device, callers must check
+ * `available` and fall back to the CPU pre-filter.
+ */
+
+import { BLOOM_SCAN_WGSL } from './bloom-shader.wgsl.js';
+
+// Minimal subset of the WebGPU types we touch. The full typings would
+// require `@types/webgpu`, which is heavy and changes often; declaring
+// only what we use keeps the public surface lean.
+type GPULike = {
+  requestAdapter(): Promise<GPUAdapterLike | null>;
+};
+type GPUAdapterLike = {
+  requestDevice(): Promise<GPUDeviceLike>;
+};
+type GPUDeviceLike = {
+  createBuffer(desc: unknown): GPUBufferLike;
+  createBindGroupLayout(desc: unknown): unknown;
+  createPipelineLayout(desc: unknown): unknown;
+  createShaderModule(desc: { code: string }): unknown;
+  createComputePipeline(desc: unknown): GPUComputePipelineLike;
+  createBindGroup(desc: unknown): unknown;
+  createCommandEncoder(): GPUCommandEncoderLike;
+  queue: { submit(commands: unknown[]): void; writeBuffer(b: GPUBufferLike, off: number, data: ArrayBufferView): void };
+  destroy(): void;
+};
+type GPUBufferLike = {
+  destroy(): void;
+  mapAsync(mode: number): Promise<void>;
+  getMappedRange(): ArrayBuffer;
+  unmap(): void;
+  size: number;
+};
+type GPUComputePipelineLike = {
+  getBindGroupLayout(idx: number): unknown;
+};
+type GPUCommandEncoderLike = {
+  beginComputePass(): GPUComputePassLike;
+  copyBufferToBuffer(src: GPUBufferLike, sOff: number, dst: GPUBufferLike, dOff: number, size: number): void;
+  finish(): unknown;
+};
+type GPUComputePassLike = {
+  setPipeline(p: unknown): void;
+  setBindGroup(idx: number, g: unknown): void;
+  dispatchWorkgroups(x: number, y?: number, z?: number): void;
+  end(): void;
+};
+
+// GPUBufferUsage / GPUMapMode constants are exposed on the global as
+// numeric enums. We mirror only the values we need to avoid pulling typings.
+const USAGE_STORAGE    = 0x80;
+const USAGE_COPY_DST   = 0x08;
+const USAGE_COPY_SRC   = 0x04;
+const USAGE_MAP_READ   = 0x01;
+const MAP_READ         = 0x01;
+
+export class BloomGpu {
+  available = false;
+  private _device: GPUDeviceLike | null = null;
+  private _pipeline: GPUComputePipelineLike | null = null;
+  private _chunkBuf: GPUBufferLike | null = null;
+  private _paramBuf: GPUBufferLike | null = null;
+  private _passBuf:  GPUBufferLike | null = null;
+  private _readBuf:  GPUBufferLike | null = null;
+  private _chunkCount = 0;
+
+  /** Try to acquire a WebGPU device. Sets `available = true` on success. */
+  async init(): Promise<boolean> {
+    try {
+      const gpu = (globalThis as unknown as { navigator?: { gpu?: GPULike } }).navigator?.gpu;
+      if (!gpu) return false;
+      const adapter = await gpu.requestAdapter();
+      if (!adapter) return false;
+      const device = await adapter.requestDevice();
+
+      const module = device.createShaderModule({ code: BLOOM_SCAN_WGSL });
+      this._pipeline = device.createComputePipeline({
+        layout: 'auto',
+        compute: { module, entryPoint: 'bloom_scan' },
+      } as unknown as Record<string, unknown>) as GPUComputePipelineLike;
+
+      this._device = device;
+      this.available = true;
+      return true;
+    } catch {
+      this.available = false;
+      return false;
+    }
+  }
+
+  /**
+   * Upload the per-chunk Bloom values to a fresh GPU buffer. Resizes if
+   * the chunk count changed since last upload. Pass a `Uint32Array` where
+   * each chunk occupies two consecutive entries (low half then high half).
+   */
+  uploadChunkBlooms(blooms: Uint32Array, chunkCount: number): void {
+    if (!this._device) throw new Error('BloomGpu not initialised');
+    if (this._chunkBuf && this._chunkCount !== chunkCount) {
+      this._chunkBuf.destroy();
+      this._chunkBuf = null;
+    }
+    if (!this._chunkBuf) {
+      this._chunkBuf = this._device.createBuffer({
+        size: blooms.byteLength,
+        usage: USAGE_STORAGE | USAGE_COPY_DST,
+      } as unknown as Record<string, unknown>);
+    }
+    this._device.queue.writeBuffer(this._chunkBuf, 0, blooms);
+    this._chunkCount = chunkCount;
+  }
+
+  /**
+   * Dispatch the Bloom scan and read back the packed pass-bitset.
+   *
+   * Returns a `Uint32Array` of length `ceil(chunkCount / 32)`. Each set
+   * bit indicates a chunk that passed the Bloom pre-filter.
+   */
+  async scan(patternLo: number, patternHi: number): Promise<Uint32Array> {
+    if (!this._device || !this._pipeline || !this._chunkBuf) {
+      throw new Error('BloomGpu not ready');
+    }
+    const device = this._device;
+    const passWords = Math.ceil(this._chunkCount / 32);
+
+    if (!this._paramBuf) {
+      this._paramBuf = device.createBuffer({
+        size: 16,
+        usage: USAGE_STORAGE | USAGE_COPY_DST,
+      } as unknown as Record<string, unknown>);
+    }
+    const paramBytes = new Uint32Array([patternLo >>> 0, patternHi >>> 0, this._chunkCount, 0]);
+    device.queue.writeBuffer(this._paramBuf, 0, paramBytes);
+
+    if (!this._passBuf || this._passBuf.size < passWords * 4) {
+      this._passBuf?.destroy();
+      this._readBuf?.destroy();
+      this._passBuf = device.createBuffer({
+        size: Math.max(4, passWords * 4),
+        usage: USAGE_STORAGE | USAGE_COPY_SRC,
+      } as unknown as Record<string, unknown>);
+      this._readBuf = device.createBuffer({
+        size: Math.max(4, passWords * 4),
+        usage: USAGE_MAP_READ | USAGE_COPY_DST,
+      } as unknown as Record<string, unknown>);
+    }
+    // Zero the pass buffer between dispatches — writeBuffer with empty
+    // suffices since atomics set bits cumulatively.
+    device.queue.writeBuffer(this._passBuf, 0, new Uint32Array(passWords));
+
+    const layout = this._pipeline.getBindGroupLayout(0);
+    const bindGroup = device.createBindGroup({
+      layout,
+      entries: [
+        { binding: 0, resource: { buffer: this._chunkBuf } },
+        { binding: 1, resource: { buffer: this._paramBuf } },
+        { binding: 2, resource: { buffer: this._passBuf  } },
+      ],
+    } as unknown as Record<string, unknown>);
+
+    const encoder = device.createCommandEncoder();
+    const pass = encoder.beginComputePass();
+    pass.setPipeline(this._pipeline);
+    pass.setBindGroup(0, bindGroup);
+    pass.dispatchWorkgroups(Math.ceil(this._chunkCount / 64));
+    pass.end();
+    encoder.copyBufferToBuffer(this._passBuf, 0, this._readBuf!, 0, passWords * 4);
+    device.queue.submit([encoder.finish()]);
+
+    await this._readBuf!.mapAsync(MAP_READ);
+    const out = new Uint32Array(this._readBuf!.getMappedRange().slice(0));
+    this._readBuf!.unmap();
+    return out;
+  }
+
+  destroy(): void {
+    this._chunkBuf?.destroy();
+    this._paramBuf?.destroy();
+    this._passBuf?.destroy();
+    this._readBuf?.destroy();
+    this._device?.destroy();
+    this._device = null;
+    this._chunkBuf = null;
+    this._paramBuf = null;
+    this._passBuf  = null;
+    this._readBuf  = null;
+    this.available = false;
+  }
+
+  /** TC39 explicit-resource-management alias of `destroy()`. */
+  [Symbol.dispose](): void { this.destroy(); }
+}
+
+/**
+ * Convenience: take a flat `Uint8Array` view of a WASM `chunks[]` array
+ * (32 bytes per Chunk) and extract just the 8-byte Bloom field at offset 0
+ * into a `Uint32Array` packed `[lo, hi, lo, hi, …]`.
+ *
+ * Used by the engine when uploading to the GPU buffer.
+ */
+export function packBloomsFromChunks(chunkBytes: Uint8Array, chunkCount: number): Uint32Array {
+  const out = new Uint32Array(chunkCount * 2);
+  const view = new DataView(chunkBytes.buffer, chunkBytes.byteOffset, chunkBytes.byteLength);
+  for (let i = 0; i < chunkCount; i++) {
+    // Chunk struct: bloom: u64 at offset 0 within a 32-byte record.
+    out[i * 2]     = view.getUint32(i * 32 + 0, true); // low
+    out[i * 2 + 1] = view.getUint32(i * 32 + 4, true); // high
+  }
+  return out;
+}

package/src/gpu/bloom-shader.wgsl.ts

@@ -0,0 +1,48 @@
+/**
+ * WGSL compute shader for batched Bloom-filter pre-checks.
+ *
+ * Inputs (storage buffers):
+ *   blooms     : array<u32> — 2 u32 per chunk encoding the chunk's u64 Bloom.
+ *                Index 2*i = low half, 2*i+1 = high half.
+ *   params     : { pattern_lo: u32, pattern_hi: u32, chunk_count: u32 }
+ *
+ * Outputs:
+ *   passes     : array<atomic<u32>> — bit i (in passes[i / 32]) set iff
+ *                chunks[i] might contain the pattern.
+ *
+ * Workgroup size of 64 matches typical occupancy on Apple GPUs and
+ * Snapdragon; larger workgroups gain little for this memory-bound op.
+ *
+ * Kept in TS so the bundler ships it inline — no extra `.wgsl` asset fetch.
+ */
+export const BLOOM_SCAN_WGSL = /* wgsl */ `
+struct Params {
+  pattern_lo:  u32,
+  pattern_hi:  u32,
+  chunk_count: u32,
+  _pad:        u32,
+};
+
+@group(0) @binding(0) var<storage, read>          blooms : array<u32>;
+@group(0) @binding(1) var<storage, read>          params : Params;
+@group(0) @binding(2) var<storage, read_write>    passes : array<atomic<u32>>;
+
+@compute @workgroup_size(64)
+fn bloom_scan(@builtin(global_invocation_id) gid: vec3<u32>) {
+  let i = gid.x;
+  if (i >= params.chunk_count) { return; }
+
+  let bloom_lo = blooms[i * 2u];
+  let bloom_hi = blooms[i * 2u + 1u];
+
+  // (bloom & pattern) == pattern  iff  pattern is a subset of bloom.
+  let match_lo = (bloom_lo & params.pattern_lo) == params.pattern_lo;
+  let match_hi = (bloom_hi & params.pattern_hi) == params.pattern_hi;
+  if (!(match_lo && match_hi)) { return; }
+
+  // Mark bit i in the packed result. Each u32 holds 32 chunk results.
+  let word = i / 32u;
+  let bit  = i % 32u;
+  atomicOr(&passes[word], 1u << bit);
+}
+`;

package/src/persistence.ts

@@ -0,0 +1,175 @@
+/**
+ * Albex persistence layer.
+ *
+ * Backends:
+ *   - OPFS  (preferred, available in Chrome 102+, Safari 15.2+, Firefox 111+)
+ *   - IndexedDB (universal fallback)
+ *
+ * Both store a single binary blob per snapshot name. The format is opaque to
+ * this layer — see the snapshot header documented in `wasm/src/lib.rs`.
+ *
+ * API surface is intentionally minimal:
+ *   savePersisted(name, bytes)  → void
+ *   loadPersisted(name)         → Uint8Array | null
+ *   deletePersisted(name)       → void
+ *   listPersisted()             → string[]
+ */
+
+const DB_NAME    = 'albex';
+const STORE_NAME = 'snapshots';
+const OPFS_DIR   = 'albex-snapshots';
+
+// ─────────────────────────────────────────────────────────────────────────────
+// Backend detection
+// ─────────────────────────────────────────────────────────────────────────────
+
+function hasOpfs(): boolean {
+  // navigator.storage.getDirectory is the OPFS entry point.
+  return typeof navigator !== 'undefined'
+    && typeof navigator.storage?.getDirectory === 'function';
+}
+
+function hasIdb(): boolean {
+  return typeof indexedDB !== 'undefined';
+}
+
+// ─────────────────────────────────────────────────────────────────────────────
+// OPFS backend
+// ─────────────────────────────────────────────────────────────────────────────
+
+async function opfsDir(): Promise<FileSystemDirectoryHandle> {
+  const root = await navigator.storage.getDirectory();
+  return root.getDirectoryHandle(OPFS_DIR, { create: true });
+}
+
+async function opfsSave(name: string, bytes: Uint8Array): Promise<void> {
+  const dir  = await opfsDir();
+  const file = await dir.getFileHandle(name, { create: true });
+  const w    = await file.createWritable();
+  // FileSystemWritableFileStream.write demands ArrayBuffer specifically (not
+  // ArrayBufferLike); copy through a plain Uint8Array to satisfy the type.
+  const plain = new Uint8Array(bytes.byteLength);
+  plain.set(bytes);
+  await w.write(plain);
+  await w.close();
+}
+
+async function opfsLoad(name: string): Promise<Uint8Array | null> {
+  try {
+    const dir  = await opfsDir();
+    const file = await dir.getFileHandle(name);
+    const f    = await file.getFile();
+    return new Uint8Array(await f.arrayBuffer());
+  } catch {
+    return null;
+  }
+}
+
+async function opfsDelete(name: string): Promise<void> {
+  try {
+    const dir = await opfsDir();
+    await dir.removeEntry(name);
+  } catch { /* not found */ }
+}
+
+async function opfsList(): Promise<string[]> {
+  const dir = await opfsDir();
+  const out: string[] = [];
+  // @ts-expect-error: async iterator on FileSystemDirectoryHandle (Web Spec)
+  for await (const [name] of dir.entries()) out.push(name as string);
+  return out;
+}
+
+// ─────────────────────────────────────────────────────────────────────────────
+// IndexedDB backend
+// ─────────────────────────────────────────────────────────────────────────────
+
+function idbOpen(): Promise<IDBDatabase> {
+  return new Promise((resolve, reject) => {
+    const req = indexedDB.open(DB_NAME, 1);
+    req.onupgradeneeded = () => {
+      const db = req.result;
+      if (!db.objectStoreNames.contains(STORE_NAME)) {
+        db.createObjectStore(STORE_NAME);
+      }
+    };
+    req.onsuccess = () => resolve(req.result);
+    req.onerror   = () => reject(req.error);
+  });
+}
+
+function idbReq<T>(req: IDBRequest<T>): Promise<T> {
+  return new Promise((resolve, reject) => {
+    req.onsuccess = () => resolve(req.result);
+    req.onerror   = () => reject(req.error);
+  });
+}
+
+async function idbSave(name: string, bytes: Uint8Array): Promise<void> {
+  const db = await idbOpen();
+  const tx = db.transaction(STORE_NAME, 'readwrite');
+  // structured-clone of a TypedArray copies the buffer once. That's the price
+  // of the IDB fallback; OPFS path avoids it.
+  tx.objectStore(STORE_NAME).put(bytes, name);
+  await new Promise<void>((res, rej) => {
+    tx.oncomplete = () => res();
+    tx.onerror    = () => rej(tx.error);
+  });
+  db.close();
+}
+
+async function idbLoad(name: string): Promise<Uint8Array | null> {
+  const db = await idbOpen();
+  const tx = db.transaction(STORE_NAME, 'readonly');
+  const got = await idbReq<unknown>(tx.objectStore(STORE_NAME).get(name));
+  db.close();
+  if (got instanceof Uint8Array) return got;
+  if (got instanceof ArrayBuffer) return new Uint8Array(got);
+  return null;
+}
+
+async function idbDelete(name: string): Promise<void> {
+  const db = await idbOpen();
+  const tx = db.transaction(STORE_NAME, 'readwrite');
+  tx.objectStore(STORE_NAME).delete(name);
+  await new Promise<void>((res, rej) => {
+    tx.oncomplete = () => res();
+    tx.onerror    = () => rej(tx.error);
+  });
+  db.close();
+}
+
+async function idbList(): Promise<string[]> {
+  const db = await idbOpen();
+  const tx = db.transaction(STORE_NAME, 'readonly');
+  const keys = await idbReq<IDBValidKey[]>(tx.objectStore(STORE_NAME).getAllKeys());
+  db.close();
+  return keys.map(k => String(k));
+}
+
+// ─────────────────────────────────────────────────────────────────────────────
+// Backend router
+// ─────────────────────────────────────────────────────────────────────────────
+
+export async function savePersisted(name: string, bytes: Uint8Array): Promise<void> {
+  if (hasOpfs()) return opfsSave(name, bytes);
+  if (hasIdb())  return idbSave(name, bytes);
+  throw new Error('No persistence backend available (OPFS or IndexedDB)');
+}
+
+export async function loadPersisted(name: string): Promise<Uint8Array | null> {
+  if (hasOpfs()) return opfsLoad(name);
+  if (hasIdb())  return idbLoad(name);
+  return null;
+}
+
+export async function deletePersisted(name: string): Promise<void> {
+  if (hasOpfs()) return opfsDelete(name);
+  if (hasIdb())  return idbDelete(name);
+}
+
+export async function listPersisted(): Promise<string[]> {
+  if (hasOpfs()) return opfsList();
+  if (hasIdb())  return idbList();
+  return [];
+}