brainy/tests/unit/hnsw/remove-reverse-adjacency.test.ts

/**
 * @module tests/unit/hnsw/remove-reverse-adjacency
 * @description Correctness guard for the reverse-adjacency index that makes
 * `removeItem` O(in-degree) instead of O(N) (and bulk delete O(N·degree) instead
 * of O(N²)). The risk of a reverse index is staleness — a missed maintenance site
 * leaves a dangling reference to a deleted node or a stale referrer. These tests
 * exercise mixed add/delete workloads and assert three invariants:
 *
 *   1. No surviving node's connections reference a deleted id (no dangling edges).
 *   2. The incrementally-maintained reverse index matches a fresh rebuild from the
 *      live adjacency — i.e. link/prune/remove maintenance is exactly consistent.
 *   3. Search still returns the true nearest neighbours of the surviving set
 *      (verified exactly against a brute-force scan on a complete small graph).
 */
import { describe, it, expect, beforeEach } from 'vitest'
import { JsHnswVectorIndex } from '../../../src/hnsw/hnswIndex.js'
import { euclideanDistance } from '../../../src/utils/index.js'
import { MemoryStorage } from '../../../src/storage/adapters/memoryStorage.js'

const DIM = 8

function seededRand(seed: number): () => number {
  let s = seed >>> 0
  return () => {
    s = (s + 0x6d2b79f5) | 0
    let t = Math.imul(s ^ (s >>> 15), 1 | s)
    t = (t + Math.imul(t ^ (t >>> 7), 61 | t)) ^ t
    return ((t ^ (t >>> 14)) >>> 0) / 4294967296
  }
}

/** A deterministic vector pointing in a pseudo-random direction (well-connected graph). */
function vec(idx: number): number[] {
  const rand = seededRand(idx + 1)
  return Array.from({ length: DIM }, () => rand() * 2 - 1)
}

type Noun = { id: string; vector: number[]; connections: Map<number, Set<string>> }

function nounsOf(index: JsHnswVectorIndex): Map<string, Noun> {
  return (index as unknown as { nouns: Map<string, Noun> }).nouns
}

/** Flatten a reverse index (or a freshly-derived one) to sorted `target|level|source` triples. */
function triplesFromIncoming(inc: Map<string, Map<number, Set<string>>>): string[] {
  const out: string[] = []
  for (const [target, byLevel] of inc) {
    for (const [level, sources] of byLevel) {
      for (const source of sources) out.push(`${target}|${level}|${source}`)
    }
  }
  return out.sort()
}

/** Derive the ground-truth reverse index directly from the live forward adjacency. */
function triplesFromAdjacency(nouns: Map<string, Noun>): string[] {
  const out: string[] = []
  for (const [nodeId, node] of nouns) {
    for (const [level, targets] of node.connections) {
      for (const target of targets) out.push(`${target}|${level}|${nodeId}`)
    }
  }
  return out.sort()
}

function assertNoDanglingRefs(nouns: Map<string, Noun>): void {
  for (const [nodeId, node] of nouns) {
    for (const [level, targets] of node.connections) {
      for (const target of targets) {
        expect(nouns.has(target), `node ${nodeId} (L${level}) links to deleted ${target}`).toBe(true)
      }
    }
  }
}

function assertReverseIndexConsistent(index: JsHnswVectorIndex): void {
  // ensureIncoming() returns the live (incrementally-maintained) index, building it
  // only if it was invalidated. The triples it yields must equal the triples derived
  // straight from the forward adjacency — empty leftover sets contribute nothing.
  const live = (index as unknown as { ensureIncoming: () => Map<string, Map<number, Set<string>>> }).ensureIncoming()
  expect(triplesFromIncoming(live)).toEqual(triplesFromAdjacency(nounsOf(index)))
}

function makeIndex(M: number): JsHnswVectorIndex {
  return new JsHnswVectorIndex(
    { M, efConstruction: 200, efSearch: 64, ml: 16 },
    euclideanDistance,
    { useParallelization: false, storage: new MemoryStorage() }
  )
}

describe('HNSW removeItem — reverse-adjacency correctness', () => {
  it('leaves no dangling references and a consistent reverse index after interleaved add/delete', async () => {
    const index = makeIndex(16)
    const live = new Set<string>()

    // Add 120, then interleave: delete every 3rd as we keep adding — exercises the
    // incremental maintenance (link + prune on add, unlink on remove) continuously.
    for (let i = 0; i < 120; i++) {
      const id = `n-${i}`
      await index.addItem({ id, vector: vec(i) })
      live.add(id)
      if (i % 3 === 0 && i > 0) {
        const victim = `n-${i - 1}`
        await index.removeItem(victim)
        live.delete(victim)
      }
    }

    // Now delete ~half of the survivors in one burst (the O(N²)-prone path).
    const survivors = [...live]
    for (let i = 0; i < survivors.length; i += 2) {
      await index.removeItem(survivors[i])
      live.delete(survivors[i])
    }

    const nouns = nounsOf(index)
    expect(nouns.size).toBe(live.size)
    assertNoDanglingRefs(nouns)
    assertReverseIndexConsistent(index)
  })

  it('removing the entry point repeatedly stays consistent', async () => {
    const index = makeIndex(8)
    for (let i = 0; i < 40; i++) await index.addItem({ id: `e-${i}`, vector: vec(i) })

    // Delete whatever is currently the entry point, 20 times.
    for (let k = 0; k < 20; k++) {
      const ep = (index as unknown as { entryPointId: string | null }).entryPointId
      if (!ep) break
      await index.removeItem(ep)
      assertNoDanglingRefs(nounsOf(index))
    }
    assertReverseIndexConsistent(index)
  })

  it('search returns the true nearest neighbours of the surviving set (exact, complete graph)', async () => {
    // M ≥ N-1 ⇒ the small graph is complete and never prunes, so HNSW recall is 100%
    // and we can assert the survivors\' exact top-k against a brute-force scan.
    const index = makeIndex(64)
    const all: Array<{ id: string; v: number[] }> = []
    for (let i = 0; i < 24; i++) {
      const v = vec(i + 1000)
      all.push({ id: `s-${i}`, v })
      await index.addItem({ id: `s-${i}`, vector: v })
    }

    const deleted = new Set(['s-2', 's-5', 's-9', 's-14', 's-20'])
    for (const id of deleted) await index.removeItem(id)

    const survivors = all.filter((x) => !deleted.has(x.id))
    const query = vec(7).map((x) => x * 0.5) // an arbitrary deterministic query point
    const k = 5

    const bruteTopK = survivors
      .map((x) => ({ id: x.id, d: euclideanDistance(query, x.v) }))
      .sort((a, b) => a.d - b.d)
      .slice(0, k)
      .map((x) => x.id)

    const got = (await index.search(query, k)).map(([id]) => id)

    expect(got).toEqual(bruteTopK)
    // And none of the deleted ids can ever surface.
    expect(got.some((id) => deleted.has(id))).toBe(false)
    assertNoDanglingRefs(nounsOf(index))
  })
})