1mod cursor;
42mod partition;
43
44pub use self::cursor::Cursor;
45use self::partition::Partition;
46#[commonware_macros::stability(ALPHA)]
47use crate::index::partitioned::{PartitionRange, Partitioned};
48use crate::{
49 index::{
50 Cursor as CursorTrait, Factory, Ordered, Unordered,
51 partitioned::partition_index_and_sub_key,
52 },
53 translator::Translator,
54};
55use commonware_runtime::{
56 Metrics,
57 telemetry::metrics::{Counter, Gauge, MetricsExt as _},
58};
59use std::{
60 collections::{BTreeMap, HashMap, btree_map, hash_map},
61 ops::Bound,
62};
63
64const SPILL_THRESHOLD: usize = 512;
70
71pub struct Index<T: Translator, V: Send + Sync, const P: usize> {
75 translator: T,
77
78 partitions: Box<[Partition<T::Key, V>]>,
83
84 spilled: HashMap<usize, BTreeMap<T::Key, Vec<V>>>,
88
89 threshold: usize,
92
93 keys: Gauge,
95
96 items: Gauge,
98
99 pruned: Counter,
101
102 spills: Counter,
106}
107
108impl<T: Translator, V: Send + Sync, const P: usize> Index<T, V, P> {
109 pub fn new(ctx: impl Metrics, translator: T) -> Self {
111 const {
112 assert!(P > 0 && P <= 3, "P must be in 1..=3");
113 }
114 let count = 1usize << (P * 8);
115 let partitions = (0..count)
116 .map(|_| Partition::default())
117 .collect::<Vec<_>>()
118 .into_boxed_slice();
119 Self {
120 translator,
121 partitions,
122 spilled: HashMap::new(),
123 threshold: SPILL_THRESHOLD,
124 keys: ctx.gauge("keys", "Number of translated keys in the index"),
125 items: ctx.gauge("items", "Number of items in the index"),
126 pruned: ctx.counter("pruned", "Number of items pruned"),
127 spills: ctx.counter("spills", "Number of partitions spilled to the side-table"),
128 }
129 }
130
131 #[cfg(test)]
135 pub(crate) fn with_threshold(ctx: impl Metrics, translator: T, threshold: usize) -> Self {
136 assert!(threshold > 0, "spill threshold must be at least 1");
137 let mut index = Self::new(ctx, translator);
138 index.threshold = threshold;
139 index
140 }
141
142 #[commonware_macros::stability(ALPHA)]
144 fn for_each_value(&self, mut f: impl FnMut(&V)) {
145 for (p, partition) in self.partitions.iter().enumerate() {
146 for v in partition.values_iter() {
147 f(v);
148 }
149 if let Some(inner) = self.spilled_partition(p) {
150 for vals in inner.values() {
151 for v in vals {
152 f(v);
153 }
154 }
155 }
156 }
157 }
158
159 fn maybe_spill(&mut self, i: usize) {
161 if self.partitions[i].len() < self.threshold {
162 return;
163 }
164 let inner: BTreeMap<T::Key, Vec<V>> = self.partitions[i].drain_runs().into_iter().collect();
165 self.spilled.insert(i, inner);
166 self.spills.inc();
167 }
168
169 fn spilled_partition(&self, i: usize) -> Option<&BTreeMap<T::Key, Vec<V>>> {
172 if self.spilled.is_empty() {
173 return None;
174 }
175 self.spilled.get(&i)
176 }
177
178 fn partition_values(&self, i: usize, k: &T::Key) -> &[V] {
181 if !self.partitions[i].is_empty() {
182 return self.partitions[i].values(k);
183 }
184 self.spilled_partition(i)
185 .and_then(|inner| inner.get(k))
186 .map_or(&[], Vec::as_slice)
187 }
188
189 fn partition_first(&self, i: usize) -> Option<&[V]> {
191 self.partitions[i].first_values().or_else(|| {
192 self.spilled_partition(i)?
193 .first_key_value()
194 .map(|(_, v)| v.as_slice())
195 })
196 }
197
198 fn partition_last(&self, i: usize) -> Option<&[V]> {
200 self.partitions[i].last_values().or_else(|| {
201 self.spilled_partition(i)?
202 .last_key_value()
203 .map(|(_, v)| v.as_slice())
204 })
205 }
206
207 fn is_empty(&self) -> bool {
209 self.keys.get() == 0
210 }
211
212 fn partition_next_after(&self, i: usize, k: &T::Key) -> Option<&[V]> {
214 self.partitions[i].next_values_after(k).or_else(|| {
215 self.spilled_partition(i)?
216 .range((Bound::Excluded(*k), Bound::Unbounded))
217 .next()
218 .map(|(_, v)| v.as_slice())
219 })
220 }
221
222 fn partition_prev_before(&self, i: usize, k: &T::Key) -> Option<&[V]> {
224 self.partitions[i].prev_values_before(k).or_else(|| {
225 self.spilled_partition(i)?
226 .range((Bound::Unbounded, Bound::Excluded(*k)))
227 .next_back()
228 .map(|(_, v)| v.as_slice())
229 })
230 }
231
232 #[cfg(test)]
234 pub(crate) fn spilled_count(&self) -> usize {
235 self.spilled.len()
236 }
237
238 #[cfg(test)]
240 fn spills(&self) -> usize {
241 self.spills.get() as usize
242 }
243
244 fn get_mut_slot(&mut self, i: usize, sub: &[u8]) -> Option<Cursor<'_, T::Key, V>> {
247 let k = self.translator.transform(sub);
248 self.maybe_spill(i);
249 if !self.partitions[i].is_empty() {
250 let run = self.partitions[i].run_range(&k);
251 if run.is_empty() {
252 return None;
253 }
254 return Some(Cursor::soa(
255 &mut self.partitions[i],
256 k,
257 run,
258 &self.keys,
259 &self.items,
260 &self.pruned,
261 ));
262 }
263
264 if self
266 .spilled_partition(i)
267 .is_some_and(|inner| inner.contains_key(&k))
268 {
269 return Some(Cursor::spilled(
270 &mut self.spilled,
271 i,
272 k,
273 &self.keys,
274 &self.items,
275 &self.pruned,
276 ));
277 }
278
279 None
281 }
282
283 fn get_mut_or_insert_slot(
287 &mut self,
288 i: usize,
289 sub: &[u8],
290 value: V,
291 ) -> Option<Cursor<'_, T::Key, V>> {
292 let k = self.translator.transform(sub);
293 self.maybe_spill(i);
294 if !self.partitions[i].is_empty() {
295 let run = self.partitions[i].run_range(&k);
296 if !run.is_empty() {
297 return Some(Cursor::soa(
298 &mut self.partitions[i],
299 k,
300 run,
301 &self.keys,
302 &self.items,
303 &self.pruned,
304 ));
305 }
306 self.partitions[i].insert_at(run.end, k, value);
307 self.keys.inc();
308 self.items.inc();
309 self.maybe_spill(i);
310 return None;
311 }
312
313 if let Some(inner) = self.spilled_partition(i) {
316 if inner.contains_key(&k) {
317 return Some(Cursor::spilled(
318 &mut self.spilled,
319 i,
320 k,
321 &self.keys,
322 &self.items,
323 &self.pruned,
324 ));
325 }
326 self.spilled.get_mut(&i).unwrap().insert(k, vec![value]);
327 self.keys.inc();
328 self.items.inc();
329 return None;
330 }
331
332 self.partitions[i].insert_at(0, k, value);
334 self.keys.inc();
335 self.items.inc();
336 self.maybe_spill(i);
337
338 None
339 }
340}
341
342#[commonware_macros::stability(ALPHA)]
343impl<T: Translator, V: Send + Sync + 'static, const P: usize> Partitioned for Index<T, V, P> {
344 type Range = RangeIndex<T, V, P>;
345
346 fn partition_count(&self) -> usize {
347 self.partitions.len()
348 }
349
350 fn partition_of(key: &[u8]) -> usize {
351 partition_index_and_sub_key::<P>(key).0
352 }
353
354 fn new_range(&self, offset: usize, count: usize) -> RangeIndex<T, V, P> {
358 let partitions = (0..count)
359 .map(|_| Partition::default())
360 .collect::<Vec<_>>()
361 .into_boxed_slice();
362 RangeIndex {
363 index: Self {
364 translator: self.translator.clone(),
365 partitions,
366 spilled: HashMap::new(),
367 threshold: self.threshold,
368 keys: self.keys.clone(),
369 items: self.items.clone(),
370 pruned: self.pruned.clone(),
371 spills: self.spills.clone(),
372 },
373 offset,
374 }
375 }
376
377 fn install_range(&mut self, mut worker: RangeIndex<T, V, P>) {
380 let lo = worker.offset;
381 let len = worker.index.partitions.len();
382
383 assert!(
385 self.spilled.keys().all(|&p| p < lo || p >= lo + len),
386 "install target range must be empty"
387 );
388 for (local, partition) in worker.index.partitions.iter_mut().enumerate() {
389 let global = lo + local;
390 assert!(
391 self.partitions[global].is_empty(),
392 "install target range must be empty"
393 );
394 self.partitions[global] = std::mem::take(partition);
395 }
396
397 for (local, inner) in worker.index.spilled.drain() {
400 self.spilled.insert(lo + local, inner);
401 }
402 }
403}
404
405#[commonware_macros::stability(ALPHA)]
411pub(crate) struct RangeIndex<T: Translator, V: Send + Sync, const P: usize> {
412 index: Index<T, V, P>,
414
415 offset: usize,
417}
418
419#[commonware_macros::stability(ALPHA)]
420impl<T: Translator, V: Send + Sync, const P: usize> PartitionRange for RangeIndex<T, V, P> {
421 type Value = V;
422 type Cursor<'a>
423 = Cursor<'a, T::Key, V>
424 where
425 Self: 'a;
426
427 fn get_mut(&mut self, key: &[u8]) -> Option<Cursor<'_, T::Key, V>> {
428 let (i, sub) = partition_index_and_sub_key::<P>(key);
429 self.index.get_mut_slot(i - self.offset, sub)
430 }
431
432 fn get_mut_or_insert(&mut self, key: &[u8], value: V) -> Option<Cursor<'_, T::Key, V>> {
433 let (i, sub) = partition_index_and_sub_key::<P>(key);
434 self.index
435 .get_mut_or_insert_slot(i - self.offset, sub, value)
436 }
437
438 fn for_each_value(&self, f: impl FnMut(&V)) {
439 self.index.for_each_value(f);
440 }
441}
442
443impl<T: Translator, V: Send + Sync, const P: usize> Factory for Index<T, V, P> {
444 type Translator = T;
445
446 fn new(ctx: impl Metrics, translator: T) -> Self {
447 Self::new(ctx, translator)
448 }
449}
450
451impl<T: Translator, V: Send + Sync, const P: usize> Unordered for Index<T, V, P> {
452 type Value = V;
453 type Cursor<'a>
454 = Cursor<'a, T::Key, V>
455 where
456 Self: 'a;
457
458 fn get<'a>(&'a self, key: &[u8]) -> impl Iterator<Item = &'a V> + Send + 'a
459 where
460 V: 'a,
461 {
462 let (i, sub) = partition_index_and_sub_key::<P>(key);
463 let k = self.translator.transform(sub);
464 self.partition_values(i, &k).iter()
465 }
466
467 fn get_many<'a, K: AsRef<[u8]>>(&'a self, keys: &[K], mut visit: impl FnMut(usize, &'a V))
468 where
469 V: 'a,
470 {
471 let mut order: Vec<(usize, T::Key, usize)> = keys
475 .iter()
476 .enumerate()
477 .map(|(key_idx, key)| {
478 let (partition, sub) = partition_index_and_sub_key::<P>(key.as_ref());
479 (partition, self.translator.transform(sub), key_idx)
480 })
481 .collect();
482 order.sort_unstable();
483 for (partition, translated, key_idx) in order {
484 for value in self.partition_values(partition, &translated) {
485 visit(key_idx, value);
486 }
487 }
488 }
489
490 fn get_mut<'a>(&'a mut self, key: &[u8]) -> Option<Self::Cursor<'a>> {
491 let (i, sub) = partition_index_and_sub_key::<P>(key);
492 self.get_mut_slot(i, sub)
493 }
494
495 fn get_mut_or_insert<'a>(
496 &'a mut self,
497 key: &[u8],
498 value: Self::Value,
499 ) -> Option<Self::Cursor<'a>> {
500 let (i, sub) = partition_index_and_sub_key::<P>(key);
501 self.get_mut_or_insert_slot(i, sub, value)
502 }
503
504 fn insert(&mut self, key: &[u8], value: Self::Value) {
505 let (i, sub) = partition_index_and_sub_key::<P>(key);
506 let k = self.translator.transform(sub);
507 self.maybe_spill(i);
508 if !self.partitions[i].is_empty() {
509 let run = self.partitions[i].run_range(&k);
510 let new_key = run.is_empty();
511 self.partitions[i].insert_at(run.end, k, value);
512 self.items.inc();
513 if new_key {
514 self.keys.inc();
515 }
516 self.maybe_spill(i);
517 return;
518 }
519
520 if !self.spilled.is_empty()
522 && let hash_map::Entry::Occupied(mut partition) = self.spilled.entry(i)
523 {
524 match partition.get_mut().entry(k) {
525 btree_map::Entry::Occupied(mut run) => run.get_mut().push(value),
526 btree_map::Entry::Vacant(run) => {
527 run.insert(vec![value]);
528 self.keys.inc();
529 }
530 }
531 self.items.inc();
532 return;
533 }
534
535 self.partitions[i].insert_at(0, k, value);
537 self.items.inc();
538 self.keys.inc();
539 self.maybe_spill(i);
540 }
541
542 fn insert_and_retain(
543 &mut self,
544 key: &[u8],
545 value: Self::Value,
546 should_retain: impl Fn(&Self::Value) -> bool,
547 ) {
548 let (i, _) = partition_index_and_sub_key::<P>(key);
549 if let Some(mut cursor) = self.get_mut(key) {
550 cursor.retain(&should_retain);
551 if should_retain(&value) {
552 cursor.insert(value);
553 }
554 } else if should_retain(&value) {
555 self.insert(key, value);
556 }
557 self.maybe_spill(i);
558 }
559
560 fn remove(&mut self, key: &[u8]) {
561 let (i, sub) = partition_index_and_sub_key::<P>(key);
562 let k = self.translator.transform(sub);
563 self.maybe_spill(i);
564 if !self.partitions[i].is_empty() {
565 let run = self.partitions[i].run_range(&k);
566 if run.is_empty() {
567 return;
568 }
569 let n = run.len();
570 self.partitions[i].remove_run(run);
571 self.keys.dec();
572 self.items.dec_by(n as i64);
573 self.pruned.inc_by(n as u64);
574 return;
575 }
576 if !self.spilled.is_empty()
579 && let hash_map::Entry::Occupied(mut partition) = self.spilled.entry(i)
580 && let Some(vals) = partition.get_mut().remove(&k)
581 {
582 let n = vals.len();
583 self.keys.dec();
584 self.items.dec_by(n as i64);
585 self.pruned.inc_by(n as u64);
586 if partition.get().is_empty() {
587 partition.remove();
588 }
589 }
590 }
591
592 #[cfg(test)]
593 fn keys(&self) -> usize {
594 self.keys.get() as usize
595 }
596
597 #[cfg(test)]
598 fn items(&self) -> usize {
599 self.items.get() as usize
600 }
601
602 #[cfg(test)]
603 fn pruned(&self) -> usize {
604 self.pruned.get() as usize
605 }
606}
607
608impl<T: Translator, V: Send + Sync, const P: usize> Ordered for Index<T, V, P> {
609 fn prev_translated_key<'a>(
610 &'a self,
611 key: &[u8],
612 ) -> Option<(impl Iterator<Item = &'a V> + Send + 'a, bool)>
613 where
614 V: 'a,
615 {
616 if self.is_empty() {
618 return None;
619 }
620
621 let (i, sub) = partition_index_and_sub_key::<P>(key);
624 let k = self.translator.transform(sub);
625 if let Some(vals) = self.partition_prev_before(i, &k) {
626 return Some((vals.iter(), false));
627 }
628 for p in (0..i).rev() {
629 if let Some(vals) = self.partition_last(p) {
630 return Some((vals.iter(), false));
631 }
632 }
633 for p in (0..self.partitions.len()).rev() {
634 if let Some(vals) = self.partition_last(p) {
635 return Some((vals.iter(), true));
636 }
637 }
638 None
639 }
640
641 fn next_translated_key<'a>(
642 &'a self,
643 key: &[u8],
644 ) -> Option<(impl Iterator<Item = &'a V> + Send + 'a, bool)>
645 where
646 V: 'a,
647 {
648 if self.is_empty() {
650 return None;
651 }
652
653 let (i, sub) = partition_index_and_sub_key::<P>(key);
656 let k = self.translator.transform(sub);
657 if let Some(vals) = self.partition_next_after(i, &k) {
658 return Some((vals.iter(), false));
659 }
660 for p in i + 1..self.partitions.len() {
661 if let Some(vals) = self.partition_first(p) {
662 return Some((vals.iter(), false));
663 }
664 }
665 for p in 0..self.partitions.len() {
666 if let Some(vals) = self.partition_first(p) {
667 return Some((vals.iter(), true));
668 }
669 }
670 None
671 }
672
673 fn first_translated_key<'a>(&'a self) -> Option<impl Iterator<Item = &'a V> + Send + 'a>
674 where
675 V: 'a,
676 {
677 if self.is_empty() {
679 return None;
680 }
681
682 for p in 0..self.partitions.len() {
684 if let Some(vals) = self.partition_first(p) {
685 return Some(vals.iter());
686 }
687 }
688 None
689 }
690
691 fn last_translated_key<'a>(&'a self) -> Option<impl Iterator<Item = &'a V> + Send + 'a>
692 where
693 V: 'a,
694 {
695 if self.is_empty() {
697 return None;
698 }
699
700 for p in (0..self.partitions.len()).rev() {
702 if let Some(vals) = self.partition_last(p) {
703 return Some(vals.iter());
704 }
705 }
706 None
707 }
708}
709
710#[cfg(test)]
711mod tests {
712 use super::*;
713 use crate::translator::OneCap;
714 use commonware_formatting::hex;
715 use commonware_macros::test_traced;
716 use commonware_runtime::{Runner, Supervisor as _, deterministic};
717
718 fn new_index(context: deterministic::Context) -> Index<OneCap, u64, 1> {
719 Index::new(context, OneCap)
720 }
721
722 fn new_index_spilling(context: deterministic::Context) -> Index<OneCap, u64, 1> {
726 Index::with_threshold(context, OneCap, 2)
727 }
728
729 #[test_traced]
730 fn test_empty_and_sparse_nav() {
731 deterministic::Runner::default().start(|context| async move {
732 let mut index = new_index(context);
733
734 assert!(index.first_translated_key().is_none());
737 assert!(index.last_translated_key().is_none());
738 assert!(index.prev_translated_key(&[0x80, 0x00]).is_none());
739 assert!(index.next_translated_key(&[0x80, 0x00]).is_none());
740
741 index.insert(&[0x05, 0x01], 1);
744 index.insert(&[0xF0, 0x02], 2);
745 assert_eq!(index.keys(), 2);
746 assert_eq!(
747 index
748 .first_translated_key()
749 .unwrap()
750 .copied()
751 .collect::<Vec<_>>(),
752 vec![1]
753 );
754 assert_eq!(
755 index
756 .last_translated_key()
757 .unwrap()
758 .copied()
759 .collect::<Vec<_>>(),
760 vec![2]
761 );
762
763 let (it, wrapped) = index.next_translated_key(&[0x05, 0x01]).unwrap();
765 assert_eq!((it.copied().collect::<Vec<_>>(), wrapped), (vec![2], false));
766 let (it, wrapped) = index.next_translated_key(&[0xF0, 0x02]).unwrap();
767 assert_eq!((it.copied().collect::<Vec<_>>(), wrapped), (vec![1], true));
768
769 let (it, wrapped) = index.prev_translated_key(&[0xF0, 0x02]).unwrap();
771 assert_eq!((it.copied().collect::<Vec<_>>(), wrapped), (vec![1], false));
772 let (it, wrapped) = index.prev_translated_key(&[0x05, 0x01]).unwrap();
773 assert_eq!((it.copied().collect::<Vec<_>>(), wrapped), (vec![2], true));
774
775 let (it, wrapped) = index.prev_translated_key(&[0x80, 0x00]).unwrap();
777 assert_eq!((it.copied().collect::<Vec<_>>(), wrapped), (vec![1], false));
778 let (it, wrapped) = index.next_translated_key(&[0x80, 0x00]).unwrap();
779 assert_eq!((it.copied().collect::<Vec<_>>(), wrapped), (vec![2], false));
780
781 index.remove(&[0x05, 0x01]);
783 index.remove(&[0xF0, 0x02]);
784 assert_eq!(index.keys(), 0);
785 assert!(index.prev_translated_key(&[0x80, 0x00]).is_none());
786 assert!(index.next_translated_key(&[0x80, 0x00]).is_none());
787 });
788 }
789
790 #[test_traced]
791 fn test_spill_transition() {
792 deterministic::Runner::default().start(|context| async move {
793 let mut index = new_index_spilling(context);
794 index.insert(&[0x10, 0x01], 1);
796 assert_eq!(index.spilled_count(), 0);
797 index.insert(&[0x10, 0x02], 2); assert_eq!(index.spilled_count(), 1);
799 index.insert(&[0x10, 0x03], 3); assert_eq!(index.spilled_count(), 1);
801 assert_eq!(index.keys(), 3);
802 assert_eq!(index.items(), 3);
803
804 assert_eq!(
806 index.get(&[0x10, 0x01]).copied().collect::<Vec<_>>(),
807 vec![1]
808 );
809 index.insert(&[0x10, 0x02], 22);
810 assert_eq!(
811 index.get(&[0x10, 0x02]).copied().collect::<Vec<_>>(),
812 vec![2, 22]
813 );
814 assert_eq!(index.items(), 4);
815
816 index.insert(&[0x20, 0x05], 5);
818 assert_eq!(index.spilled_count(), 1);
819 assert_eq!(
820 index.get(&[0x20, 0x05]).copied().collect::<Vec<_>>(),
821 vec![5]
822 );
823 });
824 }
825
826 #[test_traced]
827 fn test_spill_after_cursor_growth() {
828 deterministic::Runner::default().start(|context| async move {
829 let mut index = new_index_spilling(context);
830 let key = [0x10, 0x01];
831
832 index.insert(&key, 1);
833 {
834 let mut cursor = index.get_mut(&key).unwrap();
835 assert_eq!(cursor.next().copied(), Some(1));
836 assert_eq!(cursor.next(), None);
837 cursor.insert(2);
838 }
839 assert_eq!(index.spilled_count(), 0);
840
841 index.insert(&key, 3);
843 assert_eq!(index.spilled_count(), 1);
844 assert_eq!(index.get(&key).copied().collect::<Vec<_>>(), vec![1, 2, 3]);
845
846 let other = [0x20, 0x01];
848 index.insert(&other, 4);
849 index.insert_and_retain(&other, 5, |_| true);
850 assert_eq!(index.spilled_count(), 2);
851 assert_eq!(index.get(&other).copied().collect::<Vec<_>>(), vec![4, 5]);
852
853 let third = [0x30, 0x01];
856 index.insert(&third, 6);
857 {
858 let mut cursor = index.get_mut(&third).unwrap();
859 assert_eq!(cursor.next().copied(), Some(6));
860 assert_eq!(cursor.next(), None);
861 cursor.insert(7);
862 }
863 assert_eq!(index.spilled_count(), 2);
864 {
865 let mut cursor = index.get_mut(&third).unwrap();
866 assert_eq!(cursor.next().copied(), Some(6));
867 assert_eq!(cursor.next().copied(), Some(7));
868 assert_eq!(cursor.next(), None);
869 }
870 assert_eq!(index.spilled_count(), 3);
871 assert_eq!(index.get(&third).copied().collect::<Vec<_>>(), vec![6, 7]);
872
873 let fourth = [0x40, 0x01];
875 index.insert(&fourth, 8);
876 {
877 let mut cursor = index.get_mut(&fourth).unwrap();
878 assert_eq!(cursor.next().copied(), Some(8));
879 assert_eq!(cursor.next(), None);
880 cursor.insert(9);
881 }
882 assert_eq!(index.spilled_count(), 3);
883 index.remove(&[0x40, 0x02]);
884 assert_eq!(index.spilled_count(), 4);
885 assert_eq!(index.get(&fourth).copied().collect::<Vec<_>>(), vec![8, 9]);
886 });
887 }
888
889 #[test_traced]
890 fn test_spill_after_get_mut_or_insert_cursor_growth() {
891 deterministic::Runner::default().start(|context| async move {
892 let mut index = new_index_spilling(context);
893 let key = [0x10, 0x01];
894
895 index.insert(&key, 1);
896 {
897 let mut cursor = index.get_mut_or_insert(&key, 2).unwrap();
898 assert_eq!(cursor.next().copied(), Some(1));
899 assert_eq!(cursor.next(), None);
900 cursor.insert(2);
901 }
902 assert_eq!(index.spilled_count(), 0);
903
904 assert!(index.get_mut_or_insert(&key, 3).is_some());
906 assert_eq!(index.spilled_count(), 1);
907 assert_eq!(index.get(&key).copied().collect::<Vec<_>>(), vec![1, 2]);
908 });
909 }
910
911 #[test_traced]
912 fn test_spill_nav() {
913 deterministic::Runner::default().start(|context| async move {
914 let mut index = new_index_spilling(context);
915 index.insert(&[0x10, 0x01], 1);
918 index.insert(&[0x10, 0x02], 2);
919 index.insert(&[0x20, 0x05], 5);
920 index.insert(&[0x30, 0x07], 7);
921 index.insert(&[0x30, 0x08], 8);
922 assert_eq!(index.spilled_count(), 2); assert_eq!(index.first_translated_key().unwrap().next(), Some(&1));
925 assert_eq!(index.last_translated_key().unwrap().next(), Some(&8));
926
927 let (mut it, wrapped) = index.next_translated_key(&[0x10, 0x01]).unwrap();
929 assert!(!wrapped);
930 assert_eq!(it.next(), Some(&2));
931 let (mut it, wrapped) = index.next_translated_key(&[0x10, 0x02]).unwrap();
933 assert!(!wrapped);
934 assert_eq!(it.next(), Some(&5));
935 let (mut it, wrapped) = index.next_translated_key(&[0x20, 0x05]).unwrap();
937 assert!(!wrapped);
938 assert_eq!(it.next(), Some(&7));
939 let (mut it, wrapped) = index.prev_translated_key(&[0x30, 0x07]).unwrap();
941 assert!(!wrapped);
942 assert_eq!(it.next(), Some(&5));
943 let (mut it, wrapped) = index.prev_translated_key(&[0x20, 0x05]).unwrap();
945 assert!(!wrapped);
946 assert_eq!(it.next(), Some(&2));
947 let (mut it, wrapped) = index.next_translated_key(&[0x30, 0x08]).unwrap();
949 assert!(wrapped);
950 assert_eq!(it.next(), Some(&1));
951 });
952 }
953
954 #[test_traced]
955 fn test_spill_despill_on_full_drain() {
956 deterministic::Runner::default().start(|context| async move {
957 let mut index = new_index_spilling(context);
958 index.insert(&[0x10, 0x01], 1);
959 index.insert(&[0x10, 0x02], 2); assert_eq!(index.spilled_count(), 1);
961
962 index.remove(&[0x10, 0x01]);
963 assert_eq!(index.spilled_count(), 1); index.remove(&[0x10, 0x02]);
965 assert_eq!(index.spilled_count(), 0); assert_eq!(index.keys(), 0);
967
968 index.insert(&[0x10, 0x09], 9);
970 assert_eq!(index.spilled_count(), 0);
971 assert_eq!(
972 index.get(&[0x10, 0x09]).copied().collect::<Vec<_>>(),
973 vec![9]
974 );
975 });
976 }
977
978 #[test_traced]
979 fn test_spill_full_lifecycle() {
980 deterministic::Runner::default().start(|context| async move {
981 let mut index = new_index_spilling(context);
982
983 assert_eq!(index.spilled_count(), 0);
985 assert_eq!(index.keys(), 0);
986 assert_eq!(index.items(), 0);
987
988 index.insert(&[0x10, 0x01], 1);
990 assert_eq!(index.spilled_count(), 0);
991
992 index.insert(&[0x10, 0x02], 2);
994 assert_eq!(index.spilled_count(), 1);
995 assert_eq!(index.spills(), 1);
996 assert_eq!(index.keys(), 2);
997 assert_eq!(index.items(), 2);
998
999 {
1002 let mut cursor = index.get_mut(&[0x10, 0x01]).unwrap();
1003 assert_eq!(cursor.next().copied(), Some(1));
1004 cursor.delete();
1005 }
1006 assert_eq!(index.spilled_count(), 1); {
1008 let mut cursor = index.get_mut(&[0x10, 0x02]).unwrap();
1009 assert_eq!(cursor.next().copied(), Some(2));
1010 cursor.delete();
1011 }
1012 assert_eq!(index.spilled_count(), 0); assert_eq!(index.spills(), 1); assert_eq!(index.keys(), 0);
1015 assert_eq!(index.items(), 0);
1016
1017 index.insert(&[0x10, 0x03], 3);
1019 assert_eq!(index.spilled_count(), 0);
1020 index.insert(&[0x10, 0x04], 4);
1021 assert_eq!(index.spilled_count(), 1);
1022 assert_eq!(index.spills(), 2); assert_eq!(
1024 index.get(&[0x10, 0x03]).copied().collect::<Vec<_>>(),
1025 vec![3]
1026 );
1027 assert_eq!(
1028 index.get(&[0x10, 0x04]).copied().collect::<Vec<_>>(),
1029 vec![4]
1030 );
1031
1032 index.remove(&[0x10, 0x03]);
1034 assert_eq!(index.spilled_count(), 1); index.remove(&[0x10, 0x04]);
1036 assert_eq!(index.spilled_count(), 0);
1037 assert_eq!(index.spills(), 2); assert_eq!(index.keys(), 0);
1039 assert_eq!(index.items(), 0);
1040
1041 assert_eq!(index.pruned(), 4);
1043 });
1044 }
1045
1046 #[test_traced]
1047 fn test_spill_counts_live() {
1048 deterministic::Runner::default().start(|context| async move {
1049 let mut full = new_index_spilling(context.child("full"));
1052 assert_eq!(full.spills(), 0);
1053
1054 let mut worker = full.new_range(0, full.partition_count());
1059 worker.get_mut_or_insert(&[0x10, 0x01], 1);
1060 worker.get_mut_or_insert(&[0x10, 0x02], 2); worker.index.insert(&[0x20, 0x01], 3);
1062 worker.index.insert(&[0x20, 0x02], 4); worker.index.remove(&[0x20, 0x01]);
1064 worker.index.remove(&[0x20, 0x02]); assert_eq!(worker.index.spilled_count(), 1);
1066 assert_eq!(full.spills(), 2); full.install_range(worker);
1070 assert_eq!(full.spilled_count(), 1);
1071 assert_eq!(full.spills(), 2);
1072 });
1073 }
1074
1075 #[test_traced]
1076 fn test_worker_prunes_count_live() {
1077 deterministic::Runner::default().start(|context| async move {
1078 let mut full = new_index(context.child("full"));
1079 assert_eq!(full.pruned(), 0);
1080
1081 let mut worker = full.new_range(0, full.partition_count());
1087 worker.index.insert(&[0x10, 0x01], 1);
1088 worker.index.insert(&[0x10, 0x01], 2);
1089 {
1090 let mut cursor = worker.get_mut(&[0x10, 0x01]).unwrap();
1091 while cursor.next().is_some() {
1092 cursor.delete();
1093 }
1094 }
1095 assert_eq!(full.pruned(), 2);
1096
1097 full.install_range(worker);
1099 assert_eq!(full.pruned(), 2);
1100 });
1101 }
1102
1103 #[test_traced]
1107 fn test_range_for_each_value_visits_all_values_once() {
1108 deterministic::Runner::default().start(|context| async move {
1109 let full = new_index_spilling(context.child("full"));
1110
1111 let mut worker = full.new_range(0x80, 2);
1115 assert!(worker.get_mut_or_insert(&[0x80, 0x01], 1).is_none());
1116 assert!(worker.get_mut_or_insert(&[0x80, 0x02, 0xAA], 2).is_none());
1117 assert!(worker.get_mut_or_insert(&[0x80, 0x02, 0xBB], 3).is_some());
1118 {
1119 let mut cursor = worker.get_mut(&[0x80, 0x02, 0xBB]).unwrap();
1120 cursor.next();
1121 cursor.insert(3);
1122 }
1123 assert!(worker.get_mut_or_insert(&[0x81, 0x07], 4).is_none());
1124 assert_eq!(worker.index.spilled_count(), 1);
1125
1126 let mut seen = Vec::new();
1127 worker.for_each_value(|v| seen.push(*v));
1128 seen.sort_unstable();
1129 assert_eq!(seen, vec![1, 2, 3, 4]);
1130 });
1131 }
1132
1133 #[test_traced]
1138 fn test_install_range_nonzero_offset() {
1139 deterministic::Runner::default().start(|context| async move {
1140 let mut full = new_index_spilling(context.child("full"));
1143
1144 let mut worker = full.new_range(0x80, 2);
1147 worker.get_mut_or_insert(&[0x80, 0x01], 1);
1148 worker.get_mut_or_insert(&[0x80, 0x02], 2);
1149 worker.get_mut_or_insert(&[0x81, 0x07], 3);
1150 assert_eq!(worker.index.spilled_count(), 1);
1151 assert_eq!(worker.index.keys(), 3);
1152
1153 full.install_range(worker);
1156 assert_eq!(full.spilled_count(), 1);
1157 assert_eq!(full.keys(), 3);
1158 assert_eq!(full.items(), 3);
1159 assert_eq!(
1160 full.get(&[0x80, 0x01]).copied().collect::<Vec<_>>(),
1161 vec![1]
1162 );
1163 assert_eq!(
1164 full.get(&[0x80, 0x02]).copied().collect::<Vec<_>>(),
1165 vec![2]
1166 );
1167 assert_eq!(
1168 full.get(&[0x81, 0x07]).copied().collect::<Vec<_>>(),
1169 vec![3]
1170 );
1171
1172 assert!(full.get(&[0x00, 0x01]).next().is_none());
1175 assert!(full.get(&[0x01, 0x07]).next().is_none());
1176 });
1177 }
1178
1179 #[test_traced]
1180 fn test_spill_get_mut_or_insert() {
1181 deterministic::Runner::default().start(|context| async move {
1182 let mut index = new_index_spilling(context);
1183 index.insert(&[0x10, 0x01], 1);
1184 index.insert(&[0x10, 0x02], 2); assert_eq!(index.spilled_count(), 1);
1186 assert_eq!(index.keys(), 2);
1187 assert_eq!(index.items(), 2);
1188
1189 {
1192 let mut cursor = index.get_mut_or_insert(&[0x10, 0x01], 99).unwrap();
1193 assert_eq!(cursor.next().copied(), Some(1));
1194 assert!(cursor.next().is_none());
1195 }
1196 assert_eq!(index.keys(), 2);
1197 assert_eq!(index.items(), 2);
1198 assert_eq!(
1199 index.get(&[0x10, 0x01]).copied().collect::<Vec<_>>(),
1200 vec![1]
1201 );
1202
1203 assert!(index.get_mut_or_insert(&[0x10, 0x03], 3).is_none());
1206 assert_eq!(index.spilled_count(), 1);
1207 assert_eq!(index.keys(), 3);
1208 assert_eq!(index.items(), 3);
1209 assert_eq!(
1210 index.get(&[0x10, 0x03]).copied().collect::<Vec<_>>(),
1211 vec![3]
1212 );
1213 });
1214 }
1215
1216 #[test_traced]
1217 #[should_panic(expected = "must call Cursor::next()")]
1218 fn test_spill_cursor_delete_before_next_panics() {
1219 deterministic::Runner::default().start(|context| async move {
1220 let mut index = new_index_spilling(context);
1221 index.insert(&[0x10, 0x01], 1);
1222 index.insert(&[0x10, 0x02], 2); let mut cursor = index.get_mut(&[0x10, 0x01]).unwrap(); cursor.delete();
1225 });
1226 }
1227
1228 #[test_traced]
1229 fn test_soa_basic() {
1230 deterministic::Runner::default().start(|context| async move {
1231 let mut index = new_index(context);
1232 assert_eq!(index.keys(), 0);
1233
1234 let key = b"duplicate".as_slice();
1235 index.insert(key, 1);
1236 index.insert(key, 2);
1237 index.insert(key, 3);
1238 assert_eq!(index.keys(), 1);
1239 assert_eq!(index.items(), 3);
1240 assert_eq!(index.get(key).copied().collect::<Vec<_>>(), vec![1, 2, 3]);
1241
1242 {
1243 let mut cursor = index.get_mut(key).unwrap();
1244 assert_eq!(*cursor.next().unwrap(), 1);
1245 assert_eq!(*cursor.next().unwrap(), 2);
1246 assert_eq!(*cursor.next().unwrap(), 3);
1247 assert!(cursor.next().is_none());
1248 }
1249
1250 index.insert(key, 3);
1251 index.insert(key, 4);
1252 index.retain(key, |i| *i != 3);
1253 assert_eq!(index.get(key).copied().collect::<Vec<_>>(), vec![1, 2, 4]);
1254
1255 index.retain(key, |_| false);
1256 assert_eq!(
1257 index.get(key).copied().collect::<Vec<_>>(),
1258 Vec::<u64>::new()
1259 );
1260 assert_eq!(index.keys(), 0);
1261 assert!(index.get_mut(key).is_none());
1262
1263 index.retain(key, |_| false);
1265 });
1266 }
1267
1268 #[test_traced]
1269 fn test_soa_cursor_find() {
1270 deterministic::Runner::default().start(|context| async move {
1271 let mut index = new_index(context);
1272 let key = b"test_key";
1273 for v in [10u64, 20, 30, 40] {
1274 index.insert(key, v);
1275 }
1276
1277 {
1278 let mut cursor = index.get_mut(key).unwrap();
1279 assert!(cursor.find(|&v| v == 30));
1280 cursor.update(35);
1281 }
1282 let values: Vec<u64> = index.get(key).copied().collect();
1283 assert!(values.contains(&35) && !values.contains(&30));
1284
1285 {
1286 let mut cursor = index.get_mut(key).unwrap();
1287 assert!(!cursor.find(|&v| v == 100));
1288 assert!(cursor.next().is_none());
1289 }
1290
1291 {
1292 let mut cursor = index.get_mut(key).unwrap();
1293 assert!(cursor.find(|&v| v == 20));
1294 cursor.delete();
1295 }
1296 let values: Vec<u64> = index.get(key).copied().collect();
1297 assert!(!values.contains(&20));
1298 assert_eq!(values.len(), 3);
1299 });
1300 }
1301
1302 #[test_traced]
1303 fn test_soa_get_many_and_partitions() {
1304 deterministic::Runner::default().start(|context| async move {
1305 let mut index = new_index(context);
1306 index.insert(b"ab", 1);
1308 index.insert(b"ab", 2);
1309 index.insert(b"abX", 3);
1310 index.insert(b"zz", 4);
1311
1312 let keys: Vec<&[u8]> = vec![b"zz", b"missing", b"ab", b"zz"];
1313 let mut visits: Vec<Vec<u64>> = vec![Vec::new(); keys.len()];
1314 index.get_many(&keys, |key_idx, value| visits[key_idx].push(*value));
1315 assert_eq!(visits[0], vec![4]);
1316 assert!(visits[1].is_empty());
1317 assert_eq!(visits[2], vec![1, 2, 3]);
1318 assert_eq!(visits[3], vec![4]);
1319 });
1320 }
1321
1322 #[test_traced]
1323 fn test_soa_insert_and_retain() {
1324 deterministic::Runner::default().start(|context| async move {
1325 let mut index = new_index(context);
1326 index.insert(b"k", 1u64);
1328 index.insert_and_retain(b"k", 2, |_| true);
1329 assert_eq!(index.get(b"k").copied().collect::<Vec<_>>(), vec![1, 2]);
1330
1331 index.insert_and_retain(b"k", 9, |v| *v != 9);
1333 assert_eq!(index.get(b"k").copied().collect::<Vec<_>>(), vec![1, 2]);
1334
1335 index.insert_and_retain(b"k", 9, |_| false);
1337 assert!(index.get_mut(b"k").is_none());
1338 assert_eq!(index.keys(), 0);
1339
1340 index.insert_and_retain(b"new", 7, |_| true);
1342 assert_eq!(index.get(b"new").copied().collect::<Vec<_>>(), vec![7]);
1343 assert_eq!(index.keys(), 1);
1344 });
1345 }
1346
1347 #[test_traced]
1348 fn test_soa_remove() {
1349 deterministic::Runner::default().start(|context| async move {
1350 let mut index = new_index(context);
1351 index.insert(b"k", 1u64);
1352 index.insert(b"k", 2);
1353 index.insert(b"other", 3);
1354 assert_eq!(index.items(), 3);
1355 assert_eq!(index.keys(), 2);
1356
1357 index.remove(b"k");
1358 assert!(index.get_mut(b"k").is_none());
1359 assert_eq!(index.keys(), 1);
1360 assert_eq!(index.items(), 1);
1361 assert_eq!(index.pruned(), 2);
1362 assert_eq!(index.get(b"other").copied().collect::<Vec<_>>(), vec![3]);
1363
1364 index.remove(b"missing"); assert_eq!(index.keys(), 1);
1366 });
1367 }
1368
1369 #[test_traced]
1370 fn test_soa_ordered() {
1371 deterministic::Runner::default().start(|context| async move {
1372 let mut index = new_index(context);
1373 assert!(index.first_translated_key().is_none());
1374 assert!(index.last_translated_key().is_none());
1375 assert!(index.next_translated_key(b"key").is_none());
1376 assert!(index.prev_translated_key(b"key").is_none());
1377
1378 let k1 = &hex!("0x0b02AA"); let k2 = &hex!("0x1c04CC"); let k2_collides = &hex!("0x1c0411"); let k3 = &hex!("0x2d06EE"); index.insert(k1, 1);
1384 index.insert(k2, 21);
1385 index.insert(k2_collides, 22);
1386 index.insert(k3, 3);
1387 assert_eq!(index.keys(), 3);
1388
1389 assert_eq!(index.first_translated_key().unwrap().next(), Some(&1));
1390 assert_eq!(index.last_translated_key().unwrap().next(), Some(&3));
1391
1392 let (mut it, wrapped) = index.next_translated_key(&[0x00]).unwrap();
1394 assert!(!wrapped);
1395 assert_eq!(it.next(), Some(&1));
1396 assert_eq!(it.next(), None);
1397
1398 let (mut it, wrapped) = index.next_translated_key(&hex!("0x0b02F2")).unwrap();
1400 assert!(!wrapped);
1401 assert_eq!(it.next(), Some(&21));
1402 assert_eq!(it.next(), Some(&22));
1403 assert_eq!(it.next(), None);
1404
1405 let (mut it, wrapped) = index.next_translated_key(k3).unwrap();
1407 assert!(wrapped);
1408 assert_eq!(it.next(), Some(&1));
1409
1410 let (mut it, wrapped) = index.prev_translated_key(k1).unwrap();
1412 assert!(wrapped);
1413 assert_eq!(it.next(), Some(&3));
1414
1415 let (mut it, wrapped) = index.prev_translated_key(&hex!("0x1d0102")).unwrap();
1417 assert!(!wrapped);
1418 assert_eq!(it.next(), Some(&21));
1419 assert_eq!(it.next(), Some(&22));
1420 assert_eq!(it.next(), None);
1421 });
1422 }
1423
1424 #[test_traced]
1425 fn test_soa_ordered_exhaustive_traversal() {
1426 deterministic::Runner::default().start(|context| async move {
1427 let mut index = new_index(context);
1428
1429 let prefixes = [0x00u8, 0x05, 0xAA, 0xFF];
1433 let subkeys = [0x00u8, 0x80, 0xFF];
1434 let mut keys: Vec<[u8; 2]> = Vec::new();
1435 for &p in &prefixes {
1436 for &s in &subkeys {
1437 keys.push([p, s]);
1438 }
1439 }
1440 let value_of = |k: &[u8; 2]| ((k[0] as u64) << 8) | k[1] as u64;
1441 let n = keys.len();
1442
1443 let mut scrambled = keys.clone();
1445 scrambled.reverse();
1446 scrambled.rotate_left(5);
1447 for k in &scrambled {
1448 index.insert(k, value_of(k));
1449 }
1450 assert_eq!(index.keys(), n);
1451
1452 assert_eq!(
1453 index.first_translated_key().unwrap().next(),
1454 Some(&value_of(&keys[0]))
1455 );
1456 assert_eq!(
1457 index.last_translated_key().unwrap().next(),
1458 Some(&value_of(&keys[n - 1]))
1459 );
1460
1461 for i in 0..n {
1464 let next = value_of(&keys[(i + 1) % n]);
1465 let (mut it, wrapped) = index.next_translated_key(&keys[i]).unwrap();
1466 assert_eq!(wrapped, i + 1 == n, "next wrap at index {i}");
1467 assert_eq!(it.next(), Some(&next), "next at {i}");
1468 assert_eq!(it.next(), None);
1469
1470 let prev = value_of(&keys[(i + n - 1) % n]);
1471 let (mut it, wrapped) = index.prev_translated_key(&keys[i]).unwrap();
1472 assert_eq!(wrapped, i == 0, "prev wrap at index {i}");
1473 assert_eq!(it.next(), Some(&prev), "prev at {i}");
1474 assert_eq!(it.next(), None);
1475 }
1476 });
1477 }
1478}