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