brainy/tests/unit/hnsw/allowed-ids-pushdown.test.ts

/**
 * @module tests/unit/hnsw/allowed-ids-pushdown
 * @description 8.0 #46 (CTX-PUSHDOWN-ALLOWEDIDS) — the vector-search predicate
 * pushdown. Two layers:
 *
 * 1. The JS HNSW index honors `allowedIds: ReadonlySet<string>` by restricting the
 *    beam walk to allowed nodes INSIDE the traversal (walk-all, collect-allowed) —
 *    the semantics that recover the filtered recall a naive top-k-then-filter loses.
 *    An `allowedIds` arriving as an `OpaqueIdSet` (a native/cor roaring Buffer) is
 *    opaque to the JS index and ignored — only a native provider decodes it.
 *
 * 2. `find()` forwards the matched universe to the vector beam walk as an
 *    `OpaqueIdSet` (zero id materialization) when the active metadata provider
 *    exposes the `getIdSetForFilter` producer — the native-stack zero-crossing path.
 */

import { describe, it, expect, beforeEach, afterEach } from 'vitest'
import { JsHnswVectorIndex } from '../../../src/hnsw/hnswIndex.js'
import { euclideanDistance } from '../../../src/utils/index.js'
import { MemoryStorage } from '../../../src/storage/adapters/memoryStorage.js'
import { Brainy } from '../../../src/index.js'
import { NounType } from '../../../src/types/graphTypes.js'
import { createTestConfig } from '../../helpers/test-factory.js'

const DIM = 8

/** Deterministic PRNG (mulberry32) so the synthetic graph is identical every run. */
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 vector at the requested distance-from-origin pointing in a deterministic RANDOM
 * direction (so the HNSW graph is well-connected — diverse directions, not a
 * degenerate axis cluster). The query is the origin, so distance-from-query equals
 * `magnitude` exactly — letting the test place nodes in clean distance bands.
 */
function vecAt(magnitude: number, idx: number): number[] {
  const rand = seededRand(idx + 1)
  const dir = Array.from({ length: DIM }, () => rand() * 2 - 1)
  const norm = Math.hypot(...dir) || 1
  return dir.map((x) => (x / norm) * magnitude)
}

describe('#46 JS HNSW allowedIds pushdown (recall-preserving)', () => {
  let index: JsHnswVectorIndex
  const query = new Array(DIM).fill(0)
  // Two clean distance bands: the disallowed nodes are STRICTLY closer to the query
  // than the allowed targets. A naive top-k-then-filter returns only the disallowed
  // (then drops them ⇒ empty); the pushdown must reach the allowed targets. `M` ≥ N−1
  // makes the small graph complete, so reachability is guaranteed and the test asserts
  // the FILTER mechanism (walk-all, collect-allowed), not emergent HNSW connectivity.
  const disallowed: string[] = []
  const targets: string[] = []

  beforeEach(async () => {
    index = new JsHnswVectorIndex(
      { M: 16, efConstruction: 200, efSearch: 50, ml: 16 },
      euclideanDistance,
      { useParallelization: false, storage: new MemoryStorage() }
    )
    let n = 0
    for (let i = 0; i < 5; i++) {
      const id = `disallowed-${i}`
      disallowed.push(id)
      await index.addItem({ id, vector: vecAt(0.3, n++) }) // closer band
    }
    for (let i = 0; i < 5; i++) {
      const id = `target-${i}`
      targets.push(id)
      await index.addItem({ id, vector: vecAt(0.8, n++) }) // farther band
    }
  })

  afterEach(() => {
    disallowed.length = 0
    targets.length = 0
  })

  it('without restriction, the top-k are all disallowed (so post-filtering would lose recall)', async () => {
    const results = await index.search(query, 3)
    const ids = results.map(([id]) => id)
    // Every nearest neighbor is in the (closer) disallowed band → naive "filter after" = 0 hits.
    expect(ids.length).toBe(3)
    expect(ids.every((id) => disallowed.includes(id))).toBe(true)
  })

  it('with allowedIds, the walk reaches the allowed targets the naive path would miss', async () => {
    const allowed = new Set(targets)
    const results = await index.search(query, 3, undefined, { allowedIds: allowed })
    const ids = results.map(([id]) => id)

    expect(ids.length).toBe(3)
    expect(ids.some((id) => disallowed.includes(id))).toBe(false) // no disallowed leaks through
    expect(ids.every((id) => allowed.has(id))).toBe(true) // only allowed targets
  })

  it('ANDs allowedIds with candidateIds when both are given', async () => {
    const results = await index.search(query, 5, undefined, {
      candidateIds: [targets[0], targets[1], disallowed[0]],
      allowedIds: new Set([targets[0], disallowed[0], disallowed[1]])
    })
    const ids = results.map(([id]) => id)
    // Intersection = { target-0, disallowed-0 } — nothing outside it may appear.
    expect(ids.every((id) => id === targets[0] || id === disallowed[0])).toBe(true)
    expect(ids).toContain(targets[0])
  })

  it('ignores an OpaqueIdSet (Uint8Array) — only a native provider can decode it', async () => {
    const opaque = new Uint8Array([0xca, 0xfe, 0x01, 0x02])
    const results = await index.search(query, 3, undefined, { allowedIds: opaque })
    const ids = results.map(([id]) => id)
    // Unrestricted behavior: the opaque buffer is not interpreted as a filter.
    expect(ids.every((id) => disallowed.includes(id))).toBe(true)
  })
})

