md-semantic-search 2.0.0 → 2.2.0

package/src/hnsw.ts

@@ -11,6 +11,8 @@
  * is what lets the incremental watch rebuild reuse a persisted graph.
  */
 
+import { createWasmDotEngine } from './wasm-dot.js';
+
 /** Deterministic PRNG for level assignment (same family as ivf.ts). */
 function mulberry32(seed: number): () => number {
   let s = seed >>> 0;
@@ -63,6 +65,28 @@ export interface BuildHnswOptions {
    */
   distance: (a: Float32Array | Int8Array | number[], b: number) => number;
   vectors: Array<Float32Array | Int8Array | number[]>;
+  /**
+   * Issue #228: 'auto' routes the build's distance work through the WASM
+   * SIMD dot-batch kernel (src/wasm-dot.ts) when the vectors are a uniform
+   * Float32Array matrix and the runtime supports SIMD — ~6× faster dots on a
+   * 32-node batch, ~3.4× faster full build (RESEARCH.md §17). ONLY valid when
+   * the intended distance is exactly `1 − dot(a, vectors[b])`; any custom
+   * distance (int8 norms, asymmetric variants) must keep the default
+   * 'closure' engine. Silent fallback to 'closure' when 'auto' cannot apply.
+   */
+  engine?: 'auto' | 'closure';
+}
+
+/**
+ * Distance plumbing for the build (issue #228). The closure engine wraps the
+ * caller-supplied `distance`; the wasm engine replaces BOTH pointwise and
+ * batch distances with kernel calls so every distance on the chain uses the
+ * same SIMD accumulation order (self-consistent → deterministic sidecar).
+ */
+export interface HnswBuildEngine {
+  dist(aIdx: number, bIdx: number): number;
+  /** Distances from vectors[aIdx] to each node, same order — null = pointwise. */
+  batch?(aIdx: number, nodes: number[]): Float32Array | null;
 }
 
 export const DEFAULT_HNSW_M = 16;
@@ -82,10 +106,70 @@ export function buildHnsw(
     neighbors: [],
     deleted: new Set(),
   };
-  for (let i = 0; i < vectors.length; i++) insertHnsw(hnsw, i, options.distance, options.vectors);
+  let engine: HnswBuildEngine | undefined;
+  if (options.engine === 'auto') {
+    engine = createWasmBuildEngine(vectors) ?? undefined;
+    // undefined → SIMD unavailable or non-uniform vectors: closure engine, silently.
+  }
+  for (let i = 0; i < vectors.length; i++) insertHnsw(hnsw, i, options.distance, options.vectors, engine);
   return hnsw;
 }
 
