albex 0.1.0

package/src/albex.ts

@@ -0,0 +1,577 @@
+/**
+ * Albex — local full-text search engine.
+ *
+ * Zero-dependency TypeScript/ESM wrapper around albex_wasm_bg.wasm and
+ * (optionally) albex_pdf.wasm. All text stays in-browser; nothing is sent
+ * to any server.
+ *
+ * @example
+ * ```ts
+ * const engine = new AlbexEngine({ wasmUrl: './wasm/pkg/albex_wasm_bg.wasm' });
+ * await engine.init();
+ * await engine.indexFile(myFile);
+ * const results = engine.search('contrato marco');
+ * ```
+ */
+
+// ─────────────────────────────────────────────────────────────────────────────
+// Public types
+// ─────────────────────────────────────────────────────────────────────────────
+
+export interface AlbexOptions {
+  /** URL to albex_wasm_bg.wasm (required). */
+  wasmUrl: string;
+  /** URL to albex_pdf.wasm. Required only if you call indexFile() with PDFs. */
+  pdfWasmUrl?: string;
+}
+
+export interface IndexedDocument {
+  name: string;
+  ext: string;
+  chunks: number;
+  indexTimeMs: number;
+  textBytes: number;
+}
+
+export interface SearchResult {
+  documentName: string;
+  /** Paragraph index (DOCX/TXT) or page number (PDF, 1-based). */
+  location: number;
+  /** Relevance score 0–1000. */
+  score: number;
+  /** Raw snippet text (original, with accents). */
+  snippet: string;
+  /** Match start byte offset within snippet. */
+  matchStart: number;
+  /** Match end byte offset within snippet (exclusive). */
+  matchEnd: number;
+}
+
+export interface EngineStats {
+  documents: number;
+  chunks: number;
+  textUsed: number;
+  textCapacity: number;
+  wasmMemoryBytes: number;
+}
+
+export interface SearchStats {
+  query: string;
+  timeMs: number;
+  results: number;
+  bloomTested: number;
+  bloomPassed: number;
+  bitapMatched: number;
+}
+
+// ─────────────────────────────────────────────────────────────────────────────
+// Query parsing
+// ─────────────────────────────────────────────────────────────────────────────
+
+type SimpleQuery = { kind: 'simple'; tokens: string[] };
+type PhraseQuery = { kind: 'phrase'; tokens: string[]; raw: string };
+type OrQuery    = { kind: 'or';     branches: string[][] };
+type ParsedQuery = SimpleQuery | PhraseQuery | OrQuery;
+
+function tokenize(q: string): string[] {
+  return q.trim().split(/\s+/).filter(t => t.length > 0);
+}
+
+function parseQuery(q: string): ParsedQuery {
+  const trimmed = q.trim();
+
+  // OR: "term1 | term2" or "phrase one | phrase two"
+  if (trimmed.includes('|')) {
+    const branches = trimmed.split('|')
+      .map(p => tokenize(p.replace(/"/g, '')))
+      .filter(b => b.length > 0);
+    return { kind: 'or', branches };
+  }
+
+  // Phrase: "exact phrase here"
+  const phraseMatch = /^"(.+)"$/.exec(trimmed);
+  if (phraseMatch) {
+    const inner = phraseMatch[1] ?? '';
+    const tokens = tokenize(inner);
+    return { kind: 'phrase', tokens, raw: inner };
+  }
+
+  return { kind: 'simple', tokens: tokenize(trimmed) };
+}
+
+/**
+ * Reconstruct a WASM-compatible query string from parsed tokens.
+ * The WASM engine accepts up to 4 space-separated tokens (AND semantics).
+ */
+function tokensToWasmQuery(tokens: string[]): string {
+  return tokens.slice(0, 4).join(' ');
+}
+
+// ─────────────────────────────────────────────────────────────────────────────
+// Phrase post-filter
+// ─────────────────────────────────────────────────────────────────────────────
+
+/**
+ * Returns true if `snippet` contains the phrase formed by `tokens` in order,
+ * with at most `maxGap` characters between consecutive tokens.
+ * Comparison is case- and accent-insensitive.
+ */
+function containsPhrase(snippet: string, tokens: string[], maxGap = 30): boolean {
+  const norm = (s: string): string =>
+    s.toLowerCase().normalize('NFKD').replace(/[̀-ͯ]/g, '');
+
+  const text = norm(snippet);
+  const normTokens = tokens.map(norm).filter(t => t.length > 0);
+  if (normTokens.length === 0) return true;
+
+  let pos = 0;
+  for (let t = 0; t < normTokens.length; t++) {
+    const tok = normTokens[t] ?? '';
+    if (!tok) continue;
+    const idx = text.indexOf(tok, pos);
+    if (idx === -1) return false;
+    if (t > 0 && idx - pos > maxGap) return false;
+    pos = idx + tok.length;
+  }
+  return true;
+}
+
+// ─────────────────────────────────────────────────────────────────────────────
+// ZIP helpers (DOCX + XLSX)
+// ─────────────────────────────────────────────────────────────────────────────
+
+const _enc = new TextEncoder();
+const _dec = new TextDecoder();
+
+function zipCentralDir(bytes: Uint8Array): { v: DataView; cdOff: number; cdN: number } {
+  const v = new DataView(bytes.buffer, bytes.byteOffset, bytes.byteLength);
+  let p = bytes.length - 22;
+  while (p >= 0 && v.getUint32(p, true) !== 0x06054b50) p--;
+  if (p < 0) throw new Error('Not a ZIP file');
+  return { v, cdOff: v.getUint32(p + 16, true), cdN: v.getUint16(p + 10, true) };
+}
+
+function listZipEntries(bytes: Uint8Array): string[] {
+  const { v, cdOff, cdN } = zipCentralDir(bytes);
+  const names: string[] = [];
+  let cp = cdOff;
+  for (let i = 0; i < cdN; i++) {
+    if (v.getUint32(cp, true) !== 0x02014b50) break;
+    const nl = v.getUint16(cp + 28, true), xl = v.getUint16(cp + 30, true), cl = v.getUint16(cp + 32, true);
+    names.push(_dec.decode(bytes.subarray(cp + 46, cp + 46 + nl)));
+    cp += 46 + nl + xl + cl;
+  }
+  return names;
+}
+
+async function findZipEntry(bytes: Uint8Array, name: string): Promise<Uint8Array> {
+  const { v, cdOff, cdN } = zipCentralDir(bytes);
+  let cp = cdOff;
+  for (let i = 0; i < cdN; i++) {
+    if (v.getUint32(cp, true) !== 0x02014b50) break;
+    const nl = v.getUint16(cp + 28, true), xl = v.getUint16(cp + 30, true), cl = v.getUint16(cp + 32, true);
+    const off = v.getUint32(cp + 42, true), compSz = v.getUint32(cp + 20, true);
+    if (_dec.decode(bytes.subarray(cp + 46, cp + 46 + nl)) === name) {
+      return decompEntry(bytes, v, off, compSz);
+    }
+    cp += 46 + nl + xl + cl;
+  }
+  throw new Error(`Entry "${name}" not found in ZIP`);
+}
+
+async function decompEntry(bytes: Uint8Array, v: DataView, off: number, compSize: number): Promise<Uint8Array> {
+  const meth = v.getUint16(off + 8, true);
+  const nl   = v.getUint16(off + 26, true);
+  const xl   = v.getUint16(off + 28, true);
+  const ds   = off + 30 + nl + xl;
+  // Use central-directory size (reliable even for streaming ZIPs with local-header=0).
+  const cs   = compSize > 0 ? compSize : v.getUint32(off + 18, true);
+  const data = bytes.subarray(ds, ds + cs);
+
+  if (meth === 0) return data;
+  if (meth === 8) {
+    const s = new DecompressionStream('deflate-raw');
+    const w = s.writable.getWriter();
+    const r = s.readable.getReader();
+    // Slice to a plain ArrayBuffer-backed Uint8Array to satisfy strict DOM types.
+    const plain = data.slice();
+    w.write(plain).catch(() => {}); w.close().catch(() => {});
+    const chunks: Uint8Array[] = []; let len = 0;
+    for (;;) { const { done, value } = await r.read(); if (done) break; chunks.push(value); len += value.length; }
+    const out = new Uint8Array(len); let o = 0;
+    for (const c of chunks) { out.set(c, o); o += c.length; }
+    return out;
+  }
+  throw new Error(`Unsupported ZIP compression method ${meth}`);
+}
+
+// ─────────────────────────────────────────────────────────────────────────────
+// WASM memory helpers (internal)
+// ─────────────────────────────────────────────────────────────────────────────
+
+const FEED_SIZE = 32_768; // 32 KB — fits in 64 KB scratchpad
+
+// ─────────────────────────────────────────────────────────────────────────────
+// PDF WASM imports shim
+// ─────────────────────────────────────────────────────────────────────────────
+
+function makePdfWasmImports(getPdfMem: () => WebAssembly.Memory): WebAssembly.Imports {
+  const heap: unknown[] = [];
+  let freeIdx = -1;
+  return {
+    __wbindgen_placeholder__: {
+      __wbindgen_describe: () => {},
+      __wbg_getRandomValues_3f44b700395062e5: (ptr: number, len: number) => {
+        const mem = getPdfMem();
+        crypto.getRandomValues(new Uint8Array(mem.buffer, ptr >>> 0, len >>> 0));
+      },
+      __wbindgen_object_drop_ref: (idx: number) => {
+        heap[idx] = freeIdx; freeIdx = idx;
+      },
+    },
+    __wbindgen_externref_xform__: {
+      __wbindgen_externref_table_grow: (delta: number) => {
+        const old = heap.length;
+        for (let i = 0; i < delta; i++) heap.push(undefined);
+        return old;
+      },
+      __wbindgen_externref_table_set_null: (idx: number) => { heap[idx] = undefined; },
+    },
+  };
+}
+
+// ─────────────────────────────────────────────────────────────────────────────
+// AlbexEngine
+// ─────────────────────────────────────────────────────────────────────────────
+
+export class AlbexEngine {
+  // ── main WASM ──
+  private _wasm!: WebAssembly.Exports;
+  private _mem!: WebAssembly.Memory;
+
+  // ── PDF WASM (lazy) ──
+  private _pdfWasm: WebAssembly.Exports | null = null;
+  private _pdfMem: WebAssembly.Memory | null = null;
+
+  private _docs: IndexedDocument[] = [];
+  private _lastSearch: SearchStats | null = null;
+  private readonly _opts: AlbexOptions;
+
+  constructor(opts: AlbexOptions) {
+    this._opts = opts;
+  }
+
+  /** Load and initialise the main WASM module. Must be called before any other method. */
+  async init(): Promise<void> {
+    const res = await fetch(this._opts.wasmUrl);
+    if (!res.ok) throw new Error(`Failed to fetch WASM: ${res.status}`);
+    const { instance } = await WebAssembly.instantiateStreaming(res, {});
+    this._wasm = instance.exports;
+    this._mem  = instance.exports.memory as WebAssembly.Memory;
+    (this._wasm.init as Function)();
+  }
+
+  // ── Internal helpers ──────────────────────────────────────────────────────
+
+  private _u8(off: number, n: number): Uint8Array {
+    return new Uint8Array(this._mem.buffer, off, n);
+  }
+
+  private _writePad(b: Uint8Array): number {
+    const ptr = (this._wasm.getBuffer as Function)(b.length) as number;
+    if (!ptr) throw new Error('Scratchpad too small for this chunk');
+    this._u8(ptr, b.length).set(b);
+    return ptr;
+  }
+
+  private _writeStr(s: string): number {
+    const b = _enc.encode(s);
+    this._writePad(b);
+    return b.length;
+  }
+
+  private _readPad(n: number): string {
+    const ptr = (this._wasm.getBuffer as Function)(0) as number;
+    return _dec.decode(this._u8(ptr, n));
+  }
+
+  private _feedText(text: string): void {
+    const b = _enc.encode(text);
+    for (let i = 0; i < b.length; i += FEED_SIZE) {
+      const c = b.subarray(i, i + FEED_SIZE);
+      this._writePad(c);
+      (this._wasm.feedText as Function)(c.length);
+    }
+  }
+
+  private _feedXmlBytes(xml: Uint8Array, fn: 'feedXmlBytes' | 'feedXlsxBytes'): void {
+    for (let i = 0; i < xml.length; i += FEED_SIZE) {
+      const c = xml.subarray(i, i + FEED_SIZE);
+      this._writePad(c);
+      (this._wasm[fn] as Function)(c.length);
+    }
+  }
+
+  // ── PDF WASM (lazy load) ─────────────────────────────────────────────────
+
+  private async _ensurePdfWasm(): Promise<void> {
+    if (this._pdfWasm) return;
+    if (!this._opts.pdfWasmUrl) throw new Error('pdfWasmUrl not set in AlbexOptions');
+    const res = await fetch(this._opts.pdfWasmUrl);
+    if (!res.ok) throw new Error(`Failed to fetch PDF WASM: ${res.status}`);
+    const imports = makePdfWasmImports(() => this._pdfMem!);
+    const { instance } = await WebAssembly.instantiateStreaming(res, imports);
+    this._pdfWasm = instance.exports;
+    this._pdfMem  = instance.exports.memory as WebAssembly.Memory;
+  }
+
+  // ── Indexers ──────────────────────────────────────────────────────────────
+
+  private async _indexDocx(file: File): Promise<number> {
+    const bytes = new Uint8Array(await file.arrayBuffer());
+    const xml   = await findZipEntry(bytes, 'word/document.xml');
+    (this._wasm.setDocumentName as Function)(this._writeStr(file.name));
+    (this._wasm.beginDocument as Function)();
+    this._feedXmlBytes(xml, 'feedXmlBytes');
+    return (this._wasm.endDocument as Function)() as number;
+  }
+
+  private async _indexXlsx(file: File): Promise<number> {
+    const bytes = new Uint8Array(await file.arrayBuffer());
+    (this._wasm.setDocumentName as Function)(this._writeStr(file.name));
+    (this._wasm.beginXlsx as Function)();
+
+    try {
+      const xml = await findZipEntry(bytes, 'xl/sharedStrings.xml');
+      this._feedXmlBytes(xml, 'feedXlsxBytes');
+    } catch { /* workbook may have no shared strings */ }
+
+    const sheets = listZipEntries(bytes).filter(n => /^xl\/worksheets\/sheet\d+\.xml$/.test(n));
+    for (const name of sheets) {
+      try {
+        const xml = await findZipEntry(bytes, name);
+        this._feedXmlBytes(xml, 'feedXlsxBytes');
+      } catch { /* skip corrupt/missing sheet */ }
+    }
+
+    return (this._wasm.endDocument as Function)() as number;
+  }
+
+  private async _indexPdf(file: File): Promise<number> {
+    await this._ensurePdfWasm();
+    const pw   = this._pdfWasm!;
+    const pm   = this._pdfMem!;
+    const bytes = new Uint8Array(await file.arrayBuffer());
+
+    const inPtr    = (pw.allocInput as Function)(bytes.length) as number;
+    new Uint8Array(pm.buffer, inPtr, bytes.length).set(bytes);
+    const pageCount = (pw.extractPdf as Function)(bytes.length) as number;
+
+    (this._wasm.setDocumentName as Function)(this._writeStr(file.name));
+    (this._wasm.beginDocument as Function)();
+
+    if (pageCount === -2) {
+      // Image-only PDF — register doc with zero chunks.
+      return (this._wasm.endDocument as Function)() as number;
+    }
+    if (pageCount < 0) {
+      const errLen = (pw.getErrorLen as Function)() as number;
+      const errPtr = (pw.getErrorPtr as Function)() as number;
+      const msg = errLen > 0
+        ? new TextDecoder().decode(new Uint8Array(pm.buffer, errPtr, errLen))
+        : 'PDF parse error';
+      throw new Error(msg);
+    }
+
+    for (let p = 0; p < pageCount; p++) {
+      const len = (pw.getPageLen as Function)(p) as number;
+      if (!len) continue;
+      const text = new TextDecoder('utf-8').decode(
+        new Uint8Array(pm.buffer, (pw.getPagePtr as Function)(p) as number, len)
+      );
+      this._feedText(text);
+      (this._wasm.flushParagraph as Function)();
+    }
+
+    return (this._wasm.endDocument as Function)() as number;
+  }
+
+  private async _indexTxt(file: File): Promise<number> {
+    const text = await file.text();
+    (this._wasm.setDocumentName as Function)(this._writeStr(file.name));
+    (this._wasm.beginDocument as Function)();
+    for (const para of text.split(/\n{2,}/)) {
+      const l = para.replace(/\n/g, ' ').trim();
+      if (l) { this._feedText(l); (this._wasm.flushParagraph as Function)(); }
+    }
+    return (this._wasm.endDocument as Function)() as number;
+  }
+
+  private async _indexXml(file: File): Promise<number> {
+    const plain = (await file.text())
+      .replace(/<[^]*?>/g, '\n')
+      .replace(/&amp;/g, '&').replace(/&lt;/g, '<').replace(/&gt;/g, '>')
+      .replace(/&quot;/g, '"').replace(/&apos;/g, "'")
+      .replace(/[ \t]+/g, ' ').trim();
+    (this._wasm.setDocumentName as Function)(this._writeStr(file.name));
+    (this._wasm.beginDocument as Function)();
+    for (const seg of plain.split(/\n{2,}/)) {
+      const l = seg.replace(/\n/g, ' ').trim();
+      if (l) { this._feedText(l); (this._wasm.flushParagraph as Function)(); }
+    }
+    return (this._wasm.endDocument as Function)() as number;
+  }
+
+  private static readonly _INDEXERS: Record<string, (engine: AlbexEngine, file: File) => Promise<number>> = {
+    docx: (e, f) => e._indexDocx(f),
+    xlsx: (e, f) => e._indexXlsx(f),
+    pdf:  (e, f) => e._indexPdf(f),
+    txt:  (e, f) => e._indexTxt(f),
+    xml:  (e, f) => e._indexXml(f),
+  };
+
+  // ── Public API ────────────────────────────────────────────────────────────
+
+  /**
+   * Index a file. Supported formats: DOCX, XLSX, PDF, TXT, XML.
+   * Throws for unsupported formats or parse errors.
+   */
+  async indexFile(file: File): Promise<IndexedDocument> {
+    const ext = file.name.split('.').pop()?.toLowerCase() ?? '';
+    const indexer = AlbexEngine._INDEXERS[ext];
+    if (!indexer) throw new Error(`Unsupported format: .${ext}`);
+
+    const t0     = performance.now();
+    const textPre = (this._wasm.getTextUsed as Function)() as number;
+    const chunks  = await indexer(this, file);
+    const doc: IndexedDocument = {
+      name:        file.name,
+      ext,
+      chunks,
+      indexTimeMs: performance.now() - t0,
+      textBytes:   ((this._wasm.getTextUsed as Function)() as number) - textPre,
+    };
+    this._docs.push(doc);
+    return doc;
+  }
+
+  /**
+   * Search the index. Supports:
+   * - Simple queries:  `contrato`  (AND of tokens, accent-insensitive)
+   * - Phrase queries:  `"contrato marco"`  (must appear as phrase)
+   * - OR queries:      `contrato | acuerdo`  (union of two searches)
+   */
+  search(query: string): SearchResult[] {
+    const parsed = parseQuery(query);
+
+    if (parsed.kind === 'or') {
+      return this._searchOr(parsed.branches, query);
+    }
+
+    const results = this._runSearch(tokensToWasmQuery(parsed.tokens), query);
+
+    if (parsed.kind === 'phrase') {
+      return results.filter(r => containsPhrase(r.snippet, parsed.tokens));
+    }
+
+    return results;
+  }
+
+  private _searchOr(branches: string[][], rawQuery: string): SearchResult[] {
+    const seen = new Set<string>();
+    const all: SearchResult[] = [];
+
+    for (const tokens of branches) {
+      const q = tokensToWasmQuery(tokens);
+      if (!q) continue;
+      const results = this._runSearch(q, rawQuery);
+      for (const r of results) {
+        const key = `${r.documentName}:${r.location}:${r.matchStart}`;
+        if (!seen.has(key)) { seen.add(key); all.push(r); }
+      }
+    }
+
+    // Re-rank the merged list by score descending.
+    all.sort((a, b) => b.score - a.score);
+    return all;
+  }
+
+  private _runSearch(wasmQuery: string, displayQuery: string): SearchResult[] {
+    const ql  = this._writeStr(wasmQuery);
+    (this._wasm.setPattern as Function)(ql);
+
+    const t0    = performance.now();
+    const count = (this._wasm.search as Function)() as number;
+    const ms    = performance.now() - t0;
+
+    this._lastSearch = {
+      query:        displayQuery,
+      timeMs:       ms,
+      results:      count,
+      bloomTested:  (this._wasm.getStatBloomTested as Function)() as number,
+      bloomPassed:  (this._wasm.getStatBloomPassed as Function)() as number,
+      bitapMatched: (this._wasm.getStatBitapMatched as Function)() as number,
+    };
+
+    const results: SearchResult[] = [];
+    for (let i = 0; i < count; i++) {
+      const score    = (this._wasm.getResultScore    as Function)(i) as number;
+      const location = (this._wasm.getResultLocation as Function)(i) as number;
+      const matchStart = (this._wasm.getResultStart  as Function)(i) as number;
+      const matchEnd   = (this._wasm.getResultEnd    as Function)(i) as number;
+      const nl         = (this._wasm.getResultDocName as Function)(i) as number;
+      const name       = nl > 0 ? this._readPad(nl) : '?';
+      const sl         = (this._wasm.getSnippet      as Function)(i) as number;
+      const snippet    = sl > 0 ? this._readPad(sl) : '';
+
+      results.push({ documentName: name, location, score, snippet, matchStart, matchEnd });
+    }
+    return results;
+  }
+
+  /** Returns current engine statistics. */
+  getStats(): EngineStats {
+    return {
+      documents:      this._docs.length,
+      chunks:         (this._wasm.getChunkCount    as Function)() as number,
+      textUsed:       (this._wasm.getTextUsed      as Function)() as number,
+      textCapacity:   (this._wasm.getTextCapacity  as Function)() as number,
+      wasmMemoryBytes: this._mem.buffer.byteLength,
+    };
+  }
+
+  /** Returns stats from the most recent search, or null. */
+  getLastSearchStats(): SearchStats | null {
+    return this._lastSearch;
+  }
+
+  /** Returns the list of indexed documents. */
+  get documents(): readonly IndexedDocument[] {
+    return this._docs;
+  }
+
+  /** Supported file extensions. */
+  static get supportedExtensions(): string[] {
+    return Object.keys(AlbexEngine._INDEXERS);
+  }
+
+  /** Configure search sensitivity. */
+  setMaxErrors(errors: 0 | 1 | 2 | 3): void {
+    (this._wasm.setMaxErrors as Function)(errors);
+  }
+
+  setThreshold(threshold: number): void {
+    (this._wasm.setThreshold as Function)(Math.max(0, Math.min(1000, threshold)));
+  }
+
+  setMaxResults(max: number): void {
+    (this._wasm.setMaxResults as Function)(Math.max(1, Math.min(200, max)));
+  }
+
+  /** Full reset — clears all indexed documents and chunks. */
+  reset(): void {
+    (this._wasm.init as Function)();
+    this._docs = [];
+    this._lastSearch = null;
+  }
+}