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(crate) fn tree(&self) -> &BTreeSet<(IndexValue, Vec<u8>)> {
98 &self.tree
99 }
100
101 pub fn apply(&mut self, key: &[u8], new: Option<IndexValue>) {
105 if let Some(old) = self.back.remove(key) {
106 self.tree.remove(&(old.clone(), key.to_vec()));
107 self.stats.entries -= 1;
108 self.stats.approx_bytes = self
109 .stats
110 .approx_bytes
111 .saturating_sub((old.approx_bytes() + key.len() + ENTRY_OVERHEAD) as u64);
112 self.dec_count(&old);
113 }
114 match new {
115 Some(v) => {
116 self.stats.entries += 1;
117 self.stats.approx_bytes +=
118 (v.approx_bytes() + key.len() + ENTRY_OVERHEAD) as u64;
119 self.inc_count(&v);
120 self.back.insert(key.to_vec(), v.clone());
121 self.tree.insert((v, key.to_vec()));
122 }
123 None => {
124 self.stats.coerce_failures += 1;
125 }
126 }
127 }
128
129 pub fn remove(&mut self, key: &[u8]) {
131 if let Some(rv) = &mut self.values {
132 rv.clear(key);
133 }
134 if let Some(old) = self.back.remove(key) {
135 self.tree.remove(&(old.clone(), key.to_vec()));
136 self.stats.entries -= 1;
137 self.stats.approx_bytes = self
138 .stats
139 .approx_bytes
140 .saturating_sub((old.approx_bytes() + key.len() + ENTRY_OVERHEAD) as u64);
141 self.dec_count(&old);
142 }
143 }
144
145 fn inc_count(&mut self, v: &IndexValue) {
146 let c = self.value_counts.entry(v.clone()).or_insert(0);
147 *c += 1;
148 if *c == 2 {
149 self.stats.duplicates += 1;
150 }
151 }
152
153 fn dec_count(&mut self, v: &IndexValue) {
154 if let Some(c) = self.value_counts.get_mut(v) {
155 if *c == 2 {
156 self.stats.duplicates -= 1;
157 }
158 *c -= 1;
159 if *c == 0 {
160 self.value_counts.remove(v);
161 }
162 }
163 }
164
165 pub fn range(
170 &self,
171 min: &IndexValue,
172 max: &IndexValue,
173 cursor: Option<&Cursor>,
174 limit: usize,
175 ) -> (Vec<(Vec<u8>, IndexValue)>, Option<Cursor>) {
176 let lower: Bound<(IndexValue, Vec<u8>)> = match cursor {
177 Some(c) => Bound::Excluded((c.value.clone(), c.key.clone())),
178 None => Bound::Included((min.clone(), Vec::new())),
179 };
180 let mut out = Vec::with_capacity(limit.min(64));
183 let mut iter = self.tree.range((lower, Bound::Unbounded));
184 for (v, k) in iter.by_ref() {
185 if v > max {
186 break;
187 }
188 out.push((k.clone(), v.clone()));
189 if out.len() == limit {
190 break;
191 }
192 }
193 let next = if out.len() == limit {
194 out.last().map(|(k, v)| Cursor { value: v.clone(), key: k.clone() })
195 } else {
196 None
197 };
198 (out, next)
199 }
200
201 pub fn eq(&self, value: &IndexValue, limit: usize) -> Vec<Vec<u8>> {
204 let lower = Bound::Included((value.clone(), Vec::new()));
205 self.tree
206 .range((lower, Bound::Unbounded))
207 .take_while(|(v, _)| v == value)
208 .take(limit)
209 .map(|(_, k)| k.clone())
210 .collect()
211 }
212
213 pub fn count(&self, min: &IndexValue, max: &IndexValue) -> u64 {
215 let lower = Bound::Included((min.clone(), Vec::new()));
216 self.tree
217 .range((lower, Bound::Unbounded))
218 .take_while(|(v, _)| v <= max)
219 .count() as u64
220 }
221
222 pub fn verify_entry(&self, key: &[u8]) -> Option<&IndexValue> {
225 self.back.get(key)
226 }
227
228 pub fn scan<'s>(
234 &'s self,
235 after: Option<&Cursor>,
236 desc: bool,
237 ) -> Box<dyn Iterator<Item = (&'s IndexValue, &'s [u8])> + 's> {
238 match (after, desc) {
239 (None, false) => Box::new(self.tree.iter().map(|(v, k)| (v, k.as_slice()))),
240 (None, true) => Box::new(self.tree.iter().rev().map(|(v, k)| (v, k.as_slice()))),
241 (Some(c), false) => Box::new(
242 self.tree
243 .range((
244 Bound::Excluded((c.value.clone(), c.key.clone())),
245 Bound::Unbounded,
246 ))
247 .map(|(v, k)| (v, k.as_slice())),
248 ),
249 (Some(c), true) => Box::new(
250 self.tree
251 .range((
252 Bound::Unbounded,
253 Bound::Excluded((c.value.clone(), c.key.clone())),
254 ))
255 .rev()
256 .map(|(v, k)| (v, k.as_slice())),
257 ),
258 }
259 }
260
261 pub fn each_entry<F: FnMut(&[u8], &IndexValue)>(&self, mut f: F) {
263 for (k, v) in &self.back {
264 f(k.as_slice(), v);
265 }
266 }
267
268 pub fn stats(&self) -> SegmentStats {
271 let mut s = self.stats;
272 if let Some(rv) = &self.values {
273 s.approx_bytes += rv.approx_bytes();
274 }
275 s
276 }
277}
278
279#[cfg(test)]
280mod tests {
281 use super::*;
282
283 fn i(v: i64) -> IndexValue {
284 IndexValue::I64(v)
285 }
286
287 fn seeded() -> Segment {
288 let mut s = Segment::new();
289 for (k, v) in [("u1", 30), ("u2", 25), ("u3", 30), ("u4", 40), ("u5", 18)] {
290 s.apply(k.as_bytes(), Some(i(v)));
291 }
292 s
293 }
294
295 #[test]
296 fn apply_replace_remove_and_stats() {
297 let mut s = seeded();
298 assert_eq!(s.stats().entries, 5);
299 assert_eq!(s.stats().duplicates, 1, "30 held twice");
300 s.apply(b"u1", Some(i(31)));
302 assert_eq!(s.stats().entries, 5);
303 assert_eq!(s.stats().duplicates, 0);
304 s.apply(b"u2", None);
306 assert_eq!(s.stats().entries, 4);
307 assert_eq!(s.stats().coerce_failures, 1);
308 s.remove(b"u3");
310 assert_eq!(s.stats().entries, 3);
311 assert_eq!(s.stats().coerce_failures, 1);
312 assert!(s.verify_entry(b"u3").is_none());
313 assert_eq!(s.verify_entry(b"u4"), Some(&i(40)));
314 }
315
316 #[test]
317 fn range_scan_orders_and_paginates() {
318 let s = seeded();
319 let (page1, cur) = s.range(&i(18), &i(30), None, 2);
320 assert_eq!(page1[0], (b"u5".to_vec(), i(18)));
321 assert_eq!(page1[1], (b"u2".to_vec(), i(25)));
322 let cur = cur.expect("more pages");
323 let (page2, cur2) = s.range(&i(18), &i(30), Some(&cur), 10);
324 assert_eq!(
325 page2,
326 vec![(b"u1".to_vec(), i(30)), (b"u3".to_vec(), i(30))],
327 "value tie broken by key"
328 );
329 assert!(cur2.is_none(), "exhausted");
330 assert_eq!(s.count(&i(18), &i(30)), 4);
331 assert_eq!(s.count(&i(99), &i(100)), 0);
332 }
333
334 #[test]
335 fn eq_and_duplicate_fence() {
336 let s = seeded();
337 assert_eq!(s.eq(&i(30), 10), vec![b"u1".to_vec(), b"u3".to_vec()]);
338 assert_eq!(s.eq(&i(40), 10), vec![b"u4".to_vec()]);
339 assert!(s.eq(&i(99), 10).is_empty());
340 }
341
342 #[test]
343 fn long_keys_at_max_value_not_missed() {
344 let mut s = Segment::new();
345 let long_key = vec![0xFFu8; 80]; s.apply(&long_key, Some(i(30)));
347 s.apply(b"short", Some(i(30)));
348 let (hits, _) = s.range(&i(30), &i(30), None, 10);
349 assert_eq!(hits.len(), 2, "max-valued long key must not be missed");
350 assert_eq!(s.eq(&i(30), 10).len(), 2);
351 assert_eq!(s.count(&i(30), &i(30)), 2);
352 }
353
354 #[test]
355 fn f64_and_str_orders() {
356 let mut s = Segment::new();
357 s.apply(b"a", Some(IndexValue::F64(1.5)));
358 s.apply(b"b", Some(IndexValue::F64(-0.5)));
359 let (hits, _) = s.range(&IndexValue::F64(-1.0), &IndexValue::F64(2.0), None, 10);
360 assert_eq!(hits[0].0, b"b".to_vec());
361
362 let mut t = Segment::new();
363 t.apply(b"x", Some(IndexValue::Str(b"banana".to_vec())));
364 t.apply(b"y", Some(IndexValue::Str(b"apple".to_vec())));
365 let (hits, _) = t.range(
366 &IndexValue::Str(b"a".to_vec()),
367 &IndexValue::Str(b"z".to_vec()),
368 None,
369 10,
370 );
371 assert_eq!(hits[0].0, b"y".to_vec());
372 }
373}