+/**
+ * Issue #228: wasm-backed build engine for uniform Float32Array matrices.
+ * Returns null whenever the fast path cannot apply — the caller then keeps
+ * the caller-supplied closure (identical semantics, JS arithmetic).
+ */
+function createWasmBuildEngine(vectors: Array<Float32Array | Int8Array | number[]>): HnswBuildEngine | null {
+  for (const v of vectors) {
+    if (!(v instanceof Float32Array)) return null;
+  }
+  const wasmEngine = createWasmDotEngine(vectors as Array<Float32Array>);
+  if (!wasmEngine) return null;
+  // Hybrid distance plumbing (measured, issue #228): BATCHES (a whole node's
+  // neighbor list — the 85% cost center) go through the SIMD kernel; pointwise
+  // calls (greedy descent, singleton scans, heuristic pairwise) stay on the
+  // plain JS dot — a 1-row wasm call pays call overhead exceeding the work.
+  // Both paths are IEEE f32 with fixed orders, so the build stays
+  // deterministic; recall is insensitive to the last-ulp difference (measured
+  // identical across engines in bench/ann-recall.mjs and the #228 tests).
+  const f32Vectors = vectors as Array<Float32Array>;
+  return {
+    dist: (a, b) => 1 - jsDotF32(f32Vectors[a], f32Vectors[b]),
+    batch: (aIdx, nodes) => {
+      const dots = wasmEngine.batch(aIdx, nodes);
+      const out = new Float32Array(dots.length);
+      for (let j = 0; j < dots.length; j++) out[j] = 1 - dots[j];
+      return out;
+    },
+  };
+}
+
+function jsDotF32(a: Float32Array, b: Float32Array): number {
+  let s = 0;
+  const n = a.length;
+  const limit = n - 7;
+  let i = 0;
+  let s0 = 0, s1 = 0, s2 = 0, s3 = 0, s4 = 0, s5 = 0, s6 = 0, s7 = 0;
+  while (i < limit) {
+    s0 += a[i] * b[i];
+    s1 += a[i + 1] * b[i + 1];
+    s2 += a[i + 2] * b[i + 2];
+    s3 += a[i + 3] * b[i + 3];
+    s4 += a[i + 4] * b[i + 4];
+    s5 += a[i + 5] * b[i + 5];
+    s6 += a[i + 6] * b[i + 6];
+    s7 += a[i + 7] * b[i + 7];
+    i += 8;
+  }
+  s = s0 + s1 + (s2 + s3) + (s4 + s5) + (s6 + s7);
+  while (i < n) {
+    s += a[i] * b[i];
+    i++;
+  }
+  return s;
+}
+
 /**
  * Insert node `idx` (vector index) into the graph. Incremental — new indices
  * attach to the existing structure without a rebuild; vectors for nodes NOT
@@ -96,7 +180,8 @@ export function insertHnsw(
   hnsw: HnswIndex,
   idx: number,
   distance: (a: Float32Array | Int8Array | number[], b: number) => number,
-  vectors: Array<Float32Array | Int8Array | number[]>
+  vectors: Array<Float32Array | Int8Array | number[]>,
+  engine?: HnswBuildEngine
 ): void {
   const vec = vectors[idx];
   if (!vec) throw new Error(`hnsw insert: no vector at index ${idx}`);
@@ -121,7 +206,14 @@ export function insertHnsw(
     return;
   }
 
-  const qDist = (node: number) => distance(vec, node);
+  // Issue #228: with a build engine, all pointwise distances go through
+  // dist(aIdx, bIdx) so the whole chain shares one arithmetic order; the
+  // per-neighbor scan uses engine.batch when available (one wasm call per
+  // node's neighbor list instead of one JS call per distance).
+  const dist = engine
+    ? (aIdx: number, bIdx: number) => engine.dist(aIdx, bIdx)
+    : (aIdx: number, bIdx: number) => distance(vectors[aIdx], bIdx);
+  const qDist = (node: number) => dist(idx, node);
 
   // 1. Greedy descent through layers above the node's level
   let ep = hnsw.entryPoint;
@@ -142,10 +234,12 @@ export function insertHnsw(
   }
 
   // 2. Insert on each layer from min(level, maxLevel) down to 0
-  const pairwise = (a: number, b: number) => distance(vectors[a], b);
+  const pairwise = engine
+    ? (a: number, b: number) => engine.dist(a, b)
+    : (a: number, b: number) => distance(vectors[a], b);
   for (let l = Math.min(level, hnsw.maxLevel); l >= 0; l--) {
     const efC = Math.max(M0, DEFAULT_HNSW_EF);
-    const candidates = searchLayer(hnsw, qDist, ep, epDist, l, efC);
+    const candidates = searchLayer(hnsw, qDist, ep, epDist, l, efC, engine?.batch, idx);
     const cap = l === 0 ? M0 : M;
     // Algorithm 4: diversity-aware neighbor selection keeps the graph
     // navigable — a candidate is skipped only when it is CLOSER to an
@@ -191,7 +285,11 @@ function searchLayer(
   ep: number,
   epDist: number,
   layer: number,
-  ef: number
+  ef: number,
+  // issue #228: when present, the unvisited neighbors of each expanded node
+  // are measured in ONE batched call (same order → same heap evolution).
+  batch?: (aIdx: number, nodes: number[]) => Float32Array | null,
+  aIdx?: number
 ): Array<{ node: number; dist: number }> {
   const visited = new Set<number>([ep]);
   // max-heap by dist via simple sorted insertion (candidate lists are small)
@@ -204,10 +302,17 @@ function searchLayer(
     const worst = results.length >= ef ? results[results.length - 1].dist : Infinity;
     if (current.dist > worst) break;
 
-    for (const nb of hnsw.neighbors[current.node][layer] ?? []) {
-      if (visited.has(nb)) continue;
+    const neighborList = hnsw.neighbors[current.node][layer] ?? [];
+    const fresh: number[] = [];
+    for (const nb of neighborList) {
+      if (!visited.has(nb)) fresh.push(nb);
+    }
+    let dists: Float32Array | null = null;
+    if (fresh.length > 1 && batch && aIdx !== undefined) dists = batch(aIdx, fresh);
+    for (let j = 0; j < fresh.length; j++) {
+      const nb = fresh[j];
       visited.add(nb);
-      const d = qDist(nb);
+      const d = dists ? dists[j] : qDist(nb);
       const worstNow = results.length >= ef ? results[results.length - 1].dist : Infinity;
       if (d < worstNow || results.length < ef) {
         candidates.push({ node: nb, dist: d });
@@ -412,3 +517,79 @@ export function deserializeHnsw(raw: unknown, expectedChunksCount?: number): Hns
   }
   return { dim, M, seed, entryPoint, maxLevel, levels, neighbors, deleted };
 }
+
+/**
+ * issue #223: carry a persisted graph across a watch rebuild. `mapping[oldIdx]`
+ * is the chunk's NEW index (undefined = the chunk is gone → tombstone); nodes
+ * carry over because vector reuse is keyed by chunkHash — a surviving chunk
+ * keeps its exact vector, so every stored edge distance stays valid.
+ * Returns null when the remap cannot produce a navigable graph (e.g. the
+ * entry point did not survive) — the caller falls back to a full rebuild.
+ */
+export function remapHnsw(
+  hnsw: HnswIndex,
+  mapping: Array<number | undefined>,
+  newCount: number,
+): HnswIndex | null {
+  if (newCount <= 0) return null;
+  // every new index may receive at most one old node
+  const taken = new Set<number>();
+  const remapOf = new Map<number, number>();
+  for (let oldIdx = 0; oldIdx < mapping.length; oldIdx++) {
+    const newIdx = mapping[oldIdx];
+    if (newIdx === undefined || newIdx < 0 || newIdx >= newCount) continue;
+    if (taken.has(newIdx)) continue;
+    taken.add(newIdx);
+    remapOf.set(oldIdx, newIdx);
+  }
+  // the entry point does NOT need to survive — it re-anchors below
+
+  const levels: number[] = new Array(newCount);
+  const neighbors: number[][][] = new Array(newCount);
+  const deleted = new Set<number>();
+
+  for (const [oldIdx, newIdx] of remapOf) {
+    levels[newIdx] = hnsw.levels[oldIdx];
+    neighbors[newIdx] = (hnsw.neighbors[oldIdx] ?? []).map((links) =>
+      (links ?? [])
+        .map((nb) => remapOf.get(nb))
+        .filter((nb): nb is number => nb !== undefined),
+    );
+    if (hnsw.deleted.has(oldIdx)) deleted.add(newIdx);
+  }
+  // structural holes would break insertHnsw/serialize — unmapped new indices
+  // are placeholders the caller fills with insertHnsw immediately after
+  for (let i = 0; i < newCount; i++) {
+    if (levels[i] === undefined) {
+      levels[i] = 0;
+      neighbors[i] = [[]];
+      deleted.add(i); // invisible until a real node is inserted over it
+    }
+  }
+
+  // the entry point must sit on the TOP layer: after a partial remap the
+  // old entry may carry a lower level, or some survivor may carry a higher
+  // one — re-anchor deterministically to the lowest-index survivor at maxLevel
+  let maxLevel = -1;
+  for (const newIdx of remapOf.values()) maxLevel = Math.max(maxLevel, levels[newIdx]);
+  if (maxLevel < 0) return null;
+  let entryPoint = -1;
+  for (let i = 0; i < newCount; i++) {
+    if (!deleted.has(i) && levels[i] === maxLevel) {
+      entryPoint = i;
+      break;
+    }
+  }
+  if (entryPoint < 0) return null;
+
+  return {
+    dim: hnsw.dim,
+    M: hnsw.M,
+    seed: hnsw.seed,
+    entryPoint,
+    maxLevel,
+    levels,
+    neighbors,
+    deleted,
+  };
+}

