Skip to main content

core_rules/
index.rs

1use crate::def::{Predicate, MAX_KEYMATCH_LIST};
2use crate::hnsw::HnswIndex;
3use core_storage::{list_tokens, Value, ValueKey};
4use std::collections::{BTreeMap, BTreeSet};
5
6// ---------------------------------------------------------------------------
7// IVF-Flat constants (Plan 11 T4)
8// ---------------------------------------------------------------------------
9
10/// Minimum number of k-means clusters. Keeps probing meaningful even for
11/// small sides (< 16 vectors).
12pub const IVF_K_MIN: usize = 4;
13
14/// Maximum number of k-means clusters. Bounds centroid memory and fit time.
15pub const IVF_K_MAX: usize = 1024;
16
17/// Fixed number of k-means iterations per fit (deterministic convergence).
18pub const IVF_ITERATIONS: usize = 12;
19
20/// Probe denominator: P = max(1, ceil(k / IVF_PROBE_DENOM)) centroids queried
21/// per lookup.  k=4 → P=1; k=64 → P=4; k=1024 → P=64.
22pub const IVF_PROBE_DENOM: usize = 16;
23
24/// Rebuild an approximate rule when dst-side IVF drift exceeds this count.
25/// Drift is only known after apply, so the WAL path issues `RebuildRule` as a
26/// second commit (not a pre-WAL Batch).
27pub const IVF_DRIFT_REBUILD: u64 = 256;
28
29thread_local! {
30    static IVF_DRIFT_REBUILD_OVERRIDE: std::cell::Cell<Option<u64>> =
31        const { std::cell::Cell::new(None) };
32}
33
34pub(crate) fn ivf_drift_rebuild_threshold() -> u64 {
35    IVF_DRIFT_REBUILD_OVERRIDE.with(|c| c.get().unwrap_or(IVF_DRIFT_REBUILD))
36}
37
38/// Run `f` with a temporary IVF dst-drift rebuild threshold.
39/// Restores the previous override (including across panics).
40pub fn with_ivf_drift_rebuild<R>(threshold: u64, f: impl FnOnce() -> R) -> R {
41    IVF_DRIFT_REBUILD_OVERRIDE.with(|c| {
42        let prev = c.replace(Some(threshold));
43        let out = std::panic::catch_unwind(std::panic::AssertUnwindSafe(f));
44        c.set(prev);
45        match out {
46            Ok(v) => v,
47            Err(p) => std::panic::resume_unwind(p),
48        }
49    })
50}
51
52/// Beam ceiling for the widening loop an exact `VectorSimilar` rule runs
53/// (see [`CandidateSpec::Hnsw`]'s `floor`). Reaching it with the floor still
54/// unreached means the candidate set is the whole tracked set, which is what the
55/// rule did before 0.6.6.
56pub const EF_MAX: usize = 4_096;
57
58/// Slack on the beam's stopping comparison, covering the `f32` arithmetic the
59/// index answers with ([`crate::hnsw::HnswIndex::search`] documents ~1e-6).
60///
61/// The beam's similarities are a candidate *ordering* number and never a
62/// reported score — every score on an edge is recomputed from the `f64` store.
63/// Requiring the worst hit to be *clearly* below `min` before the beam is
64/// trusted means `f32` rounding can cost one extra doubling and can never cost
65/// a pair.
66const BEAM_FLOOR_SLACK: f64 = 1e-5;
67
68thread_local! {
69    static EF_MAX_OVERRIDE: std::cell::Cell<Option<usize>> = const { std::cell::Cell::new(None) };
70}
71
72/// The beam ceiling the widening loop consults, honouring [`with_ef_max`].
73pub fn ef_max() -> usize {
74    EF_MAX_OVERRIDE.with(|c| c.get().unwrap_or(EF_MAX)).max(1)
75}
76
77/// Run `f` with a temporary beam ceiling. Test hook, in the shape of
78/// [`with_hnsw_build_batch`] — it exists so a test can reach the ceiling with a
79/// few hundred vectors instead of the [`EF_MAX`] thousands.
80///
81/// The override is thread-local, so `f` must do its work on the calling thread.
82pub fn with_ef_max<R>(cap: usize, f: impl FnOnce() -> R) -> R {
83    EF_MAX_OVERRIDE.with(|c| {
84        let prev = c.replace(Some(cap));
85        let out = std::panic::catch_unwind(std::panic::AssertUnwindSafe(f));
86        c.set(prev);
87        match out {
88            Ok(v) => v,
89            Err(p) => std::panic::resume_unwind(p),
90        }
91    })
92}
93
94// ---------------------------------------------------------------------------
95// Sliced HNSW build (v0.6.6 T2)
96// ---------------------------------------------------------------------------
97
98/// Vectors inserted into one rule's HNSW graph per build slice. `create_rule`
99/// does one slice inline; `pump_index_build` does one slice per pending rule
100/// per call. A corpus at or below this size is built in a single commit and
101/// behaves exactly as it did before 0.6.6.
102pub const HNSW_BUILD_BATCH: usize = 2_048;
103
104thread_local! {
105    static HNSW_BUILD_BATCH_OVERRIDE: std::cell::Cell<Option<usize>> =
106        const { std::cell::Cell::new(None) };
107}
108
109pub(crate) fn hnsw_build_batch() -> usize {
110    HNSW_BUILD_BATCH_OVERRIDE
111        .with(|c| c.get().unwrap_or(HNSW_BUILD_BATCH))
112        .max(1)
113}
114
115/// Run `f` with a temporary build-slice size. Test hook, in the shape of
116/// [`with_ivf_drift_rebuild`]. Restores the previous override (including
117/// across panics).
118///
119/// The override is thread-local, so `f` must do its `create_rule` **and** its
120/// pumping on the calling thread.
121pub fn with_hnsw_build_batch<R>(batch: usize, f: impl FnOnce() -> R) -> R {
122    HNSW_BUILD_BATCH_OVERRIDE.with(|c| {
123        let prev = c.replace(Some(batch));
124        let out = std::panic::catch_unwind(std::panic::AssertUnwindSafe(f));
125        c.set(prev);
126        match out {
127            Ok(v) => v,
128            Err(p) => std::panic::resume_unwind(p),
129        }
130    })
131}
132
133/// What an insert does with the `Hnsw` leg of a candidate spec. Every variant
134/// records `hnsw_tracked` — that set is the fallback candidate list and has to
135/// cover the whole side regardless of who fills the graph.
136enum HnswLeg<'a> {
137    /// Insert the vector into the graph. The ordinary write path.
138    All,
139    /// Skip ids the adopted graph already holds. The open-time scan.
140    Skip(&'a BTreeSet<u32>),
141    /// Leave the graph alone; a build slice supplies the vector later.
142    Defer,
143}
144
145/// Whether `spec` would put at least one vector of `node`'s props into an HNSW
146/// graph — the predicate a sliced build counts with, so that the total it
147/// reports and the progress it makes are decided by the same rule.
148pub fn hnsw_vector_present(spec: &CandidateSpec, get: &dyn Fn(&str) -> Option<Value>) -> bool {
149    match spec {
150        CandidateSpec::Union(specs) | CandidateSpec::Intersect(specs) => {
151            specs.iter().any(|s| hnsw_vector_present(s, get))
152        }
153        CandidateSpec::Hnsw { field, .. } => {
154            get(field).as_ref().and_then(as_numeric_list).is_some()
155        }
156        _ => false,
157    }
158}
159
160/// k = ceil(sqrt(n)) clamped to [IVF_K_MIN, IVF_K_MAX].
161pub fn cluster_k(n: usize) -> usize {
162    if n == 0 {
163        return IVF_K_MIN;
164    }
165    let k = (n as f64).sqrt().ceil() as usize;
166    k.clamp(IVF_K_MIN, IVF_K_MAX)
167}
168
169/// P = max(1, ceil(k / IVF_PROBE_DENOM)).
170pub fn probe_count(k: usize) -> usize {
171    k.div_ceil(IVF_PROBE_DENOM).max(1)
172}
173
174/// L2-normalize `xs`. Returns `None` for the zero vector (skipped, not clustered).
175fn l2_normalize(xs: &[f64]) -> Option<Vec<f64>> {
176    let n = xs.iter().map(|x| x * x).sum::<f64>().sqrt();
177    if n == 0.0 {
178        return None;
179    }
180    Some(xs.iter().map(|x| x / n).collect())
181}
182
183/// Squared Euclidean distance between two equal-length slices.
184/// Returns `f64::MAX` on dimension mismatch so callers always have a valid order.
185fn l2_sq(a: &[f64], b: &[f64]) -> f64 {
186    if a.len() != b.len() {
187        return f64::MAX;
188    }
189    a.iter().zip(b.iter()).map(|(x, y)| (x - y) * (x - y)).sum()
190}
191
192/// Index of the nearest centroid to `xs` by L2 distance (minimum squared).
193/// Returns 0 when `centroids` is empty.
194pub fn nearest_centroid(centroids: &[Vec<f64>], xs: &[f64]) -> usize {
195    centroids
196        .iter()
197        .enumerate()
198        .min_by(|(_, a), (_, b)| {
199            l2_sq(xs, a)
200                .partial_cmp(&l2_sq(xs, b))
201                .unwrap_or(std::cmp::Ordering::Equal)
202        })
203        .map(|(i, _)| i)
204        .unwrap_or(0)
205}
206
207/// FNV-1a 64-bit hash — stable, documented, NOT DefaultHasher.
208/// Used to seed k-means so the same rule name always produces the same
209/// clusters on the same data (WAL replay identity).
210pub fn fnv1a_u64(data: &[u8]) -> u64 {
211    const FNV_OFFSET: u64 = 14_695_981_039_346_656_037;
212    const FNV_PRIME: u64 = 1_099_511_628_211;
213    let mut h = FNV_OFFSET;
214    for &b in data {
215        h ^= b as u64;
216        h = h.wrapping_mul(FNV_PRIME);
217    }
218    h
219}
220
221/// Seeded LCG step — Knuth multiplicative; used for centroid init and empty
222/// cluster reseeding.
223#[inline]
224fn lcg_next(state: u64) -> u64 {
225    state
226        .wrapping_mul(6_364_136_223_846_793_005)
227        .wrapping_add(1_442_695_040_888_963_407)
228}
229
230/// Fit k-means over `vecs` (node_id, vector) pairs.
231///
232/// - Each vector is L2-normalized before clustering (zero vectors skipped).
233/// - `k` is clamped to `min(k, vecs.len())` so we never request more centroids
234///   than vectors.
235/// - Centroids are initialised by seeded LCG selection without replacement.
236/// - 12 iterations; empty clusters are deterministically reseeded from the
237///   full dataset.
238/// - Returns a `Vec<Vec<f64>>` of k centroids (same length as `xs` entries).
239pub fn kmeans_fit(vecs: &[(u32, Vec<f64>)], k: usize, seed: u64) -> Vec<Vec<f64>> {
240    let vecs: Vec<(u32, Vec<f64>)> = vecs
241        .iter()
242        .filter_map(|(id, xs)| l2_normalize(xs).map(|n| (*id, n)))
243        .collect();
244    if vecs.is_empty() || k == 0 {
245        return vec![];
246    }
247    let n = vecs.len();
248    let k = k.min(n);
249    let dim = vecs[0].1.len();
250    if dim == 0 {
251        return vec![];
252    }
253
254    // --- Centroid initialisation: pick k distinct indices via seeded LCG ---
255    let mut state = seed;
256    let mut used = vec![false; n];
257    let mut init_idxs: Vec<usize> = Vec::with_capacity(k);
258    let mut attempts = 0usize;
259    while init_idxs.len() < k && attempts < n * 4 {
260        state = lcg_next(state);
261        let idx = (state >> 33) as usize % n;
262        if !used[idx] {
263            used[idx] = true;
264            init_idxs.push(idx);
265        }
266        attempts += 1;
267    }
268    // If LCG didn't yield k distinct indices (pathological: n very small or
269    // many collisions), fill sequentially.
270    if init_idxs.len() < k {
271        for (i, in_use) in used.iter().enumerate().take(n) {
272            if !in_use {
273                init_idxs.push(i);
274                if init_idxs.len() == k {
275                    break;
276                }
277            }
278        }
279    }
280    let mut centroids: Vec<Vec<f64>> = init_idxs.iter().map(|&i| vecs[i].1.clone()).collect();
281    let mut assignments = vec![0usize; n];
282
283    // --- k-means iterations ---
284    for iter in 0..IVF_ITERATIONS {
285        // Assignment step
286        for (j, (_, xs)) in vecs.iter().enumerate() {
287            assignments[j] = nearest_centroid(&centroids, xs);
288        }
289
290        // Update step: accumulate sums and counts
291        let mut sums = vec![vec![0.0f64; dim]; k];
292        let mut counts = vec![0usize; k];
293        for (j, (_, xs)) in vecs.iter().enumerate() {
294            let c = assignments[j];
295            counts[c] += 1;
296            for d in 0..dim {
297                sums[c][d] += xs[d];
298            }
299        }
300
301        // Compute new centroids; collect empty ones for reseed
302        let mut new_centroids = vec![vec![0.0f64; dim]; k];
303        let mut empty: Vec<usize> = Vec::new();
304        for c in 0..k {
305            if counts[c] == 0 {
306                empty.push(c);
307            } else {
308                for d in 0..dim {
309                    new_centroids[c][d] = sums[c][d] / counts[c] as f64;
310                }
311            }
312        }
313
314        // Deterministic empty-cluster reseed: pick a vector from the dataset
315        // seeded by (original seed XOR iteration XOR empty-cluster-index).
316        for (ei, ec) in empty.into_iter().enumerate() {
317            let reseed =
318                seed ^ (iter as u64).wrapping_mul(0x9E37) ^ (ei as u64).wrapping_mul(0x1234_5679);
319            let mut rs = lcg_next(reseed);
320            rs = lcg_next(rs);
321            let pick = (rs >> 33) as usize % n;
322            new_centroids[ec] = vecs[pick].1.clone();
323        }
324
325        centroids = new_centroids;
326    }
327
328    centroids
329}
330
331#[cfg(test)]
332thread_local! {
333    static VECTOR_DIM_REJECT: std::cell::Cell<bool> = const { std::cell::Cell::new(true) };
334    static VECTOR_EARLY_EXIT: std::cell::Cell<bool> = const { std::cell::Cell::new(true) };
335}
336
337fn vector_dim_reject_enabled() -> bool {
338    #[cfg(test)]
339    {
340        VECTOR_DIM_REJECT.with(|c| c.get())
341    }
342    #[cfg(not(test))]
343    {
344        true
345    }
346}
347
348pub(crate) fn vector_early_exit_enabled() -> bool {
349    #[cfg(test)]
350    {
351        VECTOR_EARLY_EXIT.with(|c| c.get())
352    }
353    #[cfg(not(test))]
354    {
355        true
356    }
357}
358
359thread_local! {
360    /// `None` until `MUSHROOMDB_VECTOR_SCAN` has been read on this thread.
361    /// [`with_vector_scan`] replaces it for the duration of a closure.
362    static VECTOR_SCAN: std::cell::Cell<Option<bool>> = const { std::cell::Cell::new(None) };
363}
364
365/// `MUSHROOMDB_VECTOR_SCAN=1` forces every `VectorSimilar` rule back onto the
366/// full-scan candidate path: O(n²) per rule, and every pair above the rule's
367/// `min` provably found.
368///
369/// Same shape as [`vector_early_exit_enabled`] and `vector_dim_reject_enabled`,
370/// except that the switch is an environment variable rather than a test-only
371/// hook — it is the documented way for a caller to buy the exactness guarantee
372/// back.
373pub fn vector_scan_forced() -> bool {
374    VECTOR_SCAN.with(|c| match c.get() {
375        Some(v) => v,
376        None => {
377            let v = std::env::var("MUSHROOMDB_VECTOR_SCAN")
378                .map(|s| s == "1" || s.eq_ignore_ascii_case("true"))
379                .unwrap_or(false);
380            c.set(Some(v));
381            v
382        }
383    })
384}
385
386/// Run `f` with the full-scan candidate path forced on or off, whatever
387/// `MUSHROOMDB_VECTOR_SCAN` says. Test hook, in the shape of
388/// [`with_hnsw_build_batch`](crate::with_hnsw_build_batch).
389///
390/// The override is thread-local, so `f` must do its work on the calling thread.
391///
392/// # The engine outlives the closure
393///
394/// This switches which candidate spec a rule is *asked for*, and several
395/// decisions are taken once and remembered: a rule created inside the closure
396/// with the scan forced on builds no HNSW graph, so using that same engine
397/// outside the closure leaves the rule answering from `hnsw_tracked` — correct,
398/// and a full scan — until something rebuilds it. Either keep the engine inside
399/// the closure, as the equivalence test does, or reopen the store afterwards.
400pub fn with_vector_scan<R>(enabled: bool, f: impl FnOnce() -> R) -> R {
401    VECTOR_SCAN.with(|c| {
402        let prev = c.replace(Some(enabled));
403        let out = std::panic::catch_unwind(std::panic::AssertUnwindSafe(f));
404        c.set(prev);
405        match out {
406            Ok(v) => v,
407            Err(p) => std::panic::resume_unwind(p),
408        }
409    })
410}
411
412/// Force the ScanAll dim fast-reject on or off. Identity-proof hook.
413#[cfg(test)]
414pub fn with_vector_dim_reject<R>(enabled: bool, f: impl FnOnce() -> R) -> R {
415    VECTOR_DIM_REJECT.with(|c| {
416        let prev = c.replace(enabled);
417        let out = std::panic::catch_unwind(std::panic::AssertUnwindSafe(f));
418        c.set(prev);
419        match out {
420            Ok(v) => v,
421            Err(p) => std::panic::resume_unwind(p),
422        }
423    })
424}
425
426/// Force the checkpointed Cauchy-Schwarz early-exit on or off. Identity-proof hook.
427#[cfg(test)]
428pub fn with_vector_early_exit<R>(enabled: bool, f: impl FnOnce() -> R) -> R {
429    VECTOR_EARLY_EXIT.with(|c| {
430        let prev = c.replace(enabled);
431        let out = std::panic::catch_unwind(std::panic::AssertUnwindSafe(f));
432        c.set(prev);
433        match out {
434            Ok(v) => v,
435            Err(p) => std::panic::resume_unwind(p),
436        }
437    })
438}
439
440#[derive(Debug, Default)]
441pub struct SideIndex {
442    by_key: BTreeMap<ValueKey, BTreeSet<u32>>,
443    /// Per-node `(dim, L2 norm)` for `ScanAll` members. Maintained by the
444    /// same insert/remove choke-points as `by_key`. Cosine still reads live
445    /// props; `dim` is a fast-reject; `norm` is the primary freshness gate for
446    /// the checkpointed Cauchy-Schwarz early-exit (Plan 11 T3).
447    vec_meta: BTreeMap<u32, (u32, f64)>,
448    /// Per-node checkpointed suffix norms for the Cauchy-Schwarz early-exit.
449    /// `ckpts[i]` = L2 norm of `xs[i * dim / 8 ..]`.
450    /// `ckpts[0]` = full L2 norm; `ckpts[7]` = norm of the last eighth.
451    /// Built at index-insert, torn out at index-remove — maintained in lockstep
452    /// with `vec_meta` by the same choke-points.
453    /// Memory: 8 × 8 = 64 bytes per indexed vector (6.4 MB at 100k vectors).
454    vec_checkpoints: BTreeMap<u32, [f64; 8]>,
455    /// Per-node first element (`xs[0]`) for heuristic permutation detection.
456    /// A permuted vector can share `(dim, norm)` with the indexed one but
457    /// differs at `xs[0]` in virtually all realistic cases, so comparing this
458    /// one extra f64 (8 bytes per vector) breaks same-norm permutation aliasing
459    /// cheaply.  This is heuristic hardening — not a proof — but eliminates
460    /// the energy-distribution construction identified in the Plan 11 T3 review.
461    vec_anchor: BTreeMap<u32, f64>,
462
463    // --- IVF-Flat fields (Plan 11 T4; only populated for VectorClusters specs) ---
464    /// Raw vectors stored for IVF fitting and assignment-on-insert.
465    /// Populated at insert-time, torn out at remove-time.
466    /// Memory: O(n × dim) per indexed side — present only for approximate rules.
467    ivf_raw: BTreeMap<u32, Vec<f64>>,
468    /// Fitted k-means centroids (empty until after first `fit_ivf_clusters` call).
469    ivf_centroids: Vec<Vec<f64>>,
470    /// Per-node cluster assignment post-fit.
471    /// `by_key[ivf_cluster_key(cluster)] → {node_ids}`.
472    ivf_clusters: BTreeMap<u32, usize>,
473    /// Count of vector inserts/removes since last fit.  When dst-side drift
474    /// exceeds [`IVF_DRIFT_REBUILD`] on an approximate rule, apply queues a
475    /// `RebuildRule` second commit (fit resets this to zero).
476    pub ivf_drift: u64,
477
478    // --- HNSW fields (default for approximate: true + VectorSimilar) ---
479    /// HNSW graph; `None` until `init_hnsw` is called.
480    hnsw: Option<HnswIndex>,
481    /// All node ids inserted via `CandidateSpec::Hnsw`.
482    /// Used as a full-scan fallback when `hnsw` is `None` or has no entry point.
483    hnsw_tracked: BTreeSet<u32>,
484}
485
486#[derive(Debug, Default)]
487pub struct RuleIndex {
488    pub src_side: SideIndex,
489    pub dst_side: SideIndex,
490}
491
492#[derive(Debug)]
493pub enum CandidateSpec<'a> {
494    ByKey,
495    Scalar {
496        field: &'a str,
497    },
498    Tokens {
499        field: &'a str,
500    },
501    /// Src side of a `KeyMatch`-rooted rule: the FK field's scalar value, or —
502    /// when that value is a list — one bucket per **string** element (the first
503    /// [`MAX_KEYMATCH_LIST`] in stored order, non-strings skipped).
504    ///
505    /// The reverse lookup always probes with a single key (the destination
506    /// node's key), so a multi-valued FK has to fan out at index time: a src
507    /// node listing n keys sits in n buckets and is found through any of them.
508    /// A scalar value indexes exactly as [`CandidateSpec::Scalar`] does.
509    ScalarOrElements {
510        field: &'a str,
511    },
512    NumericBucket {
513        field: &'a str,
514        tolerance: f64,
515    },
516    GeoGrid {
517        field: &'a str,
518        km: f64,
519    },
520    ScanAll {
521        field: &'a str,
522    },
523    /// IVF-Flat approximate candidate selection (legacy; still supported as
524    /// direct fallback — no longer the default for `approximate: true`).
525    ///
526    /// k-means fitted over the indexed side's vectors; candidates are members
527    /// of the P = `max(1, ceil(k/16))` nearest centroids to the query vector.
528    /// NOT a superset of true positives — recall floor governs correctness.
529    VectorClusters {
530        field: &'a str,
531        min: f64,
532    },
533    /// HNSW approximate candidate selection (default for `approximate: true`).
534    ///
535    /// Returns the `k` nearest vectors by cosine similarity from the in-tree
536    /// HNSW graph.  Falls back to returning all tracked nodes when the graph
537    /// has no entry point (e.g. before any node is inserted, or when used
538    /// without calling `init_hnsw`).
539    Hnsw {
540        field: &'a str,
541        /// Number of approximate candidates to return; typically
542        /// `max(max_edges, 64)` from the owning `RuleDef`.
543        k: usize,
544        /// `Some(min)` for an exact rule: widen the beam until its worst hit
545        /// falls below `min`, so a qualifying node cannot be sitting outside a
546        /// truncated beam. `None` for an approximate rule: one pass at `k`,
547        /// which is what the rule did before 0.6.6.
548        floor: Option<f64>,
549    },
550    /// Union of multiple candidate specs, used for `Any` predicates.
551    ///
552    /// Each branch of the `Any` predicate contributes its own candidate set
553    /// (key index, token index, numeric bucket, etc.); the resulting candidate
554    /// set is their union.  Insert and remove recurse into every child spec so
555    /// the index stays coherent for all branches simultaneously.
556    Union(Vec<CandidateSpec<'a>>),
557    /// Intersection of multiple candidate specs, used for `All` predicates.
558    ///
559    /// Each conjunct contributes its own candidate set; the result is their
560    /// intersection (empty child → empty). `ScanAll` children are skipped at
561    /// probe time (they are the universe); if every child is `ScanAll`, the
562    /// spec stays a full scan. Insert/remove recurse into every child.
563    Intersect(Vec<CandidateSpec<'a>>),
564}
565
566/// Returns the exact candidate strategy derived from `p`.
567///
568/// `All(parts)` returns `Intersect` of each part's spec. A leading
569/// `VectorSimilar` is `ScanAll` and is skipped at probe time when another
570/// conjunct has an index; candidates stay a superset of true matches.
571///
572/// `Any(parts)` returns `Union` of each branch's candidate spec — the correct
573/// superset for OR semantics.
574///
575/// # Panics
576///
577/// Panics on `All([])` or `Any([])`. Predicates must pass `RuleDef::validate()` first.
578pub fn candidate_spec(p: &Predicate) -> CandidateSpec<'_> {
579    match p {
580        Predicate::KeyMatch { .. } => CandidateSpec::ByKey,
581        Predicate::FieldEqual { field } => CandidateSpec::Scalar { field },
582        Predicate::Overlap { field, .. } => CandidateSpec::Tokens { field },
583        Predicate::NumericWithin { field, tolerance } => CandidateSpec::NumericBucket {
584            field,
585            tolerance: *tolerance,
586        },
587        Predicate::GeoRadius { field, km } => CandidateSpec::GeoGrid { field, km: *km },
588        Predicate::VectorSimilar { field, .. } => CandidateSpec::ScanAll { field },
589        Predicate::All(parts) => {
590            debug_assert!(
591                !parts.is_empty(),
592                "candidate_spec requires a validated predicate"
593            );
594            CandidateSpec::Intersect(parts.iter().map(candidate_spec).collect())
595        }
596        Predicate::Any(parts) => {
597            debug_assert!(
598                !parts.is_empty(),
599                "candidate_spec requires a validated predicate"
600            );
601            CandidateSpec::Union(parts.iter().map(candidate_spec).collect())
602        }
603    }
604}
605
606/// Approximate candidate strategy: like `candidate_spec` but replaces
607/// `ScanAll` with `CandidateSpec::Hnsw` for `VectorSimilar`-rooted predicates.
608///
609/// Used when `RuleDef::approximate == true`.  `All` is `Intersect` of each
610/// child's approx spec (not `parts[0]`), so `FieldEqual` / `NumericWithin`
611/// conjuncts still probe their indexes.
612///
613/// `k` is the number of HNSW candidates to return; callers should use
614/// `max(max_edges, 64)` from the owning `RuleDef`.  Use the public
615/// zero-argument wrapper (`candidate_spec_approx`) for tests that don't
616/// need a specific k (defaults to 64).
617///
618/// `Any` predicates cannot be `approximate=true` (validate() rejects them),
619/// so `Any` falls through to `candidate_spec` (exact Union path).
620///
621/// # Panics
622///
623/// Panics on `All([])` or `Any([])`. Predicates must pass `RuleDef::validate()` first.
624pub fn candidate_spec_approx(p: &Predicate) -> CandidateSpec<'_> {
625    candidate_spec_approx_with_k(p, 64)
626}
627
628/// Like `candidate_spec_approx` but with an explicit HNSW candidate count `k`.
629///
630/// The beam takes no floor, so the search is one pass at `k`: the approximate
631/// rule's behaviour. [`candidate_spec_approx_with_floor`] is the exact rule's
632/// version.
633pub fn candidate_spec_approx_with_k(p: &Predicate, k: usize) -> CandidateSpec<'_> {
634    candidate_spec_approx_with_floor(p, k, false)
635}
636
637/// [`candidate_spec_approx_with_k`], optionally taking each `VectorSimilar`'s
638/// own `min` as the beam's stopping similarity.
639///
640/// `floored` is what separates an exact rule from an approximate one: with it,
641/// the beam widens until its worst hit falls below `min`, so a node above `min`
642/// cannot be sitting outside a truncated beam.
643pub fn candidate_spec_approx_with_floor(
644    p: &Predicate,
645    k: usize,
646    floored: bool,
647) -> CandidateSpec<'_> {
648    match p {
649        Predicate::VectorSimilar { field, min } => CandidateSpec::Hnsw {
650            field,
651            k,
652            floor: floored.then_some(*min),
653        },
654        Predicate::All(parts) => {
655            debug_assert!(
656                !parts.is_empty(),
657                "candidate_spec_approx requires a validated predicate"
658            );
659            CandidateSpec::Intersect(
660                parts
661                    .iter()
662                    .map(|p| candidate_spec_approx_with_floor(p, k, floored))
663                    .collect(),
664            )
665        }
666        other => candidate_spec(other),
667    }
668}
669
670/// True when `spec` probes an HNSW graph anywhere, so the owning rule needs one
671/// built. Every `VectorSimilar`-rooted rule does, exact or approximate, unless
672/// [`vector_scan_forced`] has put it back on the full scan.
673pub fn spec_has_hnsw(spec: &CandidateSpec<'_>) -> bool {
674    match spec {
675        CandidateSpec::Hnsw { .. } => true,
676        CandidateSpec::Union(parts) | CandidateSpec::Intersect(parts) => {
677            parts.iter().any(spec_has_hnsw)
678        }
679        _ => false,
680    }
681}
682
683pub(crate) fn as_finite_f64(v: &Value) -> Option<f64> {
684    match v {
685        Value::Int(i) => Some(*i as f64),
686        Value::Float(f) if f.is_finite() => Some(*f),
687        _ => None,
688    }
689}
690
691fn as_latlon(v: &Value) -> Option<(f64, f64)> {
692    let Value::List(items) = v else {
693        return None;
694    };
695    if items.len() != 2 {
696        return None;
697    }
698    let lat = as_finite_f64(&items[0])?;
699    let lon = as_finite_f64(&items[1])?;
700    if (-90.0..=90.0).contains(&lat) && (-180.0..=180.0).contains(&lon) {
701        Some((lat, lon))
702    } else {
703        None
704    }
705}
706
707pub(crate) fn as_numeric_list(v: &Value) -> Option<Vec<f64>> {
708    let Value::List(items) = v else {
709        return None;
710    };
711    if items.is_empty() {
712        return None;
713    }
714    items.iter().map(as_finite_f64).collect()
715}
716
717fn vec_dim_norm(v: &Value) -> Option<(u32, f64)> {
718    let xs = as_numeric_list(v)?;
719    let mut n2 = 0.0;
720    for x in &xs {
721        n2 += *x * *x;
722    }
723    Some((xs.len() as u32, n2.sqrt()))
724}
725
726/// Checkpointed suffix norms for Cauchy-Schwarz early exit.
727///
728/// `ckpts[i]` = L2 norm of `xs[boundary(i)..]` where `boundary(i) = i * dim / 8`.
729/// `ckpts[0]` equals the full L2 norm; `ckpts[7]` is the last eighth's norm.
730/// Multiple checkpoints may share the same boundary for dim < 8 (correct but no-op).
731fn compute_ckpts(xs: &[f64]) -> [f64; 8] {
732    let dim = xs.len();
733    let mut ckpts = [0.0f64; 8];
734    if dim == 0 {
735        return ckpts;
736    }
737    // boundaries[i] = i * dim / 8 (integer division).
738    let boundaries: [usize; 8] = std::array::from_fn(|i| i * dim / 8);
739    let mut suffix_sq = 0.0f64;
740    // Walk right-to-left; ci is the highest checkpoint not yet recorded.
741    let mut ci = 7i32;
742    for j in (0..dim).rev() {
743        suffix_sq += xs[j] * xs[j];
744        // Assign all checkpoints whose boundary equals j.
745        while ci >= 0 && boundaries[ci as usize] == j {
746            ckpts[ci as usize] = suffix_sq.sqrt();
747            ci -= 1;
748        }
749    }
750    ckpts
751}
752
753fn floor_to_i64(x: f64) -> i64 {
754    let floored = x.floor();
755    if !floored.is_finite() {
756        return 0;
757    }
758    if floored >= i64::MAX as f64 {
759        i64::MAX
760    } else if floored <= i64::MIN as f64 {
761        i64::MIN
762    } else {
763        floored as i64
764    }
765}
766
767/// Two values within `tolerance` always land in adjacent buckets
768/// (`|floor(a/tol) − floor(b/tol)| ≤ 1`), so probing `{b−1, b, b+1}` is a
769/// superset of every evaluate-match.
770fn numeric_index_key(v: f64, tolerance: f64) -> Option<ValueKey> {
771    if !tolerance.is_finite() || tolerance < 0.0 {
772        return None;
773    }
774    if tolerance == 0.0 {
775        let v = if v == 0.0 { 0.0_f64 } else { v };
776        return Some(ValueKey::FloatBits(v.to_bits()));
777    }
778    Some(ValueKey::Int(floor_to_i64(v / tolerance)))
779}
780
781fn numeric_probe_keys(v: f64, tolerance: f64) -> BTreeSet<ValueKey> {
782    match numeric_index_key(v, tolerance) {
783        None => BTreeSet::new(),
784        Some(k @ ValueKey::FloatBits(_)) => BTreeSet::from([k]),
785        Some(ValueKey::Int(b)) => BTreeSet::from([
786            ValueKey::Int(b.saturating_sub(1)),
787            ValueKey::Int(b),
788            ValueKey::Int(b.saturating_add(1)),
789        ]),
790        Some(other) => BTreeSet::from([other]),
791    }
792}
793
794fn geo_cell(lat: f64, lon: f64, km: f64) -> Option<(i64, i64, f64, i64)> {
795    if !km.is_finite() || km <= 0.0 {
796        return None;
797    }
798    let cell_deg = (km / 111.0).max(1e-6);
799    let gx = floor_to_i64(lat / cell_deg);
800    // Longitude wraps; lat does not (validated range, no pole crossing
801    // within the supported |lat|≲87 envelope — see cos clamp below).
802    let lon_cells = (360.0 / cell_deg).ceil() as i64;
803    let lon_cells = lon_cells.max(1);
804    let gy = floor_to_i64(lon / cell_deg).rem_euclid(lon_cells);
805    Some((gx, gy, cell_deg, lon_cells))
806}
807
808fn geo_index_key(lat: f64, lon: f64, km: f64) -> Option<ValueKey> {
809    let (gx, gy, _, _) = geo_cell(lat, lon, km)?;
810    Some(ValueKey::Str(format!("{gx}|{gy}")))
811}
812
813fn geo_probe_keys(lat: f64, lon: f64, km: f64) -> BTreeSet<ValueKey> {
814    let Some((gx, gy, cell_deg, lon_cells)) = geo_cell(lat, lon, km) else {
815        return BTreeSet::new();
816    };
817    // Cos clamp keeps the probe a superset up to |lat| ≈ 87.
818    let cos_lat = lat.to_radians().cos().max(0.05);
819    let n = ((km / (111.0 * cos_lat)) / cell_deg).ceil();
820    let n = if n.is_finite() {
821        floor_to_i64(n).max(0)
822    } else {
823        0
824    };
825    let mut out = BTreeSet::new();
826    for dx in -1..=1 {
827        for dy in -n..=n {
828            let cx = gx.saturating_add(dx);
829            let cy = gy.saturating_add(dy).rem_euclid(lon_cells);
830            out.insert(ValueKey::Str(format!("{cx}|{cy}")));
831        }
832    }
833    out
834}
835
836/// Vector candidates are a deliberate full scan of opposite-side
837/// vector-bearing nodes; ANN is Plan 8+.
838const SCAN_ALL_SENTINEL: ValueKey = ValueKey::Bool(true);
839
840/// IVF cluster buckets in `by_key`. SOH prefix keeps them off the Int space
841/// used by `NumericBucket` / integer `FieldEqual` and off token/geo Str keys.
842fn ivf_cluster_key(cluster: usize) -> ValueKey {
843    ValueKey::Str(format!("\u{1}ivf:{cluster}"))
844}
845
846/// `ScanAll` is the universe in an `Intersect`: skip it when another child
847/// has an index. Nested `Intersect` of only `ScanAll` is itself a universe.
848fn spec_is_scan_all_universe(spec: &CandidateSpec<'_>) -> bool {
849    match spec {
850        CandidateSpec::ScanAll { .. } => true,
851        CandidateSpec::Intersect(parts) => {
852            !parts.is_empty() && parts.iter().all(spec_is_scan_all_universe)
853        }
854        _ => false,
855    }
856}
857
858/// `ByKey` is resolved by `compute_desired` (FK id lookup), not `by_key`.
859/// Nested `Intersect` of only `ByKey` is likewise external.
860fn spec_is_bykey_external(spec: &CandidateSpec<'_>) -> bool {
861    match spec {
862        CandidateSpec::ByKey => true,
863        CandidateSpec::Intersect(parts) => {
864            !parts.is_empty() && parts.iter().all(spec_is_bykey_external)
865        }
866        _ => false,
867    }
868}
869
870impl SideIndex {
871    fn index_keys(spec: &CandidateSpec, get: &dyn Fn(&str) -> Option<Value>) -> BTreeSet<ValueKey> {
872        match spec {
873            CandidateSpec::ByKey => BTreeSet::new(),
874            CandidateSpec::Scalar { field } => get(field)
875                .as_ref()
876                .and_then(ValueKey::from_value)
877                .into_iter()
878                .collect(),
879            CandidateSpec::Tokens { field } => get(field)
880                .as_ref()
881                .and_then(list_tokens)
882                .unwrap_or_default(),
883            CandidateSpec::ScalarOrElements { field } => match get(field) {
884                Some(Value::List(items)) => items
885                    .iter()
886                    .take(MAX_KEYMATCH_LIST)
887                    .filter(|v| matches!(v, Value::Str(_)))
888                    .filter_map(ValueKey::from_value)
889                    .collect(),
890                Some(v) => ValueKey::from_value(&v).into_iter().collect(),
891                None => BTreeSet::new(),
892            },
893            CandidateSpec::NumericBucket { field, tolerance } => get(field)
894                .as_ref()
895                .and_then(as_finite_f64)
896                .and_then(|v| numeric_index_key(v, *tolerance))
897                .into_iter()
898                .collect(),
899            CandidateSpec::GeoGrid { field, km } => get(field)
900                .as_ref()
901                .and_then(as_latlon)
902                .and_then(|(lat, lon)| geo_index_key(lat, lon, *km))
903                .into_iter()
904                .collect(),
905            CandidateSpec::ScanAll { field } => get(field)
906                .as_ref()
907                .and_then(as_numeric_list)
908                .map(|_| SCAN_ALL_SENTINEL)
909                .into_iter()
910                .collect(),
911            // VectorClusters uses ivf_raw / ivf_clusters, not by_key. The
912            // insert() path returns early before reaching index_keys for this
913            // variant, so this arm is unreachable at runtime; it must be
914            // present to satisfy exhaustiveness.
915            CandidateSpec::VectorClusters { .. } => BTreeSet::new(),
916            // Hnsw uses the separate hnsw / hnsw_tracked fields, not by_key.
917            CandidateSpec::Hnsw { .. } => BTreeSet::new(),
918            // Union: each branch contributes its own index keys; the result is
919            // their union.  VectorClusters/Hnsw children are handled by the
920            // early-return in insert()/remove().
921            CandidateSpec::Union(specs) | CandidateSpec::Intersect(specs) => {
922                let mut out = BTreeSet::new();
923                for s in specs {
924                    out.extend(Self::index_keys(s, get));
925                }
926                out
927            }
928        }
929    }
930
931    fn probe_keys(spec: &CandidateSpec, get: &dyn Fn(&str) -> Option<Value>) -> BTreeSet<ValueKey> {
932        match spec {
933            CandidateSpec::ByKey
934            | CandidateSpec::Scalar { .. }
935            | CandidateSpec::Tokens { .. }
936            | CandidateSpec::ScalarOrElements { .. } => Self::index_keys(spec, get),
937            CandidateSpec::NumericBucket { field, tolerance } => get(field)
938                .as_ref()
939                .and_then(as_finite_f64)
940                .map(|v| numeric_probe_keys(v, *tolerance))
941                .unwrap_or_default(),
942            CandidateSpec::GeoGrid { field, km } => get(field)
943                .as_ref()
944                .and_then(as_latlon)
945                .map(|(lat, lon)| geo_probe_keys(lat, lon, *km))
946                .unwrap_or_default(),
947            CandidateSpec::ScanAll { field } => get(field)
948                .as_ref()
949                .and_then(as_numeric_list)
950                .map(|_| SCAN_ALL_SENTINEL)
951                .into_iter()
952                .collect(),
953            // VectorClusters probing is handled by ivf_candidates(), not probe_keys().
954            CandidateSpec::VectorClusters { .. } => BTreeSet::new(),
955            // Hnsw probing is handled by hnsw_candidates(), not probe_keys().
956            CandidateSpec::Hnsw { .. } => BTreeSet::new(),
957            // Union / Intersect: probe each child. `candidates()` intersects
958            // Intersect node-sets; mixing keys here is only for insert/remove.
959            CandidateSpec::Union(specs) | CandidateSpec::Intersect(specs) => {
960                let mut out = BTreeSet::new();
961                for s in specs {
962                    out.extend(Self::probe_keys(s, get));
963                }
964                out
965            }
966        }
967    }
968
969    pub fn insert(&mut self, spec: &CandidateSpec, node: u32, get: &dyn Fn(&str) -> Option<Value>) {
970        self.insert_with(spec, node, &HnswLeg::All, get);
971    }
972
973    /// `insert`, but skip the HNSW graph for ids in `already` — the open-time
974    /// scan's version, where the adopted graph is the base and the scan only
975    /// has to supply what the snapshot did not carry.
976    ///
977    /// `hnsw_tracked` is still recorded for every node, adopted or not: it is
978    /// the fallback candidate set and must cover the whole side.
979    pub fn insert_skipping(
980        &mut self,
981        spec: &CandidateSpec,
982        node: u32,
983        already: &BTreeSet<u32>,
984        get: &dyn Fn(&str) -> Option<Value>,
985    ) {
986        self.insert_with(spec, node, &HnswLeg::Skip(already), get);
987    }
988
989    /// `insert`, but the HNSW graph is left untouched — the sliced-build
990    /// version, where [`SideIndex::insert_hnsw_only`] supplies the vectors a
991    /// slice at a time.
992    ///
993    /// Every other leg of `spec` (by-key buckets, IVF, `ScanAll` metadata) is
994    /// filed exactly as `insert` files it, and `hnsw_tracked` is still
995    /// recorded, so the rule's non-vector state is whole from the moment it is
996    /// created.
997    pub fn insert_deferring_hnsw(
998        &mut self,
999        spec: &CandidateSpec,
1000        node: u32,
1001        get: &dyn Fn(&str) -> Option<Value>,
1002    ) {
1003        self.insert_with(spec, node, &HnswLeg::Defer, get);
1004    }
1005
1006    /// Insert `node` into the HNSW graph only, leaving every other leg of
1007    /// `spec` alone — the second half of [`SideIndex::insert_deferring_hnsw`].
1008    ///
1009    /// Returns `true` when a vector actually went into a graph, which is how a
1010    /// build slice counts what it has done.
1011    pub fn insert_hnsw_only(
1012        &mut self,
1013        spec: &CandidateSpec,
1014        node: u32,
1015        get: &dyn Fn(&str) -> Option<Value>,
1016    ) -> bool {
1017        match spec {
1018            CandidateSpec::Union(specs) | CandidateSpec::Intersect(specs) => {
1019                let mut any = false;
1020                for s in specs {
1021                    any |= self.insert_hnsw_only(s, node, get);
1022                }
1023                any
1024            }
1025            CandidateSpec::Hnsw { field, .. } => {
1026                let Some(xs) = get(field).as_ref().and_then(as_numeric_list) else {
1027                    return false;
1028                };
1029                self.record_vector_meta(node, &xs);
1030                self.hnsw_tracked.insert(node);
1031                if let Some(h) = &mut self.hnsw {
1032                    h.insert(node, &xs);
1033                }
1034                true
1035            }
1036            _ => false,
1037        }
1038    }
1039
1040    fn insert_with(
1041        &mut self,
1042        spec: &CandidateSpec,
1043        node: u32,
1044        leg: &HnswLeg<'_>,
1045        get: &dyn Fn(&str) -> Option<Value>,
1046    ) {
1047        // Union / Intersect: recurse into each child spec. insert() is
1048        // idempotent for ScanAll metadata (same-value overwrite).
1049        if let CandidateSpec::Union(specs) | CandidateSpec::Intersect(specs) = spec {
1050            for s in specs {
1051                self.insert_with(s, node, leg, get);
1052            }
1053            return;
1054        }
1055        // Hnsw: maintain hnsw_tracked for fallback, and hnsw graph if initialized.
1056        if let CandidateSpec::Hnsw { field, .. } = spec {
1057            if let Some(xs) = get(field).as_ref().and_then(as_numeric_list) {
1058                // An exact rule takes this arm from 0.6.6 on, and the
1059                // Cauchy-Schwarz early exit in `compute_desired` reads this
1060                // metadata, so it is recorded here as well as on `ScanAll`.
1061                self.record_vector_meta(node, &xs);
1062                self.hnsw_tracked.insert(node);
1063                match leg {
1064                    // The adopted graph already holds this vector, or the build
1065                    // is sliced and a later slice will supply it.
1066                    HnswLeg::Skip(already) if already.contains(&node) => return,
1067                    HnswLeg::Defer => return,
1068                    _ => {}
1069                }
1070                if let Some(h) = &mut self.hnsw {
1071                    h.insert(node, &xs);
1072                }
1073            }
1074            return;
1075        }
1076        // VectorClusters: IVF path — separate from the by_key / ScanAll path.
1077        if let CandidateSpec::VectorClusters { field, .. } = spec {
1078            if let Some(xs) = get(field).as_ref().and_then(as_numeric_list) {
1079                self.ivf_raw.insert(node, xs.clone());
1080                if !self.ivf_centroids.is_empty() {
1081                    // Assign in cosine space (centroids are unit-norm). Skip zeros.
1082                    if let Some(unit) = l2_normalize(&xs) {
1083                        let c = nearest_centroid(&self.ivf_centroids, &unit);
1084                        self.ivf_clusters.insert(node, c);
1085                        self.by_key
1086                            .entry(ivf_cluster_key(c))
1087                            .or_default()
1088                            .insert(node);
1089                    }
1090                    self.ivf_drift = self.ivf_drift.saturating_add(1);
1091                }
1092            }
1093            return;
1094        }
1095
1096        for k in Self::index_keys(spec, get) {
1097            self.by_key.entry(k).or_default().insert(node);
1098        }
1099        if let CandidateSpec::ScanAll { field } = spec {
1100            if let Some(xs) = get(field).as_ref().and_then(as_numeric_list) {
1101                self.record_vector_meta(node, &xs);
1102            }
1103        }
1104    }
1105
1106    /// File `node`'s `(dim, norm)`, suffix-norm checkpoints and anchor — the
1107    /// three inputs the Cauchy-Schwarz early exit reads.
1108    ///
1109    /// Called from the `ScanAll` and `Hnsw` arms of `insert_with` and from
1110    /// `insert_hnsw_only`, so an exact `VectorSimilar` rule keeps the early exit
1111    /// whichever arm files its vectors. Torn out by the matching arms of
1112    /// `remove`.
1113    fn record_vector_meta(&mut self, node: u32, xs: &[f64]) {
1114        let mut n2 = 0.0f64;
1115        for x in xs {
1116            n2 += x * x;
1117        }
1118        self.vec_meta.insert(node, (xs.len() as u32, n2.sqrt()));
1119        self.vec_checkpoints.insert(node, compute_ckpts(xs));
1120        // xs is non-empty (as_numeric_list rejects empty lists).
1121        self.vec_anchor.insert(node, xs[0]);
1122    }
1123
1124    /// Drop what [`SideIndex::record_vector_meta`] filed for `node`.
1125    fn forget_vector_meta(&mut self, node: u32) {
1126        self.vec_meta.remove(&node);
1127        self.vec_checkpoints.remove(&node);
1128        self.vec_anchor.remove(&node);
1129    }
1130
1131    pub fn remove(&mut self, spec: &CandidateSpec, node: u32, get: &dyn Fn(&str) -> Option<Value>) {
1132        // Union / Intersect: recurse into each child spec.
1133        if let CandidateSpec::Union(specs) | CandidateSpec::Intersect(specs) = spec {
1134            for s in specs {
1135                self.remove(s, node, get);
1136            }
1137            return;
1138        }
1139        // Hnsw: remove from hnsw_tracked and hnsw graph.  Increment ivf_drift
1140        // as a deletion counter so maybe_queue_ivf_rebuild fires at the same
1141        // cadence it did for IVF rules; the resulting rebuild re-scans all nodes
1142        // and optionally compacts the HNSW graph.
1143        if let CandidateSpec::Hnsw { field, .. } = spec {
1144            if get(field).as_ref().and_then(as_numeric_list).is_some() {
1145                self.forget_vector_meta(node);
1146                self.hnsw_tracked.remove(&node);
1147                if let Some(h) = &mut self.hnsw {
1148                    h.remove(node);
1149                }
1150                self.ivf_drift = self.ivf_drift.saturating_add(1);
1151            }
1152            return;
1153        }
1154        // VectorClusters: remove from ivf_raw and by_key cluster bucket.
1155        // Removal shifts cluster membership (the centroid stays but its member set
1156        // shrinks), which is a form of drift; increment the counter so callers can
1157        // decide when to trigger a rebuild.
1158        if let CandidateSpec::VectorClusters { .. } = spec {
1159            if self.ivf_raw.remove(&node).is_some() {
1160                self.ivf_drift = self.ivf_drift.saturating_add(1);
1161                if let Some(c) = self.ivf_clusters.remove(&node) {
1162                    let key = ivf_cluster_key(c);
1163                    if let Some(s) = self.by_key.get_mut(&key) {
1164                        s.remove(&node);
1165                        if s.is_empty() {
1166                            self.by_key.remove(&key);
1167                        }
1168                    }
1169                }
1170            }
1171            return;
1172        }
1173
1174        for k in Self::index_keys(spec, get) {
1175            if let Some(set) = self.by_key.get_mut(&k) {
1176                set.remove(&node);
1177                if set.is_empty() {
1178                    self.by_key.remove(&k);
1179                }
1180            }
1181        }
1182        if let CandidateSpec::ScanAll { field } = spec {
1183            if get(field).as_ref().and_then(as_numeric_list).is_some() {
1184                self.forget_vector_meta(node);
1185            }
1186        }
1187    }
1188
1189    /// Cached vector dimension for a `ScanAll` member, if present.
1190    pub fn vec_dim(&self, node: u32) -> Option<u32> {
1191        self.vec_meta.get(&node).map(|(d, _)| *d)
1192    }
1193
1194    /// Cached `(dim, L2 norm)` for tests / debug.
1195    pub fn vec_meta(&self, node: u32) -> Option<(u32, f64)> {
1196        self.vec_meta.get(&node).copied()
1197    }
1198
1199    /// Cached checkpoints for tests / debug.
1200    pub fn vec_ckpts(&self, node: u32) -> Option<&[f64; 8]> {
1201        self.vec_checkpoints.get(&node)
1202    }
1203
1204    /// Returns `(cached_norm, &checkpoints)` if the cached state matches the
1205    /// live vector under all three freshness checks.
1206    ///
1207    /// # Stale-cache gate
1208    ///
1209    /// Stale checkpoints (from a vector that differs from `live`) can produce
1210    /// **false rejects** — the Cauchy-Schwarz suffix bound may be under-tight
1211    /// for the live vector's actual energy distribution.  Three guards defend
1212    /// against this in ascending selectivity order:
1213    ///
1214    /// 1. **Dim check** — `cached_dim == live.len()`.  Different lengths →
1215    ///    immediate fallback.
1216    /// 2. **Norm check** — recomputes L2 norm with the same sequential
1217    ///    accumulation used at insert time so bits are identical for an unchanged
1218    ///    vector.  Changed norm → fallback.
1219    /// 3. **Anchor check** — compares `xs[0]` against the cached first element.
1220    ///    A permuted vector can share `(dim, norm)` with the indexed one but
1221    ///    differ at `xs[0]`, breaking the most realistic same-norm aliasing
1222    ///    attack.  This is **heuristic hardening**, not a proof: a permutation
1223    ///    that preserves `xs[0]` would still pass, but is vanishingly unlikely
1224    ///    in practice.
1225    ///
1226    /// The real coherence guarantee is structural: checkpoint rebuilds flow
1227    /// through the same insert/remove choke-points as `vec_meta`, so in
1228    /// normal single-writer operation the cache is always coherent.  These
1229    /// gates are belt-and-suspenders against bugs in those choke-points.
1230    pub(crate) fn fresh_ckpts_for<'a>(
1231        &'a self,
1232        node: u32,
1233        live: &[f64],
1234    ) -> Option<(f64, &'a [f64; 8])> {
1235        let &(dim, norm) = self.vec_meta.get(&node)?;
1236        if dim != live.len() as u32 {
1237            return None;
1238        }
1239        // Compute the live norm with the same sequential accumulation used at
1240        // insert time so the bits are identical when the vector is unchanged.
1241        let live_norm = {
1242            let mut n2 = 0.0f64;
1243            for x in live {
1244                n2 += x * x;
1245            }
1246            n2.sqrt()
1247        };
1248        if norm != live_norm {
1249            return None; // stale — fall back to brute-force evaluate()
1250        }
1251        // Heuristic anchor check: first element breaks same-norm permutation
1252        // aliasing in virtually all realistic cases.  dim > 0 guaranteed (dim
1253        // was stored from non-empty xs; live.len() == dim > 0).
1254        let live_anchor = live[0];
1255        let &cached_anchor = self.vec_anchor.get(&node)?;
1256        if live_anchor != cached_anchor {
1257            return None;
1258        }
1259        let ckpts = self.vec_checkpoints.get(&node)?;
1260        Some((norm, ckpts))
1261    }
1262
1263    pub fn candidates(
1264        &self,
1265        spec: &CandidateSpec,
1266        get: &dyn Fn(&str) -> Option<Value>,
1267    ) -> BTreeSet<u32> {
1268        // Hnsw: approximate nearest-neighbor search.
1269        if let CandidateSpec::Hnsw { field, k, floor } = spec {
1270            return self.hnsw_candidates(field, *k, *floor, get);
1271        }
1272        // VectorClusters: probe the P nearest centroids.
1273        if let CandidateSpec::VectorClusters { field, .. } = spec {
1274            return self.ivf_candidates(field, get);
1275        }
1276        // Union: take the union of candidates from each child spec.
1277        if let CandidateSpec::Union(specs) = spec {
1278            return specs.iter().flat_map(|s| self.candidates(s, get)).collect();
1279        }
1280        if let CandidateSpec::Intersect(specs) = spec {
1281            return self.intersect_candidates(specs, get);
1282        }
1283
1284        let mut out = BTreeSet::new();
1285        for k in Self::probe_keys(spec, get) {
1286            if let Some(set) = self.by_key.get(&k) {
1287                out.extend(set.iter().copied());
1288            }
1289        }
1290        // Exact: VectorSimilar evaluate is None when dims differ.
1291        if vector_dim_reject_enabled() {
1292            if let CandidateSpec::ScanAll { field } = spec {
1293                if let Some((dim, _)) = get(field).as_ref().and_then(vec_dim_norm) {
1294                    out.retain(|id| self.vec_meta.get(id).is_none_or(|(d, _)| *d == dim));
1295                }
1296            }
1297        }
1298        out
1299    }
1300
1301    /// Intersect child candidate sets. `ScanAll` is the universe (skipped);
1302    /// if every child is `ScanAll`, fall back to `ScanAll`. `ByKey` is resolved
1303    /// outside the index. Empty child → empty.
1304    fn intersect_candidates(
1305        &self,
1306        specs: &[CandidateSpec<'_>],
1307        get: &dyn Fn(&str) -> Option<Value>,
1308    ) -> BTreeSet<u32> {
1309        let mut restrictive = Vec::new();
1310        let mut scan_alls = Vec::new();
1311        for s in specs {
1312            if spec_is_scan_all_universe(s) {
1313                scan_alls.push(s);
1314            } else if spec_is_bykey_external(s) {
1315                continue;
1316            } else {
1317                restrictive.push(s);
1318            }
1319        }
1320        let to_intersect: &[&CandidateSpec<'_>] = if !restrictive.is_empty() {
1321            &restrictive
1322        } else if !scan_alls.is_empty() {
1323            &scan_alls
1324        } else {
1325            return BTreeSet::new();
1326        };
1327        let mut iter = to_intersect.iter();
1328        let Some(first) = iter.next() else {
1329            return BTreeSet::new();
1330        };
1331        let mut acc = self.candidates(first, get);
1332        if acc.is_empty() {
1333            return acc;
1334        }
1335        for s in iter {
1336            let other = self.candidates(s, get);
1337            if other.is_empty() {
1338                return BTreeSet::new();
1339            }
1340            acc = acc.intersection(&other).copied().collect();
1341            if acc.is_empty() {
1342                return acc;
1343            }
1344        }
1345        acc
1346    }
1347
1348    /// IVF candidate lookup: find the P nearest centroids to the query vector,
1349    /// return the union of their cluster members.
1350    fn ivf_candidates(&self, field: &str, get: &dyn Fn(&str) -> Option<Value>) -> BTreeSet<u32> {
1351        let Some(xs) = get(field).as_ref().and_then(as_numeric_list) else {
1352            return BTreeSet::new();
1353        };
1354        if self.ivf_centroids.is_empty() {
1355            // Not yet fitted (e.g. empty side at create time, or no data).
1356            // Fall back to full scan so early crash-recovery states don't drop recall
1357            // to zero when too few vectors were inserted for IVF to be meaningful.
1358            return self.ivf_raw.keys().copied().collect();
1359        }
1360        // When n ≤ k (actual centroid count), every node is its own centroid;
1361        // P probes only return the src's own cluster (which excludes itself),
1362        // yielding zero candidates. Full scan is correct and O(n) for these
1363        // tiny sets — this covers n < IVF_K_MIN and the exact n == k edge case.
1364        if self.ivf_raw.len() <= self.ivf_centroids.len() {
1365            return self.ivf_raw.keys().copied().collect();
1366        }
1367        let k = self.ivf_centroids.len();
1368        let p = probe_count(k);
1369
1370        // Probe in cosine space (same as centroid fit). Zero query → no candidates
1371        // (cosine with a zero vector is undefined; exact evaluate also returns None).
1372        let Some(xs) = l2_normalize(&xs) else {
1373            return BTreeSet::new();
1374        };
1375
1376        // Rank centroids by L2 distance to the unit query; take top-P.
1377        let mut dists: Vec<(usize, f64)> = self
1378            .ivf_centroids
1379            .iter()
1380            .enumerate()
1381            .map(|(i, c)| (i, l2_sq(&xs, c)))
1382            .collect();
1383        dists.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
1384
1385        let mut out = BTreeSet::new();
1386        for (ci, _) in dists.iter().take(p) {
1387            let key = ivf_cluster_key(*ci);
1388            if let Some(nodes) = self.by_key.get(&key) {
1389                out.extend(nodes.iter().copied());
1390            }
1391        }
1392        out
1393    }
1394
1395    /// Fit (or re-fit) the IVF k-means index for this side using all currently
1396    /// stored raw vectors.  Called by the engine after reindexing all nodes in
1397    /// `create_rule` and `rebuild`.
1398    ///
1399    /// `rule_name` is hashed via FNV-1a to produce a stable seed, ensuring the
1400    /// same rule+data always yields the same clusters (WAL replay identity).
1401    ///
1402    /// Clears all existing cluster assignments and by_key cluster entries, then
1403    /// assigns every non-zero vector (L2-normalized) to its nearest new centroid.
1404    /// Resets `ivf_drift` to zero.
1405    pub fn fit_ivf_clusters(&mut self, rule_name: &str) {
1406        if self.ivf_raw.is_empty() {
1407            self.ivf_centroids.clear();
1408            self.ivf_clusters.clear();
1409            self.ivf_drift = 0;
1410            return;
1411        }
1412
1413        // Clear old cluster → node mappings from by_key (namespaced IVF keys).
1414        for c in self.ivf_clusters.values() {
1415            self.by_key.remove(&ivf_cluster_key(*c));
1416        }
1417        self.ivf_clusters.clear();
1418
1419        // Gather vectors in deterministic order (BTreeMap → sorted by node id).
1420        let vecs: Vec<(u32, Vec<f64>)> = self
1421            .ivf_raw
1422            .iter()
1423            .map(|(&id, xs)| (id, xs.clone()))
1424            .collect();
1425
1426        let n = vecs.len();
1427        let k = cluster_k(n);
1428        let seed = fnv1a_u64(rule_name.as_bytes());
1429
1430        self.ivf_centroids = kmeans_fit(&vecs, k, seed);
1431
1432        // Assign in cosine space (skip zeros; they stay in ivf_raw but unclustered).
1433        for (node, xs) in &vecs {
1434            let Some(unit) = l2_normalize(xs) else {
1435                continue;
1436            };
1437            let c = nearest_centroid(&self.ivf_centroids, &unit);
1438            self.ivf_clusters.insert(*node, c);
1439            self.by_key
1440                .entry(ivf_cluster_key(c))
1441                .or_default()
1442                .insert(*node);
1443        }
1444        self.ivf_drift = 0;
1445    }
1446
1447    /// Number of fitted centroids (0 = not yet fitted).
1448    pub fn ivf_k(&self) -> usize {
1449        self.ivf_centroids.len()
1450    }
1451
1452    /// Cluster assignment for a node (None if not fitted or node not in index).
1453    pub fn ivf_cluster_of(&self, node: u32) -> Option<usize> {
1454        self.ivf_clusters.get(&node).copied()
1455    }
1456
1457    /// Export IVF state for snapshot persistence: (centroids, clusters, drift).
1458    ///
1459    /// The caller stores this in the V4 snapshot and passes it back to
1460    /// `load_ivf_state` on the next open, avoiding a full k-means re-fit.
1461    pub fn export_ivf_state(&self) -> (Vec<Vec<f64>>, BTreeMap<u32, usize>, u64) {
1462        (
1463            self.ivf_centroids.clone(),
1464            self.ivf_clusters.clone(),
1465            self.ivf_drift,
1466        )
1467    }
1468
1469    /// Restore IVF state from a V4 snapshot.
1470    ///
1471    /// This must be called AFTER the normal `insert()` pass (which populates
1472    /// `ivf_raw`) but INSTEAD OF `fit_ivf_clusters`.  It:
1473    ///   1. Removes any stale cluster-key entries from `by_key`.
1474    ///   2. Installs the persisted centroids and drift counter.
1475    ///   3. Rebuilds `by_key` cluster buckets from the persisted assignments.
1476    ///
1477    /// Nodes present in `ivf_raw` but absent from `clusters` (e.g. inserted
1478    /// post-snapshot via WAL replay before this is called) are left unassigned;
1479    /// `on_node_changed` will assign them to the nearest centroid incrementally.
1480    pub fn load_ivf_state(
1481        &mut self,
1482        centroids: Vec<Vec<f64>>,
1483        clusters: BTreeMap<u32, usize>,
1484        drift: u64,
1485    ) {
1486        // Precondition: ivf_clusters is empty when called from reindex_all_load_state (indexes reset to default); loop is defensive for any future direct-call path.
1487        // Remove old cluster bucket entries from by_key.
1488        for c in self.ivf_clusters.values() {
1489            self.by_key.remove(&ivf_cluster_key(*c));
1490        }
1491        self.ivf_clusters.clear();
1492
1493        self.ivf_centroids = centroids;
1494        self.ivf_drift = drift;
1495
1496        // Rebuild by_key from persisted assignments (only for nodes still in ivf_raw).
1497        for (&node, &c) in &clusters {
1498            if !self.ivf_raw.contains_key(&node) {
1499                // Node was removed post-snapshot (WAL replay deleted it).  Skip.
1500                continue;
1501            }
1502            self.ivf_clusters.insert(node, c);
1503            self.by_key
1504                .entry(ivf_cluster_key(c))
1505                .or_default()
1506                .insert(node);
1507        }
1508    }
1509
1510    // -----------------------------------------------------------------------
1511    // HNSW methods
1512    // -----------------------------------------------------------------------
1513
1514    /// Initialise the HNSW graph for this side, seeding it with `FNV-1a(rule_name)`.
1515    ///
1516    /// Must be called before inserting nodes via `CandidateSpec::Hnsw`.
1517    /// Idempotent: calling again with the same name replaces the existing graph.
1518    pub fn init_hnsw(&mut self, rule_name: &str) {
1519        let seed = fnv1a_u64(rule_name.as_bytes());
1520        self.hnsw = Some(HnswIndex::new(seed));
1521    }
1522
1523    /// HNSW candidate lookup: `k`-nearest-neighbor search using the built graph.
1524    ///
1525    /// With `floor` of `None` — an approximate rule — this is one beam pass at
1526    /// `k`, which is what it has always been.
1527    ///
1528    /// With `floor` of `Some(min)` — an exact rule, 0.6.6 on — the beam widens,
1529    /// and **there is exactly one way it is allowed to answer**: a beam that
1530    /// came back full (`hits.len() == ef`) whose worst hit is below `min` by more
1531    /// than [`BEAM_FLOOR_SLACK`] — the beam answers in `f32`, and the slack keeps
1532    /// that rounding on the side of widening. Such a beam has proved what it did
1533    /// not return — every node it rejected is farther from the query than one
1534    /// already known to fail the predicate — so its hits are the candidate set.
1535    /// The similarities themselves are discarded here; `compute_desired` rescores
1536    /// every candidate from the `f64` store, so `min` is only ever *decided* in
1537    /// `f64`.
1538    ///
1539    /// Every other outcome hands back the whole tracked set, which is the
1540    /// pre-0.6.6 exact candidate set:
1541    ///
1542    /// * **The beam came back short of its own width.** Layer 0 need not be one
1543    ///   connected component — a corpus of identical or near-identical vectors
1544    ///   is the case that shows it — and a beam that exhausted its frontier has
1545    ///   proved nothing about the nodes it could not reach.
1546    /// * **The ceiling ([`ef_max`], [`EF_MAX`] by default) was reached with the
1547    ///   worst hit still at or above `min`.** A cluster denser than the ceiling
1548    ///   then costs a scan; it never costs recall.
1549    /// * **A beam as wide as the index itself** (`ef >= h.len()`), where walking
1550    ///   the graph cannot beat handing back every vector on the side.
1551    /// * **The index cannot answer this query at all**
1552    ///   ([`HnswIndex::can_answer`]): no graph, an empty one, a stride that is not
1553    ///   the query's dimension, or one that refused a vector it was handed. This
1554    ///   is checked *before* any beam runs, because a beam over an index that is
1555    ///   missing part of the corpus would "prove" its floor against vectors the
1556    ///   index never held. It applies to the approximate path as well.
1557    ///
1558    /// So the index is a candidate *generator* here and never a silent filter:
1559    /// the only way a node is dropped is a beam that proved it is below `min`.
1560    fn hnsw_candidates(
1561        &self,
1562        field: &str,
1563        k: usize,
1564        floor: Option<f64>,
1565        get: &dyn Fn(&str) -> Option<Value>,
1566    ) -> BTreeSet<u32> {
1567        let Some(xs) = get(field).as_ref().and_then(as_numeric_list) else {
1568            return BTreeSet::new();
1569        };
1570        if let Some(h) = &self.hnsw {
1571            // `can_answer`, not `!is_empty()`: an index that refused a vector, or
1572            // whose stride is not this query's dimension — a 3-element stray
1573            // ingested ahead of the real corpus elects one — is non-empty and
1574            // still cannot supply the candidates this query needs. It is asked
1575            // before any beam, because a beam over such an index would "conclude"
1576            // from an incomplete corpus.
1577            if h.can_answer(xs.len()) {
1578                let Some(min) = floor else {
1579                    return h.search(&xs, k).into_iter().map(|(id, _)| id).collect();
1580                };
1581                let cap = ef_max();
1582                let mut ef = h.ef_for(k);
1583                while ef < h.len() {
1584                    // `k = ef`: the answer is every hit above the floor, so
1585                    // truncating the beam to `k` would be throwing away the
1586                    // very candidates the widening is looking for.
1587                    let hits = h.search_with_ef(&xs, ef, ef);
1588                    // `search` sorts descending, so the last hit is the worst.
1589                    let full = hits.len() == ef;
1590                    if full && hits[hits.len() - 1].1 < min - BEAM_FLOOR_SLACK {
1591                        return hits.into_iter().map(|(id, _)| id).collect();
1592                    }
1593                    // Short of its width (frontier exhausted, so a wider beam
1594                    // reaches nothing new) or at the ceiling with the floor
1595                    // still unreached: neither has proved anything about what it
1596                    // did not return.
1597                    if !full || ef >= cap {
1598                        break;
1599                    }
1600                    ef = ef.saturating_mul(2);
1601                }
1602            }
1603        }
1604        // Fallback: full scan of all tracked nodes (superset of true positives).
1605        self.hnsw_tracked.clone()
1606    }
1607
1608    /// Export the HNSW graph as an opaque versioned blob.
1609    ///
1610    /// Returns an empty `Vec` when the HNSW is not initialized.
1611    ///
1612    /// `complete` is false when the rule's sliced build still owes this side
1613    /// vectors; it rides in the blob so that a reader opening the snapshot
1614    /// knows the graph is a prefix and takes its exhaustive path rather than
1615    /// answering confidently about a fraction of the corpus. The engine reads
1616    /// it from `pending_builds`, which is not itself persisted.
1617    pub fn export_hnsw_blob(&self, complete: bool) -> Vec<u8> {
1618        self.hnsw
1619            .as_ref()
1620            .and_then(|h| crate::hnsw::encode_hnsw_blob(h, complete))
1621            .unwrap_or_default()
1622    }
1623
1624    /// Restore the HNSW graph from a previously exported blob.
1625    ///
1626    /// The `hnsw_tracked` set is populated from the restored graph's node ids
1627    /// so candidates/remove work correctly after restore.
1628    /// Silently ignores empty, corrupt, or unknown-version blobs (the HNSW
1629    /// stays uninitialized and the side keeps its full-scan fallback).
1630    pub fn load_hnsw_blob(&mut self, blob: &[u8]) {
1631        if let Ok(h) = crate::hnsw::decode_hnsw_blob(blob) {
1632            self.adopt_hnsw(h);
1633        }
1634    }
1635
1636    /// Initialise this side's HNSW graph, adopting `blob` when it holds one.
1637    ///
1638    /// Returns the node ids the adopted graph already contains, so an open-time
1639    /// scan can skip re-inserting them. An empty, corrupt, or unknown-version
1640    /// blob yields an empty graph and an empty set — exactly what `init_hnsw`
1641    /// gives today — and the scan then builds the graph as it always did.
1642    ///
1643    /// `true` in the second slot means "this side was adopted, not built", which
1644    /// is what the caller counts as a skipped build.
1645    pub fn init_or_adopt_hnsw(&mut self, rule_name: &str, blob: &[u8]) -> (BTreeSet<u32>, bool) {
1646        self.hnsw = None;
1647        if !blob.is_empty() {
1648            match crate::hnsw::decode_hnsw_blob(blob) {
1649                Ok(h) => self.adopt_hnsw(h),
1650                Err(e) => eprintln!(
1651                    "[mushroomdb] rule {rule_name:?}: a persisted HNSW index failed to load \
1652                     ({e}); rebuilding it from the node scan"
1653                ),
1654            }
1655        }
1656        match &self.hnsw {
1657            // `accounted_ids`, not `node_ids`: an adopted v4 graph carries the
1658            // vectors it parked and the ids it refused, and re-offering either
1659            // undoes the state the blob just restored — a parked vector is
1660            // superseded and then refused, which loses it for good. On a pre-v4
1661            // graph both sets are empty and this is `node_ids()` exactly, which
1662            // is what keeps the older blobs re-deriving as they always did.
1663            Some(h) => (h.accounted_ids(), true),
1664            None => {
1665                self.init_hnsw(rule_name);
1666                (BTreeSet::new(), false)
1667            }
1668        }
1669    }
1670
1671    /// Install an already-deserialized HNSW graph, replacing any existing one.
1672    ///
1673    /// `hnsw_tracked` is repopulated from the graph's node ids so candidates
1674    /// and removal work against the installed graph rather than whatever the
1675    /// preceding node scan happened to record.
1676    pub fn adopt_hnsw(&mut self, mut h: HnswIndex) {
1677        // Completeness of a live index is `RuleEngine::pending_builds`, not
1678        // this flag: an unfinished build is skipped in `hnsw_search_dst`
1679        // whether the adopted graph still carries `incomplete` or not. Clearing
1680        // it here keeps the flag for the lazily-decoded read-path copy, which
1681        // has no `pending_builds` behind it. A complete blob's missing nodes
1682        // (writes after the snapshot) are still supplied by the open-time
1683        // scan; an incomplete blob's remainder is left to `pump_index_build`.
1684        h.mark_complete();
1685        self.hnsw_tracked = h.node_ids();
1686        self.hnsw = Some(h);
1687    }
1688
1689    /// True when the HNSW graph has been initialized and contains at least one node.
1690    pub fn has_hnsw(&self) -> bool {
1691        self.hnsw.as_ref().is_some_and(|h| !h.is_empty())
1692    }
1693
1694    /// Borrow the HNSW index, if initialized.
1695    pub fn hnsw_ref(&self) -> Option<&HnswIndex> {
1696        self.hnsw.as_ref()
1697    }
1698
1699    /// Remove and return this side's HNSW graph, leaving the side without one.
1700    ///
1701    /// Lets a caller that is about to reset the whole `SideIndex` carry the
1702    /// graph across — the graph is the expensive part and is not always worth
1703    /// rebuilding.
1704    pub fn take_hnsw(&mut self) -> Option<HnswIndex> {
1705        self.hnsw.take()
1706    }
1707}
1708
1709#[cfg(test)]
1710mod tests {
1711    use super::*;
1712    use crate::def::Predicate;
1713    use core_storage::Value;
1714    use std::collections::{BTreeMap, HashMap};
1715
1716    fn getter(map: &HashMap<String, Value>) -> impl Fn(&str) -> Option<Value> + '_ {
1717        move |f: &str| map.get(f).cloned()
1718    }
1719
1720    #[test]
1721    fn kmeans_centroids_are_unit_norm() {
1722        let vecs = vec![(0, vec![3.0, 0.0, 0.0]), (1, vec![0.0, 4.0, 0.0])];
1723        let cents = kmeans_fit(&vecs, 2, 1);
1724        for c in cents {
1725            let n = c.iter().map(|x| x * x).sum::<f64>().sqrt();
1726            assert!((n - 1.0).abs() < 1e-9, "{n}");
1727        }
1728    }
1729
1730    /// Raw L2 would put `[3,0,0]` on a nearby large centroid while cosine (and
1731    /// the unit vector `[1,0,0]`) prefer the x-axis centroid. Assignment must
1732    /// L2-normalize first so scale-equivalent vectors share a cluster.
1733    ///
1734    /// Tests IVF directly (via `CandidateSpec::VectorClusters`) since
1735    /// `candidate_spec_approx` now returns `CandidateSpec::Hnsw`.
1736    #[test]
1737    fn scaled_vector_joins_same_ivf_cluster_as_unit() {
1738        // Use VectorClusters directly to test IVF cluster assignment.
1739        let spec = CandidateSpec::VectorClusters {
1740            field: "emb",
1741            min: 0.5,
1742        };
1743        let mut idx = SideIndex::default();
1744        idx.load_ivf_state(
1745            vec![vec![1.0, 0.0, 0.0], vec![2.5, 0.1, 0.0]],
1746            BTreeMap::new(),
1747            0,
1748        );
1749        idx.insert(&spec, 1, &getter(&emb(&[1.0, 0.0, 0.0])));
1750        idx.insert(&spec, 2, &getter(&emb(&[3.0, 0.0, 0.0])));
1751        assert_eq!(
1752            idx.ivf_cluster_of(1),
1753            idx.ivf_cluster_of(2),
1754            "scale-equivalent vectors must share an IVF cluster; got {:?} vs {:?}",
1755            idx.ivf_cluster_of(1),
1756            idx.ivf_cluster_of(2)
1757        );
1758        assert_eq!(idx.ivf_cluster_of(1), Some(0));
1759    }
1760
1761    #[test]
1762    fn scalar_index_buckets_by_value() {
1763        let pred = Predicate::FieldEqual {
1764            field: "ind".into(),
1765        };
1766        let spec = candidate_spec(&pred);
1767        let mut idx = SideIndex::default();
1768        let a: HashMap<_, _> = [("ind".to_string(), Value::Str("arch".into()))].into();
1769        let b: HashMap<_, _> = [("ind".to_string(), Value::Str("law".into()))].into();
1770        idx.insert(&spec, 1, &getter(&a));
1771        idx.insert(&spec, 2, &getter(&b));
1772        idx.insert(&spec, 3, &getter(&a));
1773        let c = idx.candidates(&spec, &getter(&a));
1774        assert_eq!(c.into_iter().collect::<Vec<_>>(), vec![1, 3]);
1775        idx.remove(&spec, 3, &getter(&a));
1776        assert_eq!(idx.candidates(&spec, &getter(&a)).len(), 1);
1777        // node without the field indexes nothing and matches nothing
1778        let empty: HashMap<String, Value> = HashMap::new();
1779        idx.insert(&spec, 9, &getter(&empty));
1780        assert!(idx.candidates(&spec, &getter(&empty)).is_empty());
1781    }
1782
1783    #[test]
1784    fn token_index_unions_buckets() {
1785        let mk =
1786            |items: &[&str]| Value::List(items.iter().map(|s| Value::Str((*s).into())).collect());
1787        let pred = Predicate::Overlap {
1788            field: "tags".into(),
1789            min: 0.5,
1790        };
1791        let spec = candidate_spec(&pred);
1792        let mut idx = SideIndex::default();
1793        let a: HashMap<_, _> = [("tags".to_string(), mk(&["x", "y"]))].into();
1794        let b: HashMap<_, _> = [("tags".to_string(), mk(&["y", "z"]))].into();
1795        let c: HashMap<_, _> = [("tags".to_string(), mk(&["q"]))].into();
1796        idx.insert(&spec, 1, &getter(&a));
1797        idx.insert(&spec, 2, &getter(&b));
1798        idx.insert(&spec, 3, &getter(&c));
1799        let probe: HashMap<_, _> = [("tags".to_string(), mk(&["y"]))].into();
1800        assert_eq!(
1801            idx.candidates(&spec, &getter(&probe))
1802                .into_iter()
1803                .collect::<Vec<_>>(),
1804            vec![1, 2]
1805        );
1806        idx.remove(&spec, 2, &getter(&b));
1807        assert_eq!(
1808            idx.candidates(&spec, &getter(&probe))
1809                .into_iter()
1810                .collect::<Vec<_>>(),
1811            vec![1]
1812        );
1813    }
1814
1815    #[test]
1816    fn all_intersects_parts_and_bykey_indexes_nothing() {
1817        let all = Predicate::All(vec![
1818            Predicate::FieldEqual {
1819                field: "ind".into(),
1820            },
1821            Predicate::Overlap {
1822                field: "tags".into(),
1823                min: 0.5,
1824            },
1825        ]);
1826        match candidate_spec(&all) {
1827            CandidateSpec::Intersect(v) => assert_eq!(v.len(), 2),
1828            other => panic!("{other:?}"),
1829        }
1830        let km = Predicate::KeyMatch { field: "fk".into() };
1831        assert!(matches!(candidate_spec(&km), CandidateSpec::ByKey));
1832        let mut idx = SideIndex::default();
1833        let a: HashMap<_, _> = [("fk".to_string(), Value::Str("c1".into()))].into();
1834        idx.insert(&candidate_spec(&km), 1, &getter(&a));
1835        assert!(idx.candidates(&candidate_spec(&km), &getter(&a)).is_empty());
1836    }
1837
1838    fn year(v: Value) -> HashMap<String, Value> {
1839        [("year".to_string(), v)].into()
1840    }
1841
1842    fn loc(lat: f64, lon: f64) -> HashMap<String, Value> {
1843        [(
1844            "loc".to_string(),
1845            Value::List(vec![Value::Float(lat), Value::Float(lon)]),
1846        )]
1847        .into()
1848    }
1849
1850    fn emb(vals: &[f64]) -> HashMap<String, Value> {
1851        [(
1852            "emb".to_string(),
1853            Value::List(vals.iter().copied().map(Value::Float).collect()),
1854        )]
1855        .into()
1856    }
1857
1858    fn bucket_int(spec: &CandidateSpec, map: &HashMap<String, Value>) -> Option<i64> {
1859        match SideIndex::index_keys(spec, &getter(map)).into_iter().next() {
1860            Some(ValueKey::Int(b)) => Some(b),
1861            _ => None,
1862        }
1863    }
1864
1865    #[test]
1866    fn numeric_bucket_adjacency_and_far_value() {
1867        let pred = Predicate::NumericWithin {
1868            field: "year".into(),
1869            tolerance: 2.0,
1870        };
1871        let spec = candidate_spec(&pred);
1872        assert!(matches!(
1873            spec,
1874            CandidateSpec::NumericBucket {
1875                field: "year",
1876                tolerance
1877            } if tolerance == 2.0
1878        ));
1879
1880        let v10 = year(Value::Float(10.0));
1881        let v119 = year(Value::Float(11.9));
1882        let v99 = year(Value::Float(9.9));
1883        let v141 = year(Value::Float(14.1));
1884
1885        let b10 = bucket_int(&spec, &v10).unwrap();
1886        let b119 = bucket_int(&spec, &v119).unwrap();
1887        let b99 = bucket_int(&spec, &v99).unwrap();
1888        // 10.0 and 11.9 share a bucket; 9.9 is adjacent (forces ±1 probe).
1889        assert!((b10 - b119).abs() <= 1);
1890        assert!((b10 - b99).abs() <= 1);
1891
1892        let mut idx = SideIndex::default();
1893        idx.insert(&spec, 1, &getter(&v10));
1894        idx.insert(&spec, 2, &getter(&v119));
1895        idx.insert(&spec, 3, &getter(&v141));
1896        idx.insert(&spec, 4, &getter(&v99));
1897        let hits = idx.candidates(&spec, &getter(&v10));
1898        assert_eq!(hits.into_iter().collect::<Vec<_>>(), vec![1, 2, 4]);
1899    }
1900
1901    #[test]
1902    fn numeric_tol_zero_int_float_collide() {
1903        let pred = Predicate::NumericWithin {
1904            field: "year".into(),
1905            tolerance: 0.0,
1906        };
1907        let spec = candidate_spec(&pred);
1908        let mut idx = SideIndex::default();
1909        idx.insert(&spec, 1, &getter(&year(Value::Int(2))));
1910        assert_eq!(
1911            idx.candidates(&spec, &getter(&year(Value::Float(2.0))))
1912                .into_iter()
1913                .collect::<Vec<_>>(),
1914            vec![1]
1915        );
1916        assert!(idx
1917            .candidates(&spec, &getter(&year(Value::Float(2.1))))
1918            .is_empty());
1919    }
1920
1921    #[test]
1922    fn numeric_tol_zero_signed_zero_collides() {
1923        let pred = Predicate::NumericWithin {
1924            field: "year".into(),
1925            tolerance: 0.0,
1926        };
1927        let spec = candidate_spec(&pred);
1928        let neg = year(Value::Float(-0.0));
1929        let pos = year(Value::Float(0.0));
1930        let mut idx = SideIndex::default();
1931        idx.insert(&spec, 1, &getter(&neg));
1932        assert_eq!(
1933            idx.candidates(&spec, &getter(&pos))
1934                .into_iter()
1935                .collect::<Vec<_>>(),
1936            vec![1]
1937        );
1938        let mut idx2 = SideIndex::default();
1939        idx2.insert(&spec, 2, &getter(&pos));
1940        assert_eq!(
1941            idx2.candidates(&spec, &getter(&neg))
1942                .into_iter()
1943                .collect::<Vec<_>>(),
1944            vec![2]
1945        );
1946    }
1947
1948    #[test]
1949    fn geo_grid_same_cell_cross_cell_and_far_city() {
1950        let pred = Predicate::GeoRadius {
1951            field: "loc".into(),
1952            km: 400.0,
1953        };
1954        let spec = candidate_spec(&pred);
1955        assert!(matches!(
1956            spec,
1957            CandidateSpec::GeoGrid {
1958                field: "loc",
1959                km
1960            } if km == 400.0
1961        ));
1962
1963        let paris = loc(48.8566, 2.3522);
1964        let london = loc(51.5074, -0.1278);
1965        let nearby = loc(48.9, 2.4); // same cell as Paris at km=400
1966        let ny = loc(40.7128, -74.0060);
1967
1968        let mut idx = SideIndex::default();
1969        idx.insert(&spec, 1, &getter(&paris));
1970        idx.insert(&spec, 2, &getter(&london));
1971        idx.insert(&spec, 3, &getter(&nearby));
1972        idx.insert(&spec, 4, &getter(&ny));
1973
1974        let from_paris = idx.candidates(&spec, &getter(&paris));
1975        assert!(from_paris.contains(&1), "same-cell self");
1976        assert!(from_paris.contains(&3), "same-cell neighbor");
1977        assert!(from_paris.contains(&2), "cross-cell Paris↔London ~343.5 km");
1978        assert!(!from_paris.contains(&4), "New York not in 400 km probe");
1979    }
1980
1981    #[test]
1982    fn geo_grid_high_latitude_probe_is_superset() {
1983        let pred = Predicate::GeoRadius {
1984            field: "loc".into(),
1985            km: 340.0,
1986        };
1987        let spec = candidate_spec(&pred);
1988        let reyk = loc(64.1466, -21.9426);
1989        let lat = 64.0_f64;
1990        let dlon = 300.0 / (111.0 * lat.to_radians().cos());
1991        let east = loc(lat, -21.9426 + dlon);
1992
1993        let mut idx = SideIndex::default();
1994        idx.insert(&spec, 1, &getter(&reyk));
1995        idx.insert(&spec, 2, &getter(&east));
1996        let hits = idx.candidates(&spec, &getter(&reyk));
1997        assert!(
1998            hits.contains(&2),
1999            "300 km east of Reykjavik must stay in the high-lat probe"
2000        );
2001    }
2002
2003    #[test]
2004    fn geo_grid_antimeridian_wrap_and_evaluate_agree() {
2005        let pred = Predicate::GeoRadius {
2006            field: "loc".into(),
2007            km: 400.0,
2008        };
2009        let spec = candidate_spec(&pred);
2010        let east = loc(70.0, 179.9);
2011        let west = loc(70.0, -179.9);
2012
2013        let mut idx = SideIndex::default();
2014        idx.insert(&spec, 1, &getter(&east));
2015        assert!(
2016            idx.candidates(&spec, &getter(&west)).contains(&1),
2017            "±180 pair at lat 70 must land in the wrapped probe"
2018        );
2019
2020        let sp = |f: &str| east.get(f).cloned();
2021        let dp = |f: &str| west.get(f).cloned();
2022        let score = crate::def::evaluate(
2023            &pred,
2024            &crate::def::NodeView {
2025                key: "e",
2026                props: &sp,
2027            },
2028            &crate::def::NodeView {
2029                key: "w",
2030                props: &dp,
2031            },
2032        );
2033        assert!(
2034            score.is_some(),
2035            "haversine must match across the antimeridian"
2036        );
2037
2038        // Wrap must not alias distant longitudes into the Paris probe.
2039        let paris = loc(48.8566, 2.3522);
2040        let ny = loc(40.7128, -74.0060);
2041        let mut idx2 = SideIndex::default();
2042        idx2.insert(&spec, 4, &getter(&ny));
2043        assert!(
2044            !idx2.candidates(&spec, &getter(&paris)).contains(&4),
2045            "New York still not in the Paris probe after wrap"
2046        );
2047    }
2048
2049    #[test]
2050    fn scan_all_returns_vector_nodes_skips_malformed() {
2051        let pred = Predicate::VectorSimilar {
2052            field: "emb".into(),
2053            min: 0.5,
2054        };
2055        let spec = candidate_spec(&pred);
2056        assert!(matches!(spec, CandidateSpec::ScanAll { field: "emb" }));
2057
2058        let mut idx = SideIndex::default();
2059        idx.insert(&spec, 1, &getter(&emb(&[1.0, 0.0])));
2060        idx.insert(&spec, 2, &getter(&emb(&[0.0, 1.0])));
2061        idx.insert(&spec, 3, &getter(&emb(&[1.0, 2.0, 3.0])));
2062        let empty: HashMap<_, _> = [("emb".to_string(), Value::List(vec![]))].into();
2063        let text: HashMap<_, _> =
2064            [("emb".to_string(), Value::List(vec![Value::Str("x".into())]))].into();
2065        let missing: HashMap<String, Value> = HashMap::new();
2066        idx.insert(&spec, 4, &getter(&empty));
2067        idx.insert(&spec, 5, &getter(&text));
2068        idx.insert(&spec, 6, &getter(&missing));
2069
2070        let hits = idx.candidates(&spec, &getter(&emb(&[1.0, 0.0])));
2071        assert_eq!(
2072            hits.into_iter().collect::<Vec<_>>(),
2073            vec![1, 2],
2074            "dim-2 probe must drop the dim-3 member"
2075        );
2076        assert_eq!(
2077            idx.candidates(&spec, &getter(&emb(&[1.0, 2.0, 3.0])))
2078                .into_iter()
2079                .collect::<Vec<_>>(),
2080            vec![3]
2081        );
2082        with_vector_dim_reject(false, || {
2083            assert_eq!(
2084                idx.candidates(&spec, &getter(&emb(&[1.0, 0.0])))
2085                    .into_iter()
2086                    .collect::<Vec<_>>(),
2087                vec![1, 2, 3],
2088                "unfiltered ScanAll still returns every vector node"
2089            );
2090        });
2091        assert_eq!(idx.vec_dim(1), Some(2));
2092        assert_eq!(idx.vec_dim(3), Some(3));
2093        assert!(idx.vec_meta(1).is_some());
2094        assert!(idx.vec_dim(4).is_none());
2095        assert!(idx.candidates(&spec, &getter(&empty)).is_empty());
2096        assert!(idx.candidates(&spec, &getter(&text)).is_empty());
2097        assert!(idx.candidates(&spec, &getter(&missing)).is_empty());
2098        idx.remove(&spec, 1, &getter(&emb(&[1.0, 0.0])));
2099        assert!(idx.vec_dim(1).is_none());
2100    }
2101
2102    #[test]
2103    fn legacy_specs_probe_keys_equal_index_keys() {
2104        let a: HashMap<_, _> = [
2105            ("ind".to_string(), Value::Str("arch".into())),
2106            (
2107                "tags".to_string(),
2108                Value::List(vec![Value::Str("x".into()), Value::Str("y".into())]),
2109            ),
2110            ("fk".to_string(), Value::Str("c1".into())),
2111        ]
2112        .into();
2113        let get = getter(&a);
2114        for pred in [
2115            Predicate::KeyMatch { field: "fk".into() },
2116            Predicate::FieldEqual {
2117                field: "ind".into(),
2118            },
2119            Predicate::Overlap {
2120                field: "tags".into(),
2121                min: 0.5,
2122            },
2123        ] {
2124            let spec = candidate_spec(&pred);
2125            assert_eq!(
2126                SideIndex::index_keys(&spec, &get),
2127                SideIndex::probe_keys(&spec, &get)
2128            );
2129        }
2130    }
2131
2132    #[test]
2133    fn all_vector_then_field_equal_does_not_scan_all() {
2134        let p = Predicate::All(vec![
2135            Predicate::VectorSimilar {
2136                field: "e".into(),
2137                min: 0.8,
2138            },
2139            Predicate::FieldEqual {
2140                field: "industry".into(),
2141            },
2142        ]);
2143        match candidate_spec(&p) {
2144            CandidateSpec::Intersect(v) => assert_eq!(v.len(), 2),
2145            other => panic!("{other:?}"),
2146        }
2147
2148        let spec = candidate_spec(&p);
2149        let mut idx = SideIndex::default();
2150        let mk = |industry: &str, e: &[f64]| {
2151            [
2152                ("industry".to_string(), Value::Str(industry.into())),
2153                (
2154                    "e".to_string(),
2155                    Value::List(e.iter().copied().map(Value::Float).collect()),
2156                ),
2157            ]
2158            .into()
2159        };
2160        let same: HashMap<_, _> = mk("tech", &[1.0, 0.0]);
2161        let other_ind: HashMap<_, _> = mk("law", &[1.0, 0.0]);
2162        let no_vec: HashMap<_, _> = [("industry".to_string(), Value::Str("tech".into()))].into();
2163        idx.insert(&spec, 1, &getter(&same));
2164        idx.insert(&spec, 2, &getter(&other_ind));
2165        idx.insert(&spec, 3, &getter(&no_vec));
2166
2167        let hits = idx.candidates(&spec, &getter(&same));
2168        assert!(hits.contains(&1), "matching industry must stay a candidate");
2169        assert!(
2170            !hits.contains(&2),
2171            "different industry must not be scanned in via VectorSimilar"
2172        );
2173        assert!(
2174            hits.contains(&3),
2175            "ScanAll is universe: extra Scalar-only candidates are allowed"
2176        );
2177
2178        let empty_ind: HashMap<_, _> = mk("finance", &[1.0, 0.0]);
2179        assert!(
2180            idx.candidates(&spec, &getter(&empty_ind)).is_empty(),
2181            "empty Scalar child → empty intersect"
2182        );
2183    }
2184
2185    #[test]
2186    fn all_approx_vector_then_field_equal_is_intersect() {
2187        let p = Predicate::All(vec![
2188            Predicate::VectorSimilar {
2189                field: "e".into(),
2190                min: 0.8,
2191            },
2192            Predicate::FieldEqual {
2193                field: "industry".into(),
2194            },
2195        ]);
2196        match candidate_spec_approx(&p) {
2197            CandidateSpec::Intersect(v) => assert_eq!(v.len(), 2),
2198            other => panic!("{other:?}"),
2199        }
2200
2201        let spec = candidate_spec_approx(&p);
2202        let mut idx = SideIndex::default();
2203        // Initialize HNSW so insertions populate the graph.
2204        idx.init_hnsw("test-rule");
2205        let mk = |industry: &str, e: &[f64]| {
2206            [
2207                ("industry".to_string(), Value::Str(industry.into())),
2208                (
2209                    "e".to_string(),
2210                    Value::List(e.iter().copied().map(Value::Float).collect()),
2211                ),
2212            ]
2213            .into()
2214        };
2215        let same: HashMap<_, _> = mk("tech", &[1.0, 0.0]);
2216        let other_ind: HashMap<_, _> = mk("law", &[1.0, 0.0]);
2217        idx.insert(&spec, 1, &getter(&same));
2218        idx.insert(&spec, 2, &getter(&other_ind));
2219        let hits = idx.candidates(&spec, &getter(&same));
2220        assert!(hits.contains(&1), "matching industry must stay a candidate");
2221        assert!(
2222            !hits.contains(&2),
2223            "FieldEqual must be probed on the approximate All path"
2224        );
2225    }
2226
2227    #[test]
2228    fn all_of_scan_all_stays_scan_all() {
2229        let p = Predicate::All(vec![
2230            Predicate::VectorSimilar {
2231                field: "emb".into(),
2232                min: 0.5,
2233            },
2234            Predicate::VectorSimilar {
2235                field: "emb".into(),
2236                min: 0.9,
2237            },
2238        ]);
2239        match candidate_spec(&p) {
2240            CandidateSpec::Intersect(v) => assert_eq!(v.len(), 2),
2241            other => panic!("{other:?}"),
2242        }
2243        let spec = candidate_spec(&p);
2244        let mut idx = SideIndex::default();
2245        idx.insert(&spec, 1, &getter(&emb(&[1.0, 0.0])));
2246        idx.insert(&spec, 2, &getter(&emb(&[0.0, 1.0])));
2247        let hits = idx.candidates(&spec, &getter(&emb(&[1.0, 0.0])));
2248        assert_eq!(hits.into_iter().collect::<Vec<_>>(), vec![1, 2]);
2249    }
2250
2251    #[test]
2252    fn any_stays_union() {
2253        let p = Predicate::Any(vec![
2254            Predicate::FieldEqual {
2255                field: "industry".into(),
2256            },
2257            Predicate::Overlap {
2258                field: "tags".into(),
2259                min: 0.5,
2260            },
2261        ]);
2262        match candidate_spec(&p) {
2263            CandidateSpec::Union(v) => assert_eq!(v.len(), 2),
2264            other => panic!("{other:?}"),
2265        }
2266    }
2267
2268    /// Checkpoints are populated at insert, torn out at remove,
2269    /// and ckpts[0] must equal the full L2 norm.
2270    #[test]
2271    fn checkpoint_populated_and_consistent_with_norm() {
2272        let pred = Predicate::VectorSimilar {
2273            field: "emb".into(),
2274            min: 0.8,
2275        };
2276        let spec = candidate_spec(&pred);
2277        let xs = [3.0f64, 4.0]; // norm = 5.0
2278        let mut idx = SideIndex::default();
2279        idx.insert(&spec, 1, &getter(&emb(&xs)));
2280
2281        let ckpts = idx
2282            .vec_ckpts(1)
2283            .expect("checkpoints must exist after insert");
2284        let (_, norm) = idx.vec_meta(1).unwrap();
2285        assert!(
2286            (ckpts[0] - norm).abs() < 1e-12,
2287            "ckpts[0] must equal the full L2 norm; got {} vs {}",
2288            ckpts[0],
2289            norm
2290        );
2291        assert!(
2292            (norm - 5.0).abs() < 1e-12,
2293            "norm of [3,4] must be 5.0, got {norm}"
2294        );
2295
2296        // Remove must tear out checkpoints.
2297        idx.remove(&spec, 1, &getter(&emb(&xs)));
2298        assert!(
2299            idx.vec_ckpts(1).is_none(),
2300            "checkpoints must be removed after remove()"
2301        );
2302    }
2303
2304    /// fresh_ckpts_for returns None when the live vector's norm differs
2305    /// (freshness gate) and Some when it matches.
2306    #[test]
2307    fn fresh_ckpts_for_freshness_gate() {
2308        let pred = Predicate::VectorSimilar {
2309            field: "emb".into(),
2310            min: 0.8,
2311        };
2312        let spec = candidate_spec(&pred);
2313        let xs = [1.0f64, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0];
2314        let mut idx = SideIndex::default();
2315        idx.insert(&spec, 7, &getter(&emb(&xs)));
2316
2317        // Correct live vector → gate passes.
2318        let result = idx.fresh_ckpts_for(7, &xs);
2319        assert!(
2320            result.is_some(),
2321            "fresh_ckpts_for must succeed with matching live vector"
2322        );
2323        let (norm, ckpts) = result.unwrap();
2324        assert!((norm - 1.0).abs() < 1e-12);
2325        assert!((ckpts[0] - 1.0).abs() < 1e-12);
2326
2327        // Wrong norm → gate rejects.
2328        let wrong = [2.0f64, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]; // norm = 2.0
2329        assert!(
2330            idx.fresh_ckpts_for(7, &wrong).is_none(),
2331            "freshness gate must reject mismatched norm"
2332        );
2333
2334        // Wrong dim → gate rejects.
2335        let short = [1.0f64, 0.0];
2336        assert!(
2337            idx.fresh_ckpts_for(7, &short).is_none(),
2338            "freshness gate must reject mismatched dim"
2339        );
2340
2341        // Missing node → returns None.
2342        assert!(idx.fresh_ckpts_for(99, &xs).is_none());
2343    }
2344
2345    /// Checkpoints for a dim-16 vector: ckpts[i] must be non-increasing
2346    /// (suffix norms decrease as the suffix shrinks).
2347    #[test]
2348    fn checkpoint_suffix_norms_non_increasing() {
2349        let pred = Predicate::VectorSimilar {
2350            field: "emb".into(),
2351            min: 0.5,
2352        };
2353        let spec = candidate_spec(&pred);
2354        let xs: Vec<f64> = (1..=16).map(|i| i as f64).collect();
2355        let mut idx = SideIndex::default();
2356        idx.insert(&spec, 42, &getter(&emb(&xs)));
2357
2358        let ckpts = *idx.vec_ckpts(42).unwrap();
2359        for c in 0..7 {
2360            assert!(
2361                ckpts[c] >= ckpts[c + 1] - 1e-12,
2362                "suffix norm must be non-increasing: ckpts[{c}]={} < ckpts[{}]={}",
2363                ckpts[c],
2364                c + 1,
2365                ckpts[c + 1]
2366            );
2367        }
2368        // ckpts[7] = suffix norm of the last 2 elements (14..=16).
2369        let expected_last = (15.0f64 * 15.0 + 16.0 * 16.0).sqrt();
2370        assert!(
2371            (ckpts[7] - expected_last).abs() < 1e-9,
2372            "ckpts[7] should be norm of last segment; got {} vs {}",
2373            ckpts[7],
2374            expected_last
2375        );
2376    }
2377    // -----------------------------------------------------------------------
2378    // init_or_adopt_hnsw
2379    // -----------------------------------------------------------------------
2380
2381    /// A side seeded with three vectors, plus the `Hnsw` spec that indexes them.
2382    fn hnsw_side() -> (SideIndex, CandidateSpec<'static>) {
2383        let spec = CandidateSpec::Hnsw {
2384            field: "emb",
2385            k: 8,
2386            floor: None,
2387        };
2388        let mut side = SideIndex::default();
2389        side.init_hnsw("sim");
2390        for (id, xs) in [
2391            (1u32, vec![1.0, 0.0]),
2392            (2, vec![0.0, 1.0]),
2393            (3, vec![0.7, 0.7]),
2394        ] {
2395            side.insert(&spec, id, &getter(&emb(&xs)));
2396        }
2397        (side, spec)
2398    }
2399
2400    /// A usable blob is adopted before any scan, and its node ids come back so
2401    /// the scan can skip them.
2402    #[test]
2403    fn init_or_adopt_hnsw_adopts_a_usable_blob() {
2404        let (side, spec) = hnsw_side();
2405        let blob = side.export_hnsw_blob(true);
2406
2407        let mut fresh = SideIndex::default();
2408        let (ids, adopted) = fresh.init_or_adopt_hnsw("sim", &blob);
2409        assert!(adopted, "a usable blob must be adopted, not rebuilt");
2410        assert_eq!(ids, BTreeSet::from([1, 2, 3]));
2411        assert!(fresh.has_hnsw());
2412        assert_eq!(
2413            fresh.candidates(&spec, &getter(&emb(&[1.0, 0.0]))),
2414            side.candidates(&spec, &getter(&emb(&[1.0, 0.0]))),
2415            "the adopted graph must answer as the original did"
2416        );
2417    }
2418
2419    /// A blob whose version this build does not know is treated exactly as a
2420    /// corrupt one: an empty graph, an empty skip set, and the caller's node
2421    /// scan rebuilds. Until it does, the side answers from `hnsw_tracked`.
2422    #[test]
2423    fn an_unknown_version_leaves_the_graph_empty() {
2424        let (side, spec) = hnsw_side();
2425        let mut blob = side.export_hnsw_blob(true);
2426        blob[4] = 99; // the version's low byte
2427
2428        let mut fresh = SideIndex::default();
2429        let (ids, adopted) = fresh.init_or_adopt_hnsw("sim", &blob);
2430        assert!(!adopted, "an unreadable blob must not count as adopted");
2431        assert!(ids.is_empty(), "nothing may be skipped by the scan");
2432        assert!(!fresh.has_hnsw(), "the graph must be empty");
2433
2434        // The scan then fills it, and the full-scan fallback covers the gap.
2435        for (id, xs) in [
2436            (1u32, vec![1.0, 0.0]),
2437            (2, vec![0.0, 1.0]),
2438            (3, vec![0.7, 0.7]),
2439        ] {
2440            fresh.insert_skipping(&spec, id, &ids, &getter(&emb(&xs)));
2441        }
2442        assert!(fresh.has_hnsw());
2443        assert_eq!(
2444            fresh.candidates(&spec, &getter(&emb(&[1.0, 0.0]))),
2445            BTreeSet::from([1, 2, 3])
2446        );
2447    }
2448
2449    /// A truncated blob is treated exactly as an unknown version: an empty
2450    /// graph, an empty skip set, and a rebuild for the caller's node scan.
2451    #[test]
2452    fn an_unreadable_blob_leaves_the_graph_empty() {
2453        let (side, spec) = hnsw_side();
2454        let mut blob = side.export_hnsw_blob(true);
2455        blob.truncate(blob.len() / 2);
2456
2457        let mut fresh = SideIndex::default();
2458        let (ids, adopted) = fresh.init_or_adopt_hnsw("sim", &blob);
2459        assert!(!adopted, "an unreadable blob must not count as adopted");
2460        assert!(ids.is_empty(), "nothing may be skipped by the scan");
2461        assert!(!fresh.has_hnsw(), "the graph must be empty");
2462
2463        // The scan then fills it, and the full-scan fallback covers the gap.
2464        for (id, xs) in [
2465            (1u32, vec![1.0, 0.0]),
2466            (2, vec![0.0, 1.0]),
2467            (3, vec![0.7, 0.7]),
2468        ] {
2469            fresh.insert_skipping(&spec, id, &ids, &getter(&emb(&xs)));
2470        }
2471        assert!(fresh.has_hnsw());
2472        assert_eq!(
2473            fresh.candidates(&spec, &getter(&emb(&[1.0, 0.0]))),
2474            BTreeSet::from([1, 2, 3])
2475        );
2476    }
2477
2478    /// `insert_skipping` still tracks a skipped node for the fallback scan; it
2479    /// only declines to insert it into the graph a second time.
2480    #[test]
2481    fn insert_skipping_tracks_but_does_not_reinsert() {
2482        let (side, spec) = hnsw_side();
2483        let blob = side.export_hnsw_blob(true);
2484
2485        let mut fresh = SideIndex::default();
2486        let (already, _) = fresh.init_or_adopt_hnsw("sim", &blob);
2487        let before = fresh.hnsw_ref().map(|h| h.len());
2488
2489        // Node 3 is adopted; node 4 is not.
2490        fresh.insert_skipping(&spec, 3, &already, &getter(&emb(&[0.7, 0.7])));
2491        assert_eq!(
2492            fresh.hnsw_ref().map(|h| h.len()),
2493            before,
2494            "an adopted id must not be re-inserted"
2495        );
2496        fresh.insert_skipping(&spec, 4, &already, &getter(&emb(&[-1.0, 0.0])));
2497        assert_eq!(
2498            fresh.hnsw_ref().map(|h| h.len()),
2499            before.map(|n| n + 1),
2500            "a post-snapshot id must be inserted"
2501        );
2502        assert_eq!(
2503            fresh.candidates(&spec, &getter(&emb(&[1.0, 0.0]))),
2504            BTreeSet::from([1, 2, 3, 4])
2505        );
2506    }
2507}