kevy_index/
segment.rs

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

kevy_index/
segment.rs
