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