describe('#46 find() forwards the opaque universe to the vector beam walk', () => {
  let brain: Brainy

  beforeEach(async () => {
    brain = new Brainy(createTestConfig())
    await brain.init()
    await brain.add({ type: NounType.Document, data: 'machine learning notes', metadata: { author: 'alice' } })
    await brain.add({ type: NounType.Document, data: 'unrelated cooking', metadata: { author: 'bob' } })
  })

  afterEach(async () => {
    await brain.close()
  })

  it('passes the metadata getIdSetForFilter() OpaqueIdSet straight through as allowedIds', async () => {
    const sentinel = new Uint8Array([0xab, 0xcd]) // stand-in for cor's roaring Buffer
    // Stub the native producer the JS metadata manager does not implement.
    ;(brain as any).metadataIndex.getIdSetForFilter = async () => sentinel

    // Record what the vector index actually receives, then run the real search.
    const calls: any[] = []
    const realSearch = (brain as any).index.search.bind((brain as any).index)
    ;(brain as any).index.search = async (vec: any, k: any, f: any, opts: any) => {
      if (opts) calls.push(opts)
      return realSearch(vec, k, f, opts)
    }

    await brain.find({ query: 'machine learning', where: { author: 'alice' } })

    const withUniverse = calls.find((o) => o.allowedIds !== undefined)
    expect(withUniverse).toBeDefined()
    // The opaque Buffer is forwarded verbatim (same reference — never materialized/copied).
    expect(withUniverse.allowedIds).toBe(sentinel)
    // The materialized candidateIds path is still present for the JS index.
    expect(withUniverse.candidateIds).toBeDefined()
  })

  it('omits allowedIds when the metadata provider has no opaque producer (JS path)', async () => {
    // No getIdSetForFilter stub — the JS manager genuinely lacks it.
    const calls: any[] = []
    const realSearch = (brain as any).index.search.bind((brain as any).index)
    ;(brain as any).index.search = async (vec: any, k: any, f: any, opts: any) => {
      if (opts) calls.push(opts)
      return realSearch(vec, k, f, opts)
    }

    await brain.find({ query: 'machine learning', where: { author: 'alice' } })

    const withOpts = calls.find((o) => o.candidateIds !== undefined)
    expect(withOpts).toBeDefined()
    expect(withOpts.allowedIds).toBeUndefined() // JS path: materialized candidateIds only
  })
})