kevy_index/
segment.rs

1//! [`Segment`] — one shard's slice of one index (index-follows-key,
2//! RFC D3). 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::value::IndexValue;
15
16/// Opaque pagination cursor: the last `(value, key)` served. Encoded
17/// by the runtime into the wire cursor; `None` = start.
18#[derive(Debug, Clone, PartialEq, Eq)]
19pub struct Cursor {
20    /// Last value served.
21    pub value: IndexValue,
22    /// Last key served (tiebreak within a value).
23    pub key: Vec<u8>,
24}
25
26/// Sizing + health counters (`IDX.LIST` / memory formula).
27#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
28pub struct SegmentStats {
29    /// Live entries.
30    pub entries: u64,
31    /// Approximate heap bytes (RFC D7 formula's measured side).
32    pub approx_bytes: u64,
33    /// Rows excluded because the field failed coercion / was missing.
34    pub coerce_failures: u64,
35    /// Values currently held by more than one key (unique fence).
36    pub duplicates: u64,
37}
38
39/// Per-entry structural overhead in the memory formula (RFC D7).
40const ENTRY_OVERHEAD: usize = 48;
41
42/// One shard's slice of one index.
43#[derive(Debug, Default)]
44pub struct Segment {
45    tree: BTreeSet<(IndexValue, Vec<u8>)>,
46    back: HashMap<Vec<u8>, IndexValue>,
47    value_counts: BTreeMap<IndexValue, u32>,
48    stats: SegmentStats,
49}
50
51impl Segment {
52    /// Empty segment.
53    pub fn new() -> Self {
54        Self::default()
55    }
56
57    /// Synchronous write-path maintenance: the row at `key` now
58    /// coerces to `new` (`None` = excluded / row deleted). Replaces
59    /// any previous entry for the key.
60    pub fn apply(&mut self, key: &[u8], new: Option<IndexValue>) {
61        if let Some(old) = self.back.remove(key) {
62            self.tree.remove(&(old.clone(), key.to_vec()));
63            self.stats.entries -= 1;
64            self.stats.approx_bytes = self
65                .stats
66                .approx_bytes
67                .saturating_sub((old.approx_bytes() + key.len() + ENTRY_OVERHEAD) as u64);
68            self.dec_count(&old);
69        }
70        match new {
71            Some(v) => {
72                self.stats.entries += 1;
73                self.stats.approx_bytes +=
74                    (v.approx_bytes() + key.len() + ENTRY_OVERHEAD) as u64;
75                self.inc_count(&v);
76                self.back.insert(key.to_vec(), v.clone());
77                self.tree.insert((v, key.to_vec()));
78            }
79            None => {
80                self.stats.coerce_failures += 1;
81            }
82        }
83    }
84
85    /// Row deleted (no coercion involved — not a coerce failure).
86    pub fn remove(&mut self, key: &[u8]) {
87        if let Some(old) = self.back.remove(key) {
88            self.tree.remove(&(old.clone(), key.to_vec()));
89            self.stats.entries -= 1;
90            self.stats.approx_bytes = self
91                .stats
92                .approx_bytes
93                .saturating_sub((old.approx_bytes() + key.len() + ENTRY_OVERHEAD) as u64);
94            self.dec_count(&old);
95        }
96    }
97
98    fn inc_count(&mut self, v: &IndexValue) {
99        let c = self.value_counts.entry(v.clone()).or_insert(0);
100        *c += 1;
101        if *c == 2 {
102            self.stats.duplicates += 1;
103        }
104    }
105
106    fn dec_count(&mut self, v: &IndexValue) {
107        if let Some(c) = self.value_counts.get_mut(v) {
108            if *c == 2 {
109                self.stats.duplicates -= 1;
110            }
111            *c -= 1;
112            if *c == 0 {
113                self.value_counts.remove(v);
114            }
115        }
116    }
117
118    /// Ordered scan of `[min, max]` (inclusive), resuming after
119    /// `cursor`, up to `limit` hits. Returns `(key, value)` pairs in
120    /// `(value, key)` order plus the cursor to resume from (`None` =
121    /// exhausted).
122    pub fn range(
123        &self,
124        min: &IndexValue,
125        max: &IndexValue,
126        cursor: Option<&Cursor>,
127        limit: usize,
128    ) -> (Vec<(Vec<u8>, IndexValue)>, Option<Cursor>) {
129        let lower: Bound<(IndexValue, Vec<u8>)> = match cursor {
130            Some(c) => Bound::Excluded((c.value.clone(), c.key.clone())),
131            None => Bound::Included((min.clone(), Vec::new())),
132        };
133        // Upper bound is exact via take-while — a synthetic sentinel
134        // key would MISS max-valued keys sorting above it.
135        let mut out = Vec::with_capacity(limit.min(64));
136        let mut iter = self.tree.range((lower, Bound::Unbounded));
137        for (v, k) in iter.by_ref() {
138            if v > max {
139                break;
140            }
141            out.push((k.clone(), v.clone()));
142            if out.len() == limit {
143                break;
144            }
145        }
146        let next = if out.len() == limit {
147            out.last().map(|(k, v)| Cursor { value: v.clone(), key: k.clone() })
148        } else {
149            None
150        };
151        (out, next)
152    }
153
154    /// Point lookup: every key holding exactly `value` (unique kind's
155    /// read; >1 hit = the declarative fence's `-DUPLICATE` signal).
156    pub fn eq(&self, value: &IndexValue, limit: usize) -> Vec<Vec<u8>> {
157        let lower = Bound::Included((value.clone(), Vec::new()));
158        self.tree
159            .range((lower, Bound::Unbounded))
160            .take_while(|(v, _)| v == value)
161            .take(limit)
162            .map(|(_, k)| k.clone())
163            .collect()
164    }
165
166    /// Count within `[min, max]` without materializing keys.
167    pub fn count(&self, min: &IndexValue, max: &IndexValue) -> u64 {
168        let lower = Bound::Included((min.clone(), Vec::new()));
169        self.tree
170            .range((lower, Bound::Unbounded))
171            .take_while(|(v, _)| v <= max)
172            .count() as u64
173    }
174
175    /// Verify hook: what value does the segment hold for `key`?
176    /// (`IDX.VERIFY` compares this against a fresh row coercion.)
177    pub fn verify_entry(&self, key: &[u8]) -> Option<&IndexValue> {
178        self.back.get(key)
179    }
180
181    /// Ordered streaming scan over the WHOLE segment: ascending (or
182    /// descending) `(value, key)` order, resuming exclusively past
183    /// `after`. The virtual-view pager drives this and probes
184    /// membership per candidate — O(limit × selectivity⁻¹) instead of
185    /// materializing the full member set.
186    pub fn scan<'s>(
187        &'s self,
188        after: Option<&Cursor>,
189        desc: bool,
190    ) -> Box<dyn Iterator<Item = (&'s IndexValue, &'s [u8])> + 's> {
191        match (after, desc) {
192            (None, false) => Box::new(self.tree.iter().map(|(v, k)| (v, k.as_slice()))),
193            (None, true) => Box::new(self.tree.iter().rev().map(|(v, k)| (v, k.as_slice()))),
194            (Some(c), false) => Box::new(
195                self.tree
196                    .range((
197                        Bound::Excluded((c.value.clone(), c.key.clone())),
198                        Bound::Unbounded,
199                    ))
200                    .map(|(v, k)| (v, k.as_slice())),
201            ),
202            (Some(c), true) => Box::new(
203                self.tree
204                    .range((
205                        Bound::Unbounded,
206                        Bound::Excluded((c.value.clone(), c.key.clone())),
207                    ))
208                    .rev()
209                    .map(|(v, k)| (v, k.as_slice())),
210            ),
211        }
212    }
213
214    /// Visit every `(key, value)` entry (verify / audit walks).
215    pub fn each_entry<F: FnMut(&[u8], &IndexValue)>(&self, mut f: F) {
216        for (k, v) in &self.back {
217            f(k.as_slice(), v);
218        }
219    }
220
221    /// Live counters.
222    pub fn stats(&self) -> SegmentStats {
223        self.stats
224    }
225}
226
227#[cfg(test)]
228mod tests {
229    use super::*;
230
231    fn i(v: i64) -> IndexValue {
232        IndexValue::I64(v)
233    }
234
235    fn seeded() -> Segment {
236        let mut s = Segment::new();
237        for (k, v) in [("u1", 30), ("u2", 25), ("u3", 30), ("u4", 40), ("u5", 18)] {
238            s.apply(k.as_bytes(), Some(i(v)));
239        }
240        s
241    }
242
243    #[test]
244    fn apply_replace_remove_and_stats() {
245        let mut s = seeded();
246        assert_eq!(s.stats().entries, 5);
247        assert_eq!(s.stats().duplicates, 1, "30 held twice");
248        // replace u1's value: 30 no longer duplicated
249        s.apply(b"u1", Some(i(31)));
250        assert_eq!(s.stats().entries, 5);
251        assert_eq!(s.stats().duplicates, 0);
252        // coerce-failure excludes and counts
253        s.apply(b"u2", None);
254        assert_eq!(s.stats().entries, 4);
255        assert_eq!(s.stats().coerce_failures, 1);
256        // remove is not a coerce failure
257        s.remove(b"u3");
258        assert_eq!(s.stats().entries, 3);
259        assert_eq!(s.stats().coerce_failures, 1);
260        assert!(s.verify_entry(b"u3").is_none());
261        assert_eq!(s.verify_entry(b"u4"), Some(&i(40)));
262    }
263
264    #[test]
265    fn range_scan_orders_and_paginates() {
266        let s = seeded();
267        let (page1, cur) = s.range(&i(18), &i(30), None, 2);
268        assert_eq!(page1[0], (b"u5".to_vec(), i(18)));
269        assert_eq!(page1[1], (b"u2".to_vec(), i(25)));
270        let cur = cur.expect("more pages");
271        let (page2, cur2) = s.range(&i(18), &i(30), Some(&cur), 10);
272        assert_eq!(
273            page2,
274            vec![(b"u1".to_vec(), i(30)), (b"u3".to_vec(), i(30))],
275            "value tie broken by key"
276        );
277        assert!(cur2.is_none(), "exhausted");
278        assert_eq!(s.count(&i(18), &i(30)), 4);
279        assert_eq!(s.count(&i(99), &i(100)), 0);
280    }
281
282    #[test]
283    fn eq_and_duplicate_fence() {
284        let s = seeded();
285        assert_eq!(s.eq(&i(30), 10), vec![b"u1".to_vec(), b"u3".to_vec()]);
286        assert_eq!(s.eq(&i(40), 10), vec![b"u4".to_vec()]);
287        assert!(s.eq(&i(99), 10).is_empty());
288    }
289
290    #[test]
291    fn long_keys_at_max_value_not_missed() {
292        let mut s = Segment::new();
293        let long_key = vec![0xFFu8; 80]; // sorts above any 64-byte sentinel
294        s.apply(&long_key, Some(i(30)));
295        s.apply(b"short", Some(i(30)));
296        let (hits, _) = s.range(&i(30), &i(30), None, 10);
297        assert_eq!(hits.len(), 2, "max-valued long key must not be missed");
298        assert_eq!(s.eq(&i(30), 10).len(), 2);
299        assert_eq!(s.count(&i(30), &i(30)), 2);
300    }
301
302    #[test]
303    fn f64_and_str_orders() {
304        let mut s = Segment::new();
305        s.apply(b"a", Some(IndexValue::F64(1.5)));
306        s.apply(b"b", Some(IndexValue::F64(-0.5)));
307        let (hits, _) = s.range(&IndexValue::F64(-1.0), &IndexValue::F64(2.0), None, 10);
308        assert_eq!(hits[0].0, b"b".to_vec());
309
310        let mut t = Segment::new();
311        t.apply(b"x", Some(IndexValue::Str(b"banana".to_vec())));
312        t.apply(b"y", Some(IndexValue::Str(b"apple".to_vec())));
313        let (hits, _) = t.range(
314            &IndexValue::Str(b"a".to_vec()),
315            &IndexValue::Str(b"z".to_vec()),
316            None,
317            10,
318        );
319        assert_eq!(hits[0].0, b"y".to_vec());
320    }
321}
kevy_index/segment.rs

kevy_index/
segment.rs
