Skip to main content

kevy_index/
segment.rs

1//! [`Segment`] — one shard's slice of one index (index-follows-key).
2//! Range = `BTreeSet<(value, key)>`; Unique = the same
3//! tree (point lookups are a 1-value range) plus a duplicate counter
4//! for the declarative fence.
5//!
6//! The runtime keeps a reverse map `key → value` inside the segment so
7//! `apply` can remove a row's OLD entry without re-reading history.
8
9use std::collections::BTreeMap;
10use std::collections::BTreeSet;
11use std::collections::HashMap;
12use std::ops::Bound;
13
14use crate::rowvalues::RowValues;
15use crate::value::IndexValue;
16
17/// Opaque pagination cursor: the last `(value, key)` served. Encoded
18/// by the runtime into the wire cursor; `None` = start.
19#[derive(Debug, Clone, PartialEq, Eq)]
20pub struct Cursor {
21    /// Last value served.
22    pub value: IndexValue,
23    /// Last key served (tiebreak within a value).
24    pub key: Vec<u8>,
25}
26
27/// Sizing + health counters (`IDX.LIST` / memory formula).
28#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
29pub struct SegmentStats {
30    /// Live entries.
31    pub entries: u64,
32    /// Approximate heap bytes (the measured side of the documented
33    /// memory formula).
34    pub approx_bytes: u64,
35    /// Rows excluded because the field failed coercion / was missing.
36    pub coerce_failures: u64,
37    /// Values currently held by more than one key (unique fence).
38    pub duplicates: u64,
39}
40
41/// Per-entry structural overhead in the memory formula.
42const ENTRY_OVERHEAD: usize = 48;
43
44/// One shard's slice of one index.
45#[derive(Debug, Default)]
46pub struct Segment {
47    tree: BTreeSet<(IndexValue, Vec<u8>)>,
48    back: HashMap<Vec<u8>, IndexValue>,
49    value_counts: BTreeMap<IndexValue, u32>,
50    stats: SegmentStats,
51    /// The stored-value side-channel — `Some` only when the index
52    /// declared `VALUES`. An index without the declaration holds `None`
53    /// and pays nothing: every values touch below is a never-taken
54    /// `if let` (the `Option<Positions>` physical-bypass pattern, A5).
55    values: Option<RowValues>,
56}
57
58impl Segment {
59    /// Empty segment.
60    pub fn new() -> Self {
61        Self::default()
62    }
63
64    /// Empty segment carrying the stored-value side-channel for `n`
65    /// declared `VALUES` fields (`n` = the spec's `values.len()`).
66    pub fn with_values(n: usize) -> Self {
67        Segment { values: Some(RowValues::new(n)), ..Self::default() }
68    }
69
70    /// [`Segment::apply`] plus the row's declared stored values. The
71    /// values follow the entry: an indexed row stores them, an excluded
72    /// or deleted one drops them.
73    pub fn apply_with_values(
74        &mut self,
75        key: &[u8],
76        new: Option<IndexValue>,
77        vals: &[Option<&[u8]>],
78    ) {
79        let indexed = new.is_some();
80        self.apply(key, new);
81        if let Some(rv) = &mut self.values {
82            if indexed {
83                rv.set(key, vals);
84            } else {
85                rv.clear(key);
86            }
87        }
88    }
89
90    /// `key`'s stored value for declared `VALUES` field `field`, or
91    /// `None` when the row has none (or the index declared none).
92    pub fn stored(&self, key: &[u8], field: usize) -> Option<&[u8]> {
93        self.values.as_ref()?.get(key, field)
94    }
95
96    /// Every stored value of one row, aligned with the declared
97    /// `VALUES` order — what an eviction carries into a cold entry's
98    /// payload so the clause-carrying cold path never re-reads the
99    /// row. Empty when the index declared no values.
100    pub fn stored_row(&self, key: &[u8]) -> Vec<Option<&[u8]>> {
101        match self.values.as_ref() {
102            Some(rv) => (0..rv.arity()).map(|f| rv.get(key, f)).collect(),
103            None => Vec::new(),
104        }
105    }
106
107    /// The `(value, key)` tree, for the clause-carrying scan.
108    pub(crate) fn tree(&self) -> &BTreeSet<(IndexValue, Vec<u8>)> {
109        &self.tree
110    }
111
112    /// Synchronous write-path maintenance: the row at `key` now
113    /// coerces to `new` (`None` = excluded / row deleted). Replaces
114    /// any previous entry for the key.
115    pub fn apply(&mut self, key: &[u8], new: Option<IndexValue>) {
116        if let Some(old) = self.back.remove(key) {
117            self.tree.remove(&(old.clone(), key.to_vec()));
118            self.stats.entries -= 1;
119            self.stats.approx_bytes = self
120                .stats
121                .approx_bytes
122                .saturating_sub((old.approx_bytes() + key.len() + ENTRY_OVERHEAD) as u64);
123            self.dec_count(&old);
124        }
125        match new {
126            Some(v) => {
127                self.stats.entries += 1;
128                self.stats.approx_bytes +=
129                    (v.approx_bytes() + key.len() + ENTRY_OVERHEAD) as u64;
130                self.inc_count(&v);
131                self.back.insert(key.to_vec(), v.clone());
132                self.tree.insert((v, key.to_vec()));
133            }
134            None => {
135                self.stats.coerce_failures += 1;
136            }
137        }
138    }
139
140    /// The largest value present, if any — the window boundary's
141    /// tree-tail read.
142    pub fn max_value(&self) -> Option<&IndexValue> {
143        self.tree.last().map(|(v, _)| v)
144    }
145
146    /// Entries strictly below `bound`, tree order — the read-only
147    /// preview of [`Self::split_off_below`]'s batch (the slide builds
148    /// its segment from this BEFORE cutting, so an I/O failure leaves
149    /// the tree untouched).
150    pub fn iter_below(&self, bound: &IndexValue) -> impl Iterator<Item = (&IndexValue, &[u8])> {
151        let end = (bound.clone(), Vec::new());
152        self.tree
153            .range((core::ops::Bound::Unbounded, core::ops::Bound::Excluded(end)))
154            .map(|(v, k)| (v, k.as_slice()))
155    }
156
157    /// Detach every entry whose value sorts below `bound`, in tree
158    /// order — the window-eviction cut. The detached batch leaves all
159    /// of the segment's books (tree, reverse map, value counts, stored
160    /// values, stats) exactly as if each entry had been removed one by
161    /// one; the empty-key sentinel keeps every entry AT `bound` in the
162    /// hot tree, so the cut is strictly `< bound`.
163    pub fn split_off_below(&mut self, bound: &IndexValue) -> Vec<(IndexValue, Vec<u8>)> {
164        let kept = self.tree.split_off(&(bound.clone(), Vec::new()));
165        let evicted: Vec<(IndexValue, Vec<u8>)> =
166            core::mem::replace(&mut self.tree, kept).into_iter().collect();
167        for (v, k) in &evicted {
168            self.back.remove(k);
169            self.stats.entries -= 1;
170            self.stats.approx_bytes = self
171                .stats
172                .approx_bytes
173                .saturating_sub((v.approx_bytes() + k.len() + ENTRY_OVERHEAD) as u64);
174            self.dec_count(v);
175            if let Some(rv) = &mut self.values {
176                rv.clear(k);
177            }
178        }
179        evicted
180    }
181
182    /// Row deleted (no coercion involved — not a coerce failure).
183    pub fn remove(&mut self, key: &[u8]) {
184        if let Some(rv) = &mut self.values {
185            rv.clear(key);
186        }
187        if let Some(old) = self.back.remove(key) {
188            self.tree.remove(&(old.clone(), key.to_vec()));
189            self.stats.entries -= 1;
190            self.stats.approx_bytes = self
191                .stats
192                .approx_bytes
193                .saturating_sub((old.approx_bytes() + key.len() + ENTRY_OVERHEAD) as u64);
194            self.dec_count(&old);
195        }
196    }
197
198    fn inc_count(&mut self, v: &IndexValue) {
199        let c = self.value_counts.entry(v.clone()).or_insert(0);
200        *c += 1;
201        if *c == 2 {
202            self.stats.duplicates += 1;
203        }
204    }
205
206    fn dec_count(&mut self, v: &IndexValue) {
207        if let Some(c) = self.value_counts.get_mut(v) {
208            if *c == 2 {
209                self.stats.duplicates -= 1;
210            }
211            *c -= 1;
212            if *c == 0 {
213                self.value_counts.remove(v);
214            }
215        }
216    }
217
218    /// Ordered scan of `[min, max]` (inclusive), resuming after
219    /// `cursor`, up to `limit` hits. Returns `(key, value)` pairs in
220    /// `(value, key)` order plus the cursor to resume from (`None` =
221    /// exhausted).
222    pub fn range(
223        &self,
224        min: &IndexValue,
225        max: &IndexValue,
226        cursor: Option<&Cursor>,
227        limit: usize,
228    ) -> (Vec<(Vec<u8>, IndexValue)>, Option<Cursor>) {
229        let lower: Bound<(IndexValue, Vec<u8>)> = match cursor {
230            Some(c) => Bound::Excluded((c.value.clone(), c.key.clone())),
231            None => Bound::Included((min.clone(), Vec::new())),
232        };
233        // Upper bound is exact via take-while — a synthetic sentinel
234        // key would MISS max-valued keys sorting above it.
235        let mut out = Vec::with_capacity(limit.min(64));
236        let mut iter = self.tree.range((lower, Bound::Unbounded));
237        for (v, k) in iter.by_ref() {
238            if v > max {
239                break;
240            }
241            out.push((k.clone(), v.clone()));
242            if out.len() == limit {
243                break;
244            }
245        }
246        let next = if out.len() == limit {
247            out.last().map(|(k, v)| Cursor { value: v.clone(), key: k.clone() })
248        } else {
249            None
250        };
251        (out, next)
252    }
253
254    /// Point lookup: every key holding exactly `value` (unique kind's
255    /// read; >1 hit = the declarative fence's `-DUPLICATE` signal).
256    pub fn eq(&self, value: &IndexValue, limit: usize) -> Vec<Vec<u8>> {
257        let lower = Bound::Included((value.clone(), Vec::new()));
258        self.tree
259            .range((lower, Bound::Unbounded))
260            .take_while(|(v, _)| v == value)
261            .take(limit)
262            .map(|(_, k)| k.clone())
263            .collect()
264    }
265
266    /// Count within `[min, max]` without materializing keys.
267    pub fn count(&self, min: &IndexValue, max: &IndexValue) -> u64 {
268        let lower = Bound::Included((min.clone(), Vec::new()));
269        self.tree
270            .range((lower, Bound::Unbounded))
271            .take_while(|(v, _)| v <= max)
272            .count() as u64
273    }
274
275    /// Verify hook: what value does the segment hold for `key`?
276    /// (`IDX.VERIFY` compares this against a fresh row coercion.)
277    pub fn verify_entry(&self, key: &[u8]) -> Option<&IndexValue> {
278        self.back.get(key)
279    }
280
281    /// Ordered streaming scan over the WHOLE segment: ascending (or
282    /// descending) `(value, key)` order, resuming exclusively past
283    /// `after`. The virtual-view pager drives this and probes
284    /// membership per candidate — O(limit × selectivity⁻¹) instead of
285    /// materializing the full member set.
286    pub fn scan<'s>(
287        &'s self,
288        after: Option<&Cursor>,
289        desc: bool,
290    ) -> Box<dyn Iterator<Item = (&'s IndexValue, &'s [u8])> + 's> {
291        match (after, desc) {
292            (None, false) => Box::new(self.tree.iter().map(|(v, k)| (v, k.as_slice()))),
293            (None, true) => Box::new(self.tree.iter().rev().map(|(v, k)| (v, k.as_slice()))),
294            (Some(c), false) => Box::new(
295                self.tree
296                    .range((
297                        Bound::Excluded((c.value.clone(), c.key.clone())),
298                        Bound::Unbounded,
299                    ))
300                    .map(|(v, k)| (v, k.as_slice())),
301            ),
302            (Some(c), true) => Box::new(
303                self.tree
304                    .range((
305                        Bound::Unbounded,
306                        Bound::Excluded((c.value.clone(), c.key.clone())),
307                    ))
308                    .rev()
309                    .map(|(v, k)| (v, k.as_slice())),
310            ),
311        }
312    }
313
314    /// Visit every `(key, value)` entry (verify / audit walks).
315    pub fn each_entry<F: FnMut(&[u8], &IndexValue)>(&self, mut f: F) {
316        for (k, v) in &self.back {
317            f(k.as_slice(), v);
318        }
319    }
320
321    /// Live counters. The stored-value column's heap joins the memory
322    /// term when (and only when) the index declared `VALUES`.
323    pub fn stats(&self) -> SegmentStats {
324        let mut s = self.stats;
325        if let Some(rv) = &self.values {
326            s.approx_bytes += rv.approx_bytes();
327        }
328        s
329    }
330}
331
332#[cfg(test)]
333mod tests {
334    use super::*;
335
336    fn i(v: i64) -> IndexValue {
337        IndexValue::I64(v)
338    }
339
340    fn seeded() -> Segment {
341        let mut s = Segment::new();
342        for (k, v) in [("u1", 30), ("u2", 25), ("u3", 30), ("u4", 40), ("u5", 18)] {
343            s.apply(k.as_bytes(), Some(i(v)));
344        }
345        s
346    }
347
348    #[test]
349    fn apply_replace_remove_and_stats() {
350        let mut s = seeded();
351        assert_eq!(s.stats().entries, 5);
352        assert_eq!(s.stats().duplicates, 1, "30 held twice");
353        // replace u1's value: 30 no longer duplicated
354        s.apply(b"u1", Some(i(31)));
355        assert_eq!(s.stats().entries, 5);
356        assert_eq!(s.stats().duplicates, 0);
357        // coerce-failure excludes and counts
358        s.apply(b"u2", None);
359        assert_eq!(s.stats().entries, 4);
360        assert_eq!(s.stats().coerce_failures, 1);
361        // remove is not a coerce failure
362        s.remove(b"u3");
363        assert_eq!(s.stats().entries, 3);
364        assert_eq!(s.stats().coerce_failures, 1);
365        assert!(s.verify_entry(b"u3").is_none());
366        assert_eq!(s.verify_entry(b"u4"), Some(&i(40)));
367    }
368
369    #[test]
370    fn range_scan_orders_and_paginates() {
371        let s = seeded();
372        let (page1, cur) = s.range(&i(18), &i(30), None, 2);
373        assert_eq!(page1[0], (b"u5".to_vec(), i(18)));
374        assert_eq!(page1[1], (b"u2".to_vec(), i(25)));
375        let cur = cur.expect("more pages");
376        let (page2, cur2) = s.range(&i(18), &i(30), Some(&cur), 10);
377        assert_eq!(
378            page2,
379            vec![(b"u1".to_vec(), i(30)), (b"u3".to_vec(), i(30))],
380            "value tie broken by key"
381        );
382        assert!(cur2.is_none(), "exhausted");
383        assert_eq!(s.count(&i(18), &i(30)), 4);
384        assert_eq!(s.count(&i(99), &i(100)), 0);
385    }
386
387    #[test]
388    fn eq_and_duplicate_fence() {
389        let s = seeded();
390        assert_eq!(s.eq(&i(30), 10), vec![b"u1".to_vec(), b"u3".to_vec()]);
391        assert_eq!(s.eq(&i(40), 10), vec![b"u4".to_vec()]);
392        assert!(s.eq(&i(99), 10).is_empty());
393    }
394
395    #[test]
396    fn long_keys_at_max_value_not_missed() {
397        let mut s = Segment::new();
398        let long_key = vec![0xFFu8; 80]; // sorts above any 64-byte sentinel
399        s.apply(&long_key, Some(i(30)));
400        s.apply(b"short", Some(i(30)));
401        let (hits, _) = s.range(&i(30), &i(30), None, 10);
402        assert_eq!(hits.len(), 2, "max-valued long key must not be missed");
403        assert_eq!(s.eq(&i(30), 10).len(), 2);
404        assert_eq!(s.count(&i(30), &i(30)), 2);
405    }
406
407    #[test]
408    fn f64_and_str_orders() {
409        let mut s = Segment::new();
410        s.apply(b"a", Some(IndexValue::F64(1.5)));
411        s.apply(b"b", Some(IndexValue::F64(-0.5)));
412        let (hits, _) = s.range(&IndexValue::F64(-1.0), &IndexValue::F64(2.0), None, 10);
413        assert_eq!(hits[0].0, b"b".to_vec());
414
415        let mut t = Segment::new();
416        t.apply(b"x", Some(IndexValue::Str(b"banana".to_vec())));
417        t.apply(b"y", Some(IndexValue::Str(b"apple".to_vec())));
418        let (hits, _) = t.range(
419            &IndexValue::Str(b"a".to_vec()),
420            &IndexValue::Str(b"z".to_vec()),
421            None,
422            10,
423        );
424        assert_eq!(hits[0].0, b"y".to_vec());
425    }
426}