package/src/indexer.ts

@@ -36,7 +36,7 @@ import {
 } from './lexical.js';
 import { embeddingAdapterFingerprint, legacyEmbeddingAdapterFingerprint } from './models.js';
 import { trainIVF, serializeIVF, ANN_THRESHOLD } from './ivf.js';
-import { buildHnsw, serializeHnsw, HnswIndex } from './hnsw.js';
+import { buildHnsw, insertHnsw, remapHnsw, serializeHnsw, deserializeHnsw, HnswIndex } from './hnsw.js';
 import { DocumentMetadata } from './frontmatter.js';
 import { serializeBinaryIndex } from './binary-format.js';
 import { defaultChunkTokenBudget } from './markdown-parser.js';
@@ -658,16 +658,55 @@ async function _buildIndexInner({
   if (buildSidecar) {
     const rawVecs = chunks.map((c) => (c.vec instanceof Float32Array ? c.vec : Float32Array.from(c.vec as number[])));
     if (annEngine === 'hnsw') {
-      const hnswIndex: HnswIndex = buildHnsw(rawVecs, {
-        vectors: rawVecs,
-        // cosine distance (lower = closer) over the flat normalized vectors
-        distance: (a, b) => {
-          const bv = rawVecs[b];
+      const hnswDistance = (vectors: Array<Float32Array | number[]>) =>
+        (a: Float32Array | Int8Array | number[], b: number) => {
+          const bv = vectors[b];
           let dot = 0;
           for (let i = 0; i < a.length; i++) dot += (a as ArrayLike<number>)[i] * (bv as ArrayLike<number>)[i];
           return 1 - dot;
-        },
-      });
+        };
+
+      // issue #223: reuse the persisted graph across watch rebuilds — remap
+      // surviving nodes onto the new chunk numbering via chunkHash (vector
+      // reuse is hash-keyed, so every carried edge distance stays valid),
+      // tombstone the gone ones, insert the new ones. Full rebuild only when
+      // no valid previous graph exists (first run, engine/budget switch) or
+      // the remap cannot anchor an entry point.
+      let hnswIndex: HnswIndex | null = null;
+      if (!chunkerChanged && oldIndex.chunks.length > 0) {
+        try {
+          const prev = deserializeHnsw(JSON.parse(fs.readFileSync(hnswPath, 'utf8')), oldIndex.chunks.length);
+          if (prev.levels.length === oldIndex.chunks.length) {
+            const newIdxByHash = new Map<string, number>();
+            chunks.forEach((c, newIdx) => newIdxByHash.set(c.chunkHash || chunkHash(model, c), newIdx));
+            const mapping: Array<number | undefined> = oldIndex.chunks.map((c) =>
+              newIdxByHash.get(c.chunkHash || chunkHash(model, c)),
+            );
+            hnswIndex = remapHnsw(prev, mapping, chunks.length);
+            if (hnswIndex) {
+              const inserted: number[] = [];
+              for (let newIdx = 0; newIdx < chunks.length; newIdx++) {
+                if (hnswIndex.levels[newIdx] !== undefined && !hnswIndex.deleted.has(newIdx)) continue;
+                insertHnsw(hnswIndex, newIdx, hnswDistance(rawVecs), rawVecs);
+                inserted.push(newIdx);
+              }
+              log(
+                `hnsw: reused persisted graph — ` +
+                  `${chunks.length - inserted.length} surviving node(s), ${inserted.length} inserted`,
+              );
+            }
+          }
+        } catch {
+          hnswIndex = null; // corrupt/foreign sidecar — full rebuild below
+        }
+      }
+      if (!hnswIndex) {
+        // issue #228: 'auto' routes the full build through the WASM SIMD
+        // dot-batch kernel when the vectors are a uniform Float32Array matrix
+        // (the hnswDistance closure below is exactly 1 − dot, so semantics
+        // match); otherwise buildHnsw silently keeps the closure path.
+        hnswIndex = buildHnsw(rawVecs, { vectors: rawVecs, distance: hnswDistance(rawVecs), engine: 'auto' });
+      }
       atomicWrite(hnswPath, JSON.stringify(serializeHnsw(hnswIndex)));
       if (fs.existsSync(ivfPath)) {
         try {

package/src/models.ts

@@ -131,7 +131,12 @@ export const MODELS: Record<string, ModelAdapter> = {
   },
 };
 
-export const DEFAULT_MODEL = 'e5-base';
+// Issue #227: granite-multilingual-311m is the default — Apache-2.0 (the
+// measured leader embeddinggemma-300m is license-blocked from being an npm
+// default by the Gemma Terms of Use) and measured +0.078 nDCG@10 over e5-base
+// on the 186-query golden set v2 (RESEARCH.md §15). Same size class (~298 MB
+// q8), 8k ModernBERT context. e5-base stays a registered alias.
+export const DEFAULT_MODEL = 'granite-multilingual-311m';
 export const POOLING: Pooling[] = ['mean', 'last_token', 'none', 'cls'];
 const POOLING_SET = new Set<string>(POOLING);
 

package/src/wasm-dot.ts

@@ -0,0 +1,115 @@
+/**
+ * WASM SIMD128 dot-product engine for the HNSW graph build (issue #228).
+ *
+ * The build's cost profile is ~85% distance computation (measured: 16.5k
+ * distance calls per node at n=5000, dim=256; RESEARCH.md §17), and those
+ * calls sit on a sequential dependency chain — worker_threads parallelism is
+ * structurally blocked (per-insert graph state + the byte-identical-sidecar
+ * determinism contract). The achievable lever is making each dot faster:
+ * this module batches f32 dot products into ONE WebAssembly call
+ * (`dot_batch`: N row indices → N dot products against a shared query row)
+ * with a SIMD128 mul/add inner loop — measured ~6× over the unrolled JS dot
+ * at dim=256 on a 32-dot batch.
+ *
+ * Zero native dependencies: the module is 363 bytes of hand-written WAT
+ * compiled once via wabt and embedded below (source of record:
+ * bench/ann-recall.mjs history + RESEARCH.md §17). Compile failure (no SIMD
+ * in the runtime) degrades to `null` and the caller keeps the JS closure path.
+ *
+ * Determinism: all arithmetic is IEEE f32 with a fixed lane order — repeated
+ * builds on the same corpus produce a byte-identical sidecar per engine. The
+ * SIMD accumulation order differs from the JS closure path in the last ulp,
+ * so the two engines are each self-consistent but not cross-identical; the
+ * sidecar records which engine built it (same engine → same bytes).
+ */
+
+/** The compiled `dot_batch` module (see module docstring for provenance). */
+const KERNEL_B64 =
+  'AGFzbQEAAAABCQFgBX9/f39/AAMCAQAFAwEAAQcTAgNtZW0CAAlkb3RfYmF0Y2gAAAq3AgG0AgQCfwF9An8EeyADQQRsIQlBACEF' +
+  'AkADQCAFIAJODQEgASAFQQRsaigCACAJbCEG/QwAAAAAAAAAAAAAAAAAAAAAIQr9DAAAAAAAAAAAAAAAAAAAAAAhC/0MAAAAAAAAA' +
+  'AAAAAAAAAAAACEM/QwAAAAAAAAAAAAAAAAAAAAAIQ1DAAAAACEHQQAhCAJAA0AgCEEQaiAJSg0BIAogACAIav0AAgAgBiAIav0A' +
+  'AgD95gH95AEhCiAIQRBqIQgMAAsLAkADQCAIIAlODQEgByAAIAhqKgIAIAYgCGoqAgCUkiEHIAhBBGohCAwACwsgBCAFQQRsaiAK' +
+  '/R8AIAr9HwGSIAr9HwIgCv0fA5KSIAcgC/0fACAL/R8BkiAL/R8CIAv9HwOSkpKSOAIAIAVBAWohBQwACwsL';
+
+export interface WasmDotEngine {
+  /** dot(vectors[a], vectors[b]) — one-row batch, same arithmetic as batch(). */
+  point(aIdx: number, bIdx: number): number;
+  /** result[j] = dot(vectors[aIdx], vectors[nodes[j]]) — a view into wasm memory, valid until the next call. */
+  batch(aIdx: number, nodes: number[]): Float32Array;
+  readonly dim: number;
+}
+
+/** Max nodes per batch call — HNSW neighbor lists cap at M0 = 2*M ≤ 64. */
+const MAX_BATCH = 64;
+
+/**
+ * Build the engine over a uniform Float32Array matrix. Returns null when the
+ * runtime cannot compile the SIMD module (then callers keep the JS path).
+ */
+export function createWasmDotEngine(vectors: Array<Float32Array>): WasmDotEngine | null {
+  if (vectors.length === 0) return null;
+  const dim = vectors[0].length;
+  if (dim === 0) return null;
+  for (const v of vectors) {
+    if (!(v instanceof Float32Array) || v.length !== dim) return null;
+  }
+
+  let kernelBuffer: ArrayBuffer;
+  try {
+    const bin = Buffer.from(KERNEL_B64, 'base64');
+    kernelBuffer = new ArrayBuffer(bin.byteLength);
+    new Uint8Array(kernelBuffer).set(bin);
+    if (bin.byteLength === 0) return null;
+  } catch {
+    return null;
+  }
+  let instance: WebAssembly.Instance;
+  try {
+    instance = new WebAssembly.Instance(new WebAssembly.Module(kernelBuffer), {});
+  } catch {
+    return null; // no SIMD support — caller falls back to the JS closure path
+  }
+  const exp = instance.exports as unknown as {
+    mem: WebAssembly.Memory;
+    dot_batch: (q: number, idx: number, n: number, dim: number, out: number) => void;
+  };
+  const mem = exp.mem;
+
+  // Linear memory layout: [vectors row-major][idx scratch u32 × MAX_BATCH][out f32 × MAX_BATCH]
+  const vecBytes = vectors.length * dim * 4;
+  const idxByte = vecBytes;
+  const outByte = idxByte + MAX_BATCH * 4;
+  const needBytes = outByte + MAX_BATCH * 4 + 4096;
+  const pages = Math.ceil(needBytes / 65536);
+  try {
+    if (mem.buffer.byteLength < needBytes) mem.grow(pages);
+  } catch {
+    return null;
+  }
+  // Views are created AFTER the final grow — growing detaches old buffers.
+  const f32 = new Float32Array(mem.buffer);
+  const u32 = new Uint32Array(mem.buffer);
+  for (let i = 0; i < vectors.length; i++) f32.set(vectors[i], i * dim);
+
+  const qByteOf = (aIdx: number) => {
+    if (!Number.isSafeInteger(aIdx) || aIdx < 0 || aIdx >= vectors.length) {
+      throw new Error(`wasm-dot: row index out of range: ${aIdx}`);
+    }
+    return aIdx * dim * 4;
+  };
+
+  return {
+    dim,
+    point(aIdx: number, bIdx: number): number {
+      u32[idxByte >> 2] = bIdx;
+      exp.dot_batch(qByteOf(aIdx), idxByte, 1, dim, outByte);
+      return f32[outByte >> 2];
+    },
+    batch(aIdx: number, nodes: number[]): Float32Array {
+      const n = Math.min(nodes.length, MAX_BATCH);
+      for (let j = 0; j < n; j++) u32[(idxByte >> 2) + j] = nodes[j];
+      exp.dot_batch(qByteOf(aIdx), idxByte, n, dim, outByte);
+      return f32.subarray(outByte >> 2, (outByte >> 2) + n);
+    },
+  };
+}