1use 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#[derive(Debug, Clone, PartialEq, Eq)]
20pub struct Cursor {
21 pub value: IndexValue,
23 pub key: Vec<u8>,
25}
26
27#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
29pub struct SegmentStats {
30 pub entries: u64,
32 pub approx_bytes: u64,
35 pub coerce_failures: u64,
37 pub duplicates: u64,
39}
40
41const ENTRY_OVERHEAD: usize = 48;
43
44#[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 values: Option<RowValues>,
56}
57
58impl Segment {
59 pub fn new() -> Self {
61 Self::default()
62 }
63
64 pub fn with_values(n: usize) -> Self {
67 Segment { values: Some(RowValues::new(n)), ..Self::default() }
68 }
69
70 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 pub fn stored(&self, key: &[u8], field: usize) -> Option<&[u8]> {
93 self.values.as_ref()?.get(key, field)
94 }
95
96 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 pub(crate) fn tree(&self) -> &BTreeSet<(IndexValue, Vec<u8>)> {
109 &self.tree
110 }
111
112 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 += (v.approx_bytes() + key.len() + ENTRY_OVERHEAD) as u64;
129 self.inc_count(&v);
130 self.back.insert(key.to_vec(), v.clone());
131 self.tree.insert((v, key.to_vec()));
132 }
133 None => {
134 self.stats.coerce_failures += 1;
135 }
136 }
137 }
138
139 pub fn max_value(&self) -> Option<&IndexValue> {
142 self.tree.last().map(|(v, _)| v)
143 }
144
145 pub fn iter_below(&self, bound: &IndexValue) -> impl Iterator<Item = (&IndexValue, &[u8])> {
150 let end = (bound.clone(), Vec::new());
151 self.tree
152 .range((core::ops::Bound::Unbounded, core::ops::Bound::Excluded(end)))
153 .map(|(v, k)| (v, k.as_slice()))
154 }
155
156 pub fn split_off_below(&mut self, bound: &IndexValue) -> Vec<(IndexValue, Vec<u8>)> {
163 let kept = self.tree.split_off(&(bound.clone(), Vec::new()));
164 let evicted: Vec<(IndexValue, Vec<u8>)> =
165 core::mem::replace(&mut self.tree, kept).into_iter().collect();
166 for (v, k) in &evicted {
167 self.back.remove(k);
168 self.stats.entries -= 1;
169 self.stats.approx_bytes = self
170 .stats
171 .approx_bytes
172 .saturating_sub((v.approx_bytes() + k.len() + ENTRY_OVERHEAD) as u64);
173 self.dec_count(v);
174 if let Some(rv) = &mut self.values {
175 rv.clear(k);
176 }
177 }
178 evicted
179 }
180
181 pub fn remove(&mut self, key: &[u8]) {
183 if let Some(rv) = &mut self.values {
184 rv.clear(key);
185 }
186 if let Some(old) = self.back.remove(key) {
187 self.tree.remove(&(old.clone(), key.to_vec()));
188 self.stats.entries -= 1;
189 self.stats.approx_bytes = self
190 .stats
191 .approx_bytes
192 .saturating_sub((old.approx_bytes() + key.len() + ENTRY_OVERHEAD) as u64);
193 self.dec_count(&old);
194 }
195 }
196
197 fn inc_count(&mut self, v: &IndexValue) {
198 let c = self.value_counts.entry(v.clone()).or_insert(0);
199 *c += 1;
200 if *c == 2 {
201 self.stats.duplicates += 1;
202 }
203 }
204
205 fn dec_count(&mut self, v: &IndexValue) {
206 if let Some(c) = self.value_counts.get_mut(v) {
207 if *c == 2 {
208 self.stats.duplicates -= 1;
209 }
210 *c -= 1;
211 if *c == 0 {
212 self.value_counts.remove(v);
213 }
214 }
215 }
216
217 pub fn range(
222 &self,
223 min: &IndexValue,
224 max: &IndexValue,
225 cursor: Option<&Cursor>,
226 limit: usize,
227 ) -> (Vec<(Vec<u8>, IndexValue)>, Option<Cursor>) {
228 let lower: Bound<(IndexValue, Vec<u8>)> = match cursor {
229 Some(c) => Bound::Excluded((c.value.clone(), c.key.clone())),
230 None => Bound::Included((min.clone(), Vec::new())),
231 };
232 let mut out = Vec::with_capacity(limit.min(64));
235 let mut iter = self.tree.range((lower, Bound::Unbounded));
236 for (v, k) in iter.by_ref() {
237 if v > max {
238 break;
239 }
240 out.push((k.clone(), v.clone()));
241 if out.len() == limit {
242 break;
243 }
244 }
245 let next = if out.len() == limit {
246 out.last().map(|(k, v)| Cursor { value: v.clone(), key: k.clone() })
247 } else {
248 None
249 };
250 (out, next)
251 }
252
253 pub fn eq(&self, value: &IndexValue, limit: usize) -> Vec<Vec<u8>> {
256 let lower = Bound::Included((value.clone(), Vec::new()));
257 self.tree
258 .range((lower, Bound::Unbounded))
259 .take_while(|(v, _)| v == value)
260 .take(limit)
261 .map(|(_, k)| k.clone())
262 .collect()
263 }
264
265 pub fn count(&self, min: &IndexValue, max: &IndexValue) -> u64 {
267 let lower = Bound::Included((min.clone(), Vec::new()));
268 self.tree.range((lower, Bound::Unbounded)).take_while(|(v, _)| v <= max).count() as u64
269 }
270
271 pub fn verify_entry(&self, key: &[u8]) -> Option<&IndexValue> {
274 self.back.get(key)
275 }
276
277 pub fn scan<'s>(
283 &'s self,
284 after: Option<&Cursor>,
285 desc: bool,
286 ) -> Box<dyn Iterator<Item = (&'s IndexValue, &'s [u8])> + 's> {
287 match (after, desc) {
288 (None, false) => Box::new(self.tree.iter().map(|(v, k)| (v, k.as_slice()))),
289 (None, true) => Box::new(self.tree.iter().rev().map(|(v, k)| (v, k.as_slice()))),
290 (Some(c), false) => Box::new(
291 self.tree
292 .range((Bound::Excluded((c.value.clone(), c.key.clone())), Bound::Unbounded))
293 .map(|(v, k)| (v, k.as_slice())),
294 ),
295 (Some(c), true) => Box::new(
296 self.tree
297 .range((Bound::Unbounded, Bound::Excluded((c.value.clone(), c.key.clone()))))
298 .rev()
299 .map(|(v, k)| (v, k.as_slice())),
300 ),
301 }
302 }
303
304 pub fn each_entry<F: FnMut(&[u8], &IndexValue)>(&self, mut f: F) {
306 for (k, v) in &self.back {
307 f(k.as_slice(), v);
308 }
309 }
310
311 pub fn stats(&self) -> SegmentStats {
314 let mut s = self.stats;
315 if let Some(rv) = &self.values {
316 s.approx_bytes += rv.approx_bytes();
317 }
318 s
319 }
320}
321
322#[cfg(test)]
323mod tests {
324 use super::*;
325
326 fn i(v: i64) -> IndexValue {
327 IndexValue::I64(v)
328 }
329
330 fn seeded() -> Segment {
331 let mut s = Segment::new();
332 for (k, v) in [("u1", 30), ("u2", 25), ("u3", 30), ("u4", 40), ("u5", 18)] {
333 s.apply(k.as_bytes(), Some(i(v)));
334 }
335 s
336 }
337
338 #[test]
339 fn apply_replace_remove_and_stats() {
340 let mut s = seeded();
341 assert_eq!(s.stats().entries, 5);
342 assert_eq!(s.stats().duplicates, 1, "30 held twice");
343 s.apply(b"u1", Some(i(31)));
345 assert_eq!(s.stats().entries, 5);
346 assert_eq!(s.stats().duplicates, 0);
347 s.apply(b"u2", None);
349 assert_eq!(s.stats().entries, 4);
350 assert_eq!(s.stats().coerce_failures, 1);
351 s.remove(b"u3");
353 assert_eq!(s.stats().entries, 3);
354 assert_eq!(s.stats().coerce_failures, 1);
355 assert!(s.verify_entry(b"u3").is_none());
356 assert_eq!(s.verify_entry(b"u4"), Some(&i(40)));
357 }
358
359 #[test]
360 fn range_scan_orders_and_paginates() {
361 let s = seeded();
362 let (page1, cur) = s.range(&i(18), &i(30), None, 2);
363 assert_eq!(page1[0], (b"u5".to_vec(), i(18)));
364 assert_eq!(page1[1], (b"u2".to_vec(), i(25)));
365 let cur = cur.expect("more pages");
366 let (page2, cur2) = s.range(&i(18), &i(30), Some(&cur), 10);
367 assert_eq!(
368 page2,
369 vec![(b"u1".to_vec(), i(30)), (b"u3".to_vec(), i(30))],
370 "value tie broken by key"
371 );
372 assert!(cur2.is_none(), "exhausted");
373 assert_eq!(s.count(&i(18), &i(30)), 4);
374 assert_eq!(s.count(&i(99), &i(100)), 0);
375 }
376
377 #[test]
378 fn eq_and_duplicate_fence() {
379 let s = seeded();
380 assert_eq!(s.eq(&i(30), 10), vec![b"u1".to_vec(), b"u3".to_vec()]);
381 assert_eq!(s.eq(&i(40), 10), vec![b"u4".to_vec()]);
382 assert!(s.eq(&i(99), 10).is_empty());
383 }
384
385 #[test]
386 fn long_keys_at_max_value_not_missed() {
387 let mut s = Segment::new();
388 let long_key = vec![0xFFu8; 80]; s.apply(&long_key, Some(i(30)));
390 s.apply(b"short", Some(i(30)));
391 let (hits, _) = s.range(&i(30), &i(30), None, 10);
392 assert_eq!(hits.len(), 2, "max-valued long key must not be missed");
393 assert_eq!(s.eq(&i(30), 10).len(), 2);
394 assert_eq!(s.count(&i(30), &i(30)), 2);
395 }
396
397 #[test]
398 fn f64_and_str_orders() {
399 let mut s = Segment::new();
400 s.apply(b"a", Some(IndexValue::F64(1.5)));
401 s.apply(b"b", Some(IndexValue::F64(-0.5)));
402 let (hits, _) = s.range(&IndexValue::F64(-1.0), &IndexValue::F64(2.0), None, 10);
403 assert_eq!(hits[0].0, b"b".to_vec());
404
405 let mut t = Segment::new();
406 t.apply(b"x", Some(IndexValue::Str(b"banana".to_vec())));
407 t.apply(b"y", Some(IndexValue::Str(b"apple".to_vec())));
408 let (hits, _) =
409 t.range(&IndexValue::Str(b"a".to_vec()), &IndexValue::Str(b"z".to_vec()), None, 10);
410 assert_eq!(hits[0].0, b"y".to_vec());
411 }
412}