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 +=
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 pub fn max_value(&self) -> Option<&IndexValue> {
143 self.tree.last().map(|(v, _)| v)
144 }
145
146 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 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 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 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 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 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 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 pub fn verify_entry(&self, key: &[u8]) -> Option<&IndexValue> {
278 self.back.get(key)
279 }
280
281 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 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 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 s.apply(b"u1", Some(i(31)));
355 assert_eq!(s.stats().entries, 5);
356 assert_eq!(s.stats().duplicates, 0);
357 s.apply(b"u2", None);
359 assert_eq!(s.stats().entries, 4);
360 assert_eq!(s.stats().coerce_failures, 1);
361 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]; 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}