HNSW (Hierarchical Navigable Small World)

/**
 * HNSW (Hierarchical Navigable Small World) — Multi-layer proximity graph for high-D embeddings
 * Build: O(N log N) empirical via efConstruction beam search | Query: O(log N) via efSearch
 * Distance: Cosine dissimilarity (1 - dot(u, v) / (||u|| * ||v||))
 */
export interface HNSWHit { id: number; distance: number; }

export class HNSWIndex {
  private readonly vectors = new Map<number, Float32Array>();
  private readonly norms = new Map<number, number>();
  private readonly nodeLevels = new Map<number, number>();
  private readonly layers: Array<Map<number, number[]>> = [];
  private entryPoint: number | null = null;
  private maxLevel = -1;
  private readonly levelMult: number;
  private rngState = 0x9e3779b9;

  constructor(
    private readonly M = 24,
    private readonly efConstruction = 96,
    private readonly efSearch = 80
  ) {
    this.levelMult = 1 / Math.log(Math.max(2, M));
  }

  private randomLevel(): number {
    this.rngState = (Math.imul(this.rngState, 1664525) + 1013904223) >>> 0;
    const u = Math.max(1e-9, (this.rngState >>> 8) / 16777216);
    return Math.min(8, Math.floor(-Math.log(u) * this.levelMult));
  }

  private static norm(v: Float32Array): number {
    let s = 0;
    for (let i = 0; i < v.length; i++) s += v[i] * v[i];
    return Math.sqrt(s) || 1;
  }

  private cosineDist(a: Float32Array, normA: number, bId: number): number {
    const b = this.vectors.get(bId)!;
    const normB = this.norms.get(bId)!;
    let dot = 0;
    for (let i = 0; i < a.length; i++) dot += a[i] * b[i];
    return 1 - dot / (normA * normB);
  }

  private searchLayer(query: Float32Array, qNorm: number, epId: number, ef: number, layer: number): HNSWHit[] {
    const visited = new Set<number>([epId]);
    const epDist = this.cosineDist(query, qNorm, epId);
    const candidates: HNSWHit[] = [{ id: epId, distance: epDist }];
    const results: HNSWHit[] = [{ id: epId, distance: epDist }];
    const adj = this.layers[layer];

    while (candidates.length > 0) {
      candidates.sort((a, b) => a.distance - b.distance);
      const curr = candidates.shift()!;
      results.sort((a, b) => a.distance - b.distance);
      const worstResult = results[results.length - 1];
      if (curr.distance > worstResult.distance && results.length >= ef) break;

      for (const nbId of adj.get(curr.id) ?? []) {
        if (visited.has(nbId)) continue;
        visited.add(nbId);
        const d = this.cosineDist(query, qNorm, nbId);
        if (results.length < ef || d < results[results.length - 1].distance) {
          candidates.push({ id: nbId, distance: d });
          results.push({ id: nbId, distance: d });
          if (results.length > ef) {
            results.sort((a, b) => a.distance - b.distance);
            results.pop();
          }
        }
      }
    }
    return results.sort((a, b) => a.distance - b.distance);
  }

  add(id: number, vector: Float32Array): void {
    const vNorm = HNSWIndex.norm(vector);
    this.vectors.set(id, vector);
    this.norms.set(id, vNorm);
    const level = this.randomLevel();
    this.nodeLevels.set(id, level);
    while (this.layers.length <= level) this.layers.push(new Map());
    for (let l = 0; l <= level; l++) this.layers[l].set(id, []);

    if (this.entryPoint === null) {
      this.entryPoint = id;
      this.maxLevel = level;
      return;
    }

    let currEp = this.entryPoint;
    let currDist = this.cosineDist(vector, vNorm, currEp);

    for (let l = this.maxLevel; l > level; l--) {
      let changed = true;
      while (changed) {
        changed = false;
        for (const nb of this.layers[l].get(currEp) ?? []) {
          const d = this.cosineDist(vector, vNorm, nb);
          if (d < currDist) {
            currDist = d;
            currEp = nb;
            changed = true;
          }
        }
      }
    }

    for (let l = Math.min(level, this.maxLevel); l >= 0; l--) {
      const topCandidates = this.searchLayer(vector, vNorm, currEp, this.efConstruction, l);
      const mMax = l === 0 ? this.M * 2 : this.M;
      const selected = topCandidates.slice(0, this.M).map((c) => c.id);
      this.layers[l].set(id, selected);
      for (const nb of selected) {
        const nbList = this.layers[l].get(nb)!;
        nbList.push(id);
        if (nbList.length > mMax) {
          const nbVec = this.vectors.get(nb)!;
          const nbNorm = this.norms.get(nb)!;
          nbList.sort((a, b) => this.cosineDist(nbVec, nbNorm, a) - this.cosineDist(nbVec, nbNorm, b));
          nbList.length = mMax;
        }
      }
      if (topCandidates.length > 0) currEp = topCandidates[0].id;
    }

    if (level > this.maxLevel) {
      this.maxLevel = level;
      this.entryPoint = id;
    }
  }

  searchKNN(query: Float32Array, k = 10): HNSWHit[] {
    if (this.entryPoint === null) return [];
    const qNorm = HNSWIndex.norm(query);
    let currEp = this.entryPoint;
    let currDist = this.cosineDist(query, qNorm, currEp);

    for (let l = this.maxLevel; l >= 1; l--) {
      let changed = true;
      while (changed) {
        changed = false;
        for (const nb of this.layers[l].get(currEp) ?? []) {
          const d = this.cosineDist(query, qNorm, nb);
          if (d < currDist) {
            currDist = d;
            currEp = nb;
            changed = true;
          }
        }
      }
    }

    const hits = this.searchLayer(query, qNorm, currEp, Math.max(this.efSearch, k), 0);
    return hits.slice(0, k);
  }
}