1mod allocator;
52mod iter;
53mod node;
54use crate::btreemap::iter::{IterInternal, KeysIter, ValuesIter};
55use crate::{
56 storable::Bound as StorableBound,
57 types::{Address, NULL},
58 Memory, Storable,
59};
60use allocator::Allocator;
61pub use iter::Iter;
62use node::{DerivedPageSize, Entry, Node, NodeType, PageSize, Version};
63use std::borrow::Cow;
64use std::marker::PhantomData;
65use std::ops::{Bound, RangeBounds};
66
67#[cfg(test)]
68mod proptests;
69
70const MAGIC: &[u8; 3] = b"BTR";
71const LAYOUT_VERSION: u8 = 1;
72const LAYOUT_VERSION_2: u8 = 2;
73const PACKED_HEADER_SIZE: usize = 28;
75const ALLOCATOR_OFFSET: usize = 52;
77
78const DEFAULT_PAGE_SIZE: u32 = 1024;
80
81const PAGE_SIZE_VALUE_MARKER: u32 = u32::MAX;
83
84pub struct BTreeMap<K, V, M>
228where
229 K: Storable + Ord + Clone,
230 V: Storable,
231 M: Memory,
232{
233 root_addr: Address,
236
237 version: Version,
238
239 allocator: Allocator<M>,
241
242 length: u64,
244
245 _phantom: PhantomData<(K, V)>,
247}
248
249#[derive(PartialEq, Debug)]
250struct BTreeHeader {
252 version: Version,
253 root_addr: Address,
254 length: u64,
255 }
257
258impl<K, V, M> BTreeMap<K, V, M>
259where
260 K: Storable + Ord + Clone,
261 V: Storable,
262 M: Memory,
263{
264 pub fn init(memory: M) -> Self {
269 if memory.size() == 0 {
270 return BTreeMap::new(memory);
272 }
273
274 let mut dst = vec![0; 3];
276 memory.read(0, &mut dst);
277 if dst != MAGIC {
278 BTreeMap::new(memory)
280 } else {
281 BTreeMap::load(memory)
283 }
284 }
285
286 #[cfg(test)]
290 pub fn init_v1(memory: M) -> Self {
291 if memory.size() == 0 {
292 return BTreeMap::new_v1(memory);
294 }
295
296 let mut dst = vec![0; 3];
298 memory.read(0, &mut dst);
299 if dst != MAGIC {
300 BTreeMap::new_v1(memory)
302 } else {
303 BTreeMap::load_helper(memory, false)
306 }
307 }
308
309 pub fn new(memory: M) -> Self {
321 let page_size = match (K::BOUND, V::BOUND) {
322 (
324 StorableBound::Bounded {
325 max_size: max_key_size,
326 ..
327 },
328 StorableBound::Bounded {
329 max_size: max_value_size,
330 ..
331 },
332 ) => {
333 let max_node_size = Node::<K>::max_size(max_key_size, max_value_size).get();
335
336 PageSize::Value((max_node_size * 3 / 4) as u32)
341 }
342 _ => PageSize::Value(DEFAULT_PAGE_SIZE),
344 };
345
346 let btree = Self {
347 root_addr: NULL,
348 allocator: Allocator::new(
349 memory,
350 Address::from(ALLOCATOR_OFFSET as u64),
351 page_size.get().into(),
352 ),
353 version: Version::V2(page_size),
354 length: 0,
355 _phantom: PhantomData,
356 };
357
358 btree.save_header();
359 btree
360 }
361
362 #[cfg(test)]
366 pub fn new_v1(memory: M) -> Self {
367 let max_key_size = K::BOUND.max_size();
368 let max_value_size = V::BOUND.max_size();
369
370 let btree = Self {
371 root_addr: NULL,
372 allocator: Allocator::new(
373 memory,
374 Address::from(ALLOCATOR_OFFSET as u64),
375 Node::<K>::max_size(max_key_size, max_value_size),
376 ),
377 version: Version::V1(DerivedPageSize {
378 max_key_size,
379 max_value_size,
380 }),
381 length: 0,
382 _phantom: PhantomData,
383 };
384
385 btree.save_header();
386 btree
387 }
388
389 pub fn load(memory: M) -> Self {
391 Self::load_helper(memory, true)
392 }
393
394 fn load_helper(memory: M, migrate_to_v2: bool) -> Self {
396 let header = Self::read_header(&memory);
398
399 let version = match header.version {
400 Version::V1(derived_page_size) => {
401 if migrate_to_v2 {
403 Version::V2(PageSize::Derived(derived_page_size))
404 } else {
405 let max_key_size = derived_page_size.max_key_size;
407 let max_value_size = derived_page_size.max_value_size;
408
409 assert!(
410 K::BOUND.max_size() <= max_key_size,
411 "max_key_size must be <= {max_key_size}",
412 );
413
414 assert!(
415 V::BOUND.max_size() <= max_value_size,
416 "max_value_size must be <= {max_value_size}"
417 );
418
419 Version::V1(derived_page_size)
420 }
421 }
422 other => other,
423 };
424
425 let allocator_addr = Address::from(ALLOCATOR_OFFSET as u64);
426 Self {
427 root_addr: header.root_addr,
428 allocator: Allocator::load(memory, allocator_addr),
429 version,
430 length: header.length,
431 _phantom: PhantomData,
432 }
433 }
434
435 fn read_header(memory: &M) -> BTreeHeader {
437 let mut buf = [0; PACKED_HEADER_SIZE];
439 memory.read(0, &mut buf);
440
441 assert_eq!(&buf[0..3], MAGIC, "Bad magic.");
442
443 match buf[3] {
444 LAYOUT_VERSION => {
445 BTreeHeader {
447 version: Version::V1(DerivedPageSize {
448 max_key_size: u32::from_le_bytes(buf[4..8].try_into().unwrap()),
449 max_value_size: u32::from_le_bytes(buf[8..12].try_into().unwrap()),
450 }),
451 root_addr: Address::from(u64::from_le_bytes(buf[12..20].try_into().unwrap())),
452 length: u64::from_le_bytes(buf[20..28].try_into().unwrap()),
453 }
454 }
455 LAYOUT_VERSION_2 => {
456 let page_size = {
458 let a = u32::from_le_bytes(buf[4..8].try_into().unwrap());
460 let b = u32::from_le_bytes(buf[8..12].try_into().unwrap());
461
462 if b == PAGE_SIZE_VALUE_MARKER {
463 PageSize::Value(a)
465 } else {
466 PageSize::Derived(DerivedPageSize {
468 max_key_size: a,
469 max_value_size: b,
470 })
471 }
472 };
473
474 BTreeHeader {
476 version: Version::V2(page_size),
477 root_addr: Address::from(u64::from_le_bytes(buf[12..20].try_into().unwrap())),
478 length: u64::from_le_bytes(buf[20..28].try_into().unwrap()),
479 }
480 }
481 version => {
482 panic!("Unsupported version: {version}.");
483 }
484 }
485 }
486
487 pub fn insert(&mut self, key: K, value: V) -> Option<V> {
495 let value = value.to_bytes_checked().into_owned();
496
497 let root = if self.root_addr == NULL {
498 let node = self.allocate_node(NodeType::Leaf);
500 self.root_addr = node.address();
501 self.save_header();
502 node
503 } else {
504 let mut root = self.load_node(self.root_addr);
506
507 if let Ok(idx) = root.search(&key, self.memory()) {
509 return Some(V::from_bytes(Cow::Owned(
511 self.update_value(&mut root, idx, value),
512 )));
513 }
514
515 if root.is_full() {
521 let mut new_root = self.allocate_node(NodeType::Internal);
523
524 new_root.push_child(self.root_addr);
526
527 self.root_addr = new_root.address();
529 self.save_header();
530
531 self.split_child(&mut new_root, 0);
533
534 new_root
535 } else {
536 root
537 }
538 };
539
540 self.insert_nonfull(root, key, value)
541 .map(Cow::Owned)
542 .map(V::from_bytes)
543 }
544
545 fn insert_nonfull(&mut self, mut node: Node<K>, key: K, value: Vec<u8>) -> Option<Vec<u8>> {
547 assert!(!node.is_full());
549
550 match node.search(&key, self.memory()) {
552 Ok(idx) => {
553 Some(self.update_value(&mut node, idx, value))
555 }
556 Err(idx) => {
557 match node.node_type() {
560 NodeType::Leaf => {
561 node.insert_entry(idx, (key, value));
564 self.save_node(&mut node);
565
566 self.length += 1;
568 self.save_header();
569
570 None
572 }
573 NodeType::Internal => {
574 let mut child = self.load_node(node.child(idx));
577
578 if child.is_full() {
579 if let Ok(idx) = child.search(&key, self.memory()) {
581 return Some(self.update_value(&mut child, idx, value));
583 }
584
585 self.split_child(&mut node, idx);
587
588 let idx = node.search(&key, self.memory()).unwrap_or_else(|idx| idx);
591 child = self.load_node(node.child(idx));
592 }
593
594 assert!(!child.is_full());
596
597 self.insert_nonfull(child, key, value)
598 }
599 }
600 }
601 }
602 }
603
604 fn split_child(&mut self, node: &mut Node<K>, full_child_idx: usize) {
622 assert!(!node.is_full());
624
625 let mut full_child = self.load_node(node.child(full_child_idx));
627 assert!(full_child.is_full());
628
629 let mut sibling = self.allocate_node(full_child.node_type());
631 assert_eq!(sibling.node_type(), full_child.node_type());
632
633 node.insert_child(full_child_idx + 1, sibling.address());
635
636 let (median_key, median_value) = full_child.split(&mut sibling, self.memory());
637
638 node.insert_entry(full_child_idx, (median_key, median_value));
639
640 self.save_node(&mut sibling);
641 self.save_node(&mut full_child);
642 self.save_node(node);
643 }
644
645 pub fn get(&self, key: &K) -> Option<V> {
647 if self.root_addr == NULL {
648 return None;
649 }
650 self.traverse(self.root_addr, key, |node, idx| {
651 node.into_entry(idx, self.memory()).1 })
653 .map(Cow::Owned)
654 .map(V::from_bytes)
655 }
656
657 pub fn contains_key(&self, key: &K) -> bool {
659 self.root_addr != NULL && self.traverse(self.root_addr, key, |_, _| ()).is_some()
661 }
662
663 fn traverse<F, R>(&self, node_addr: Address, key: &K, f: F) -> Option<R>
665 where
666 F: Fn(Node<K>, usize) -> R + Clone,
667 {
668 let node = self.load_node(node_addr);
669 match node.search(key, self.memory()) {
671 Ok(idx) => Some(f(node, idx)), Err(idx) => match node.node_type() {
673 NodeType::Leaf => None, NodeType::Internal => self.traverse(node.child(idx), key, f), },
676 }
677 }
678
679 pub fn is_empty(&self) -> bool {
681 self.length == 0
682 }
683
684 pub fn len(&self) -> u64 {
686 self.length
687 }
688
689 pub fn into_memory(self) -> M {
691 self.allocator.into_memory()
692 }
693
694 #[deprecated(since = "0.6.3", note = "please use `clear_new` instead")]
696 pub fn clear(self) -> Self {
697 let mem = self.allocator.into_memory();
698 Self::new(mem)
699 }
700
701 pub fn clear_new(&mut self) {
703 self.root_addr = NULL;
704 self.length = 0;
705 self.allocator.clear();
706 self.save_header();
707 }
708
709 pub fn first_key_value(&self) -> Option<(K, V)> {
712 if self.root_addr == NULL {
713 return None;
714 }
715 let root = self.load_node(self.root_addr);
716 let (k, encoded_v) = root.get_min(self.memory());
717 Some((k, V::from_bytes(Cow::Owned(encoded_v))))
718 }
719
720 pub fn last_key_value(&self) -> Option<(K, V)> {
723 if self.root_addr == NULL {
724 return None;
725 }
726 let root = self.load_node(self.root_addr);
727 let (k, encoded_v) = root.get_max(self.memory());
728 Some((k, V::from_bytes(Cow::Owned(encoded_v))))
729 }
730
731 fn memory(&self) -> &M {
732 self.allocator.memory()
733 }
734
735 fn allocator_mut(&mut self) -> &mut Allocator<M> {
736 &mut self.allocator
737 }
738
739 pub fn remove(&mut self, key: &K) -> Option<V> {
741 if self.root_addr == NULL {
742 return None;
743 }
744
745 let root_node = self.load_node(self.root_addr);
746 self.remove_helper(root_node, key)
747 .map(Cow::Owned)
748 .map(V::from_bytes)
749 }
750
751 pub fn pop_last(&mut self) -> Option<(K, V)> {
753 if self.root_addr == NULL {
754 return None;
755 }
756
757 let root = self.load_node(self.root_addr);
758 let (max_key, _) = root.get_max(self.memory());
759 self.remove_helper(root, &max_key)
760 .map(|v| (max_key, V::from_bytes(Cow::Owned(v))))
761 }
762
763 pub fn pop_first(&mut self) -> Option<(K, V)> {
765 if self.root_addr == NULL {
766 return None;
767 }
768
769 let root = self.load_node(self.root_addr);
770 let (min_key, _) = root.get_min(self.memory());
771 self.remove_helper(root, &min_key)
772 .map(|v| (min_key, V::from_bytes(Cow::Owned(v))))
773 }
774
775 fn remove_helper(&mut self, mut node: Node<K>, key: &K) -> Option<Vec<u8>> {
777 if node.address() != self.root_addr {
778 assert!(node.can_remove_entry_without_merging());
783 }
784
785 match node.node_type() {
786 NodeType::Leaf => {
787 match node.search(key, self.memory()) {
788 Ok(idx) => {
789 let value = node.remove_entry(idx, self.memory()).1;
792 self.length -= 1;
793
794 if node.entries_len() == 0 {
795 assert_eq!(
796 node.address(), self.root_addr,
797 "Removal can only result in an empty leaf node if that node is the root"
798 );
799
800 self.deallocate_node(node);
802 self.root_addr = NULL;
803 } else {
804 self.save_node(&mut node);
805 }
806
807 self.save_header();
808 Some(value)
809 }
810 _ => None, }
812 }
813 NodeType::Internal => {
814 match node.search(key, self.memory()) {
815 Ok(idx) => {
816 let left_child = self.load_node(node.child(idx));
819 if left_child.can_remove_entry_without_merging() {
820 let predecessor = left_child.get_max(self.memory());
843 self.remove_helper(left_child, &predecessor.0)?;
844
845 let (_, old_value) = node.swap_entry(idx, predecessor, self.memory());
847
848 self.save_node(&mut node);
850 return Some(old_value);
851 }
852
853 let right_child = self.load_node(node.child(idx + 1));
854 if right_child.can_remove_entry_without_merging() {
855 let successor = right_child.get_min(self.memory());
878 self.remove_helper(right_child, &successor.0)?;
879
880 let (_, old_value) = node.swap_entry(idx, successor, self.memory());
882
883 self.save_node(&mut node);
885 return Some(old_value);
886 }
887
888 assert!(left_child.at_minimum());
908 assert!(right_child.at_minimum());
909
910 let mut new_child = self.merge(
912 right_child,
913 left_child,
914 node.remove_entry(idx, self.memory()),
915 );
916
917 node.remove_child(idx + 1);
919
920 if node.entries_len() == 0 {
921 assert_eq!(node.address(), self.root_addr);
923 assert_eq!(node.child(0), new_child.address());
924 assert_eq!(node.children_len(), 1);
925
926 self.root_addr = new_child.address();
927
928 self.deallocate_node(node);
930 self.save_header();
931 } else {
932 self.save_node(&mut node);
933 }
934
935 self.save_node(&mut new_child);
936
937 self.remove_helper(new_child, key)
939 }
940 Err(idx) => {
941 let mut child = self.load_node(node.child(idx));
946
947 if child.can_remove_entry_without_merging() {
948 return self.remove_helper(child, key);
951 }
952
953 let mut left_sibling = if idx > 0 {
956 Some(self.load_node(node.child(idx - 1)))
957 } else {
958 None
959 };
960
961 let mut right_sibling = if idx + 1 < node.children_len() {
962 Some(self.load_node(node.child(idx + 1)))
963 } else {
964 None
965 };
966
967 if let Some(ref mut left_sibling) = left_sibling {
968 if left_sibling.can_remove_entry_without_merging() {
969 let (left_sibling_key, left_sibling_value) =
992 left_sibling.pop_entry(self.memory()).unwrap();
993
994 let (parent_key, parent_value) = node.swap_entry(
996 idx - 1,
997 (left_sibling_key, left_sibling_value),
998 self.memory(),
999 );
1000
1001 child.insert_entry(0, (parent_key, parent_value));
1003
1004 if let Some(last_child) = left_sibling.pop_child() {
1006 assert_eq!(left_sibling.node_type(), NodeType::Internal);
1007 assert_eq!(child.node_type(), NodeType::Internal);
1008
1009 child.insert_child(0, last_child);
1010 } else {
1011 assert_eq!(left_sibling.node_type(), NodeType::Leaf);
1012 assert_eq!(child.node_type(), NodeType::Leaf);
1013 }
1014
1015 self.save_node(left_sibling);
1016 self.save_node(&mut child);
1017 self.save_node(&mut node);
1018 return self.remove_helper(child, key);
1019 }
1020 }
1021
1022 if let Some(right_sibling) = &mut right_sibling {
1023 if right_sibling.can_remove_entry_without_merging() {
1024 let (right_sibling_key, right_sibling_value) =
1047 right_sibling.remove_entry(0, self.memory());
1048
1049 let parent_entry = node.swap_entry(
1051 idx,
1052 (right_sibling_key, right_sibling_value),
1053 self.memory(),
1054 );
1055
1056 child.push_entry(parent_entry);
1058
1059 match right_sibling.node_type() {
1061 NodeType::Internal => {
1062 assert_eq!(child.node_type(), NodeType::Internal);
1063 child.push_child(right_sibling.remove_child(0));
1064 }
1065 NodeType::Leaf => {
1066 assert_eq!(child.node_type(), NodeType::Leaf);
1067 }
1068 }
1069
1070 self.save_node(right_sibling);
1071 self.save_node(&mut child);
1072 self.save_node(&mut node);
1073 return self.remove_helper(child, key);
1074 }
1075 }
1076
1077 if let Some(left_sibling) = left_sibling {
1080 assert!(left_sibling.at_minimum());
1083 let left_sibling = self.merge(
1084 child,
1085 left_sibling,
1086 node.remove_entry(idx - 1, self.memory()),
1087 );
1088 node.remove_child(idx);
1090
1091 if node.entries_len() == 0 {
1092 let node_address = node.address();
1093 self.deallocate_node(node);
1094
1095 if node_address == self.root_addr {
1096 self.root_addr = left_sibling.address();
1098 self.save_header();
1099 }
1100 } else {
1101 self.save_node(&mut node);
1102 }
1103
1104 return self.remove_helper(left_sibling, key);
1105 }
1106
1107 if let Some(right_sibling) = right_sibling {
1108 assert!(right_sibling.at_minimum());
1111 let right_sibling = self.merge(
1112 child,
1113 right_sibling,
1114 node.remove_entry(idx, self.memory()),
1115 );
1116
1117 node.remove_child(idx);
1119
1120 if node.entries_len() == 0 {
1121 let node_address = node.address();
1122 self.deallocate_node(node);
1123
1124 if node_address == self.root_addr {
1125 self.root_addr = right_sibling.address();
1127 self.save_header();
1128 }
1129 } else {
1130 self.save_node(&mut node);
1131 }
1132
1133 return self.remove_helper(right_sibling, key);
1134 }
1135
1136 unreachable!("At least one of the siblings must exist.");
1137 }
1138 }
1139 }
1140 }
1141 }
1142
1143 pub fn iter(&self) -> Iter<K, V, M> {
1159 self.iter_internal().into()
1160 }
1161
1162 pub fn range(&self, key_range: impl RangeBounds<K>) -> Iter<K, V, M> {
1182 self.range_internal(key_range).into()
1183 }
1184
1185 pub fn iter_upper_bound(&self, bound: &K) -> Iter<K, V, M> {
1188 if let Some((start_key, _)) = self.range(..bound).next_back() {
1189 IterInternal::new_in_range(self, (Bound::Included(start_key), Bound::Unbounded)).into()
1190 } else {
1191 IterInternal::null(self).into()
1192 }
1193 }
1194
1195 pub fn keys(&self) -> KeysIter<K, V, M> {
1197 self.iter_internal().into()
1198 }
1199
1200 pub fn keys_range(&self, key_range: impl RangeBounds<K>) -> KeysIter<K, V, M> {
1202 self.range_internal(key_range).into()
1203 }
1204
1205 pub fn values(&self) -> ValuesIter<K, V, M> {
1207 self.iter_internal().into()
1208 }
1209
1210 pub fn values_range(&self, key_range: impl RangeBounds<K>) -> ValuesIter<K, V, M> {
1213 self.range_internal(key_range).into()
1214 }
1215
1216 fn iter_internal(&self) -> IterInternal<K, V, M> {
1217 IterInternal::new(self)
1218 }
1219
1220 fn range_internal(&self, key_range: impl RangeBounds<K>) -> IterInternal<K, V, M> {
1221 if self.root_addr == NULL {
1222 return IterInternal::null(self);
1224 }
1225
1226 let range = (
1227 key_range.start_bound().cloned(),
1228 key_range.end_bound().cloned(),
1229 );
1230
1231 IterInternal::new_in_range(self, range)
1232 }
1233
1234 fn merge(&mut self, source: Node<K>, mut into: Node<K>, median: Entry<K>) -> Node<K> {
1247 into.merge(source, median, &mut self.allocator);
1248 into
1249 }
1250
1251 fn allocate_node(&mut self, node_type: NodeType) -> Node<K> {
1253 match self.version {
1254 Version::V1(page_size) => Node::new_v1(self.allocator.allocate(), node_type, page_size),
1255 Version::V2(page_size) => Node::new_v2(self.allocator.allocate(), node_type, page_size),
1256 }
1257 }
1258
1259 #[inline]
1261 fn deallocate_node(&mut self, node: Node<K>) {
1262 node.deallocate(self.allocator_mut());
1263 }
1264
1265 #[inline]
1267 fn load_node(&self, address: Address) -> Node<K> {
1268 Node::load(address, self.version.page_size(), self.memory())
1269 }
1270
1271 #[inline]
1273 fn save_node(&mut self, node: &mut Node<K>) {
1274 node.save(self.allocator_mut());
1275 }
1276
1277 fn update_value(&mut self, node: &mut Node<K>, idx: usize, new_value: Vec<u8>) -> Vec<u8> {
1279 let old_value = node.swap_value(idx, new_value, self.memory());
1280 self.save_node(node);
1281 old_value
1282 }
1283
1284 fn save_header(&self) {
1286 let header = BTreeHeader {
1287 version: self.version,
1288 root_addr: self.root_addr,
1289 length: self.length,
1290 };
1291
1292 Self::write_header(&header, self.memory());
1293 }
1294
1295 fn write_header(header: &BTreeHeader, memory: &M) {
1297 let mut buf = [0; PACKED_HEADER_SIZE];
1299 buf[0..3].copy_from_slice(MAGIC.as_slice());
1300 match header.version {
1301 Version::V1(DerivedPageSize {
1302 max_key_size,
1303 max_value_size,
1304 })
1305 | Version::V2(PageSize::Derived(DerivedPageSize {
1306 max_key_size,
1307 max_value_size,
1308 })) => {
1309 buf[3] = LAYOUT_VERSION;
1310 buf[4..8].copy_from_slice(&max_key_size.to_le_bytes());
1311 buf[8..12].copy_from_slice(&max_value_size.to_le_bytes());
1312 }
1313 Version::V2(PageSize::Value(page_size)) => {
1314 buf[3] = LAYOUT_VERSION_2;
1315 buf[4..8].copy_from_slice(&page_size.to_le_bytes());
1316 buf[8..12].copy_from_slice(&PAGE_SIZE_VALUE_MARKER.to_le_bytes());
1317 }
1318 };
1319 buf[12..20].copy_from_slice(&header.root_addr.get().to_le_bytes());
1320 buf[20..28].copy_from_slice(&header.length.to_le_bytes());
1321 crate::write(memory, 0, &buf);
1323 }
1324}
1325
1326#[cfg(test)]
1327mod test {
1328 use super::*;
1329 use crate::{
1330 storable::{Blob, Bound as StorableBound},
1331 VectorMemory,
1332 };
1333 use std::cell::RefCell;
1334 use std::convert::TryFrom;
1335 use std::rc::Rc;
1336
1337 fn collect_kv<'a, K: Clone + 'a, V: Clone + 'a>(
1339 iter: impl Iterator<Item = (&'a K, &'a V)>,
1340 ) -> Vec<(K, V)> {
1341 iter.map(|(k, v)| (k.clone(), v.clone())).collect()
1342 }
1343
1344 fn collect<K: Clone, V: Clone>(it: impl Iterator<Item = (K, V)>) -> Vec<(K, V)> {
1346 it.collect()
1347 }
1348
1349 fn monotonic_buffer<const N: usize>(i: u32) -> [u8; N] {
1351 let mut buf = [0u8; N];
1352 let bytes = i.to_be_bytes();
1353 buf[N - bytes.len()..].copy_from_slice(&bytes);
1354 buf
1355 }
1356
1357 #[test]
1358 fn test_monotonic_buffer() {
1359 let cases: &[(u32, [u8; 4])] = &[
1360 (1, [0, 0, 0, 1]),
1361 (2, [0, 0, 0, 2]),
1362 ((1 << 8) - 1, [0, 0, 0, 255]),
1363 ((1 << 8), [0, 0, 1, 0]),
1364 ((1 << 16) - 1, [0, 0, 255, 255]),
1365 (1 << 16, [0, 1, 0, 0]),
1366 ((1 << 24) - 1, [0, 255, 255, 255]),
1367 (1 << 24, [1, 0, 0, 0]),
1368 ];
1369
1370 for &(input, expected) in cases {
1371 let output = monotonic_buffer::<4>(input);
1372 assert_eq!(
1373 output, expected,
1374 "monotonic_buffer::<4>({}) returned {:?}, expected {:?}",
1375 input, output, expected
1376 );
1377 }
1378 }
1379
1380 trait Builder {
1382 fn build(i: u32) -> Self;
1383 fn empty() -> Self;
1384 }
1385
1386 impl Builder for () {
1387 fn build(_i: u32) -> Self {}
1388 fn empty() -> Self {}
1389 }
1390
1391 impl Builder for u32 {
1392 fn build(i: u32) -> Self {
1393 i
1394 }
1395 fn empty() -> Self {
1396 0
1397 }
1398 }
1399
1400 impl<const N: usize> Builder for Blob<N> {
1401 fn build(i: u32) -> Self {
1402 Blob::try_from(&monotonic_buffer::<N>(i)[..]).unwrap()
1403 }
1404 fn empty() -> Self {
1405 Blob::try_from(&[][..]).unwrap()
1406 }
1407 }
1408
1409 type MonotonicVec32 = Vec<u8>;
1410 impl Builder for MonotonicVec32 {
1411 fn build(i: u32) -> Self {
1412 monotonic_buffer::<32>(i).to_vec()
1413 }
1414 fn empty() -> Self {
1415 Vec::new()
1416 }
1417 }
1418
1419 type MonotonicString32 = String;
1420 impl Builder for MonotonicString32 {
1421 fn build(i: u32) -> Self {
1422 format!("{i:0>32}")
1423 }
1424 fn empty() -> Self {
1425 String::new()
1426 }
1427 }
1428
1429 fn encode<T: Storable>(object: T) -> Vec<u8> {
1431 object.to_bytes_checked().into_owned()
1432 }
1433
1434 pub(crate) fn b(x: &[u8]) -> Blob<10> {
1436 Blob::<10>::try_from(x).unwrap()
1437 }
1438
1439 pub(crate) fn make_memory() -> Rc<RefCell<Vec<u8>>> {
1441 Rc::new(RefCell::new(Vec::new()))
1442 }
1443
1444 pub fn run_btree_test<K, V, R, F>(f: F)
1446 where
1447 K: Storable + Ord + Clone,
1448 V: Storable,
1449 F: Fn(BTreeMap<K, V, VectorMemory>) -> R,
1450 {
1451 if K::BOUND != StorableBound::Unbounded && V::BOUND != StorableBound::Unbounded {
1453 let mem = make_memory();
1455 let tree_v1 = BTreeMap::new_v1(mem);
1456 f(tree_v1);
1457
1458 let mem = make_memory();
1460 let tree_v1: BTreeMap<K, V, _> = BTreeMap::new_v1(mem);
1461 let tree_v2_migrated = BTreeMap::load_helper(tree_v1.into_memory(), true);
1462 f(tree_v2_migrated);
1463 }
1464
1465 let mem = make_memory();
1467 let tree_v2 = BTreeMap::new(mem);
1468 f(tree_v2);
1469 }
1470
1471 fn verify_monotonic<T: Builder + PartialOrd>() {
1474 for shift_bits in [8, 16, 24] {
1475 let i = (1 << shift_bits) - 1;
1476 assert!(
1477 T::build(i) < T::build(i + 1),
1478 "Monotonicity failed at i: {i}",
1479 );
1480 }
1481 }
1482
1483 macro_rules! assert_bounded {
1485 ($ty:ty) => {
1486 assert_ne!(<$ty>::BOUND, StorableBound::Unbounded, "Must be Bounded");
1487 };
1488 }
1489
1490 macro_rules! assert_unbounded {
1492 ($ty:ty) => {
1493 assert_eq!(<$ty>::BOUND, StorableBound::Unbounded, "Must be Unbounded");
1494 };
1495 }
1496
1497 macro_rules! run_with_key {
1498 ($runner:ident, $K:ty) => {{
1499 verify_monotonic::<$K>();
1500
1501 $runner::<$K, ()>();
1503
1504 assert_bounded!(u32);
1506 $runner::<$K, u32>();
1507
1508 assert_bounded!(Blob<20>);
1509 $runner::<$K, Blob<20>>();
1510
1511 assert_unbounded!(MonotonicVec32);
1513 $runner::<$K, MonotonicVec32>();
1514
1515 assert_unbounded!(MonotonicString32);
1516 $runner::<$K, MonotonicString32>();
1517 }};
1518 }
1519
1520 macro_rules! btree_test {
1522 ($name:ident, $runner:ident) => {
1523 #[test]
1524 fn $name() {
1525 assert_bounded!(u32);
1527 run_with_key!($runner, u32);
1528
1529 assert_bounded!(Blob<10>);
1530 run_with_key!($runner, Blob<10>);
1531
1532 assert_unbounded!(MonotonicVec32);
1534 run_with_key!($runner, MonotonicVec32);
1535
1536 assert_unbounded!(MonotonicString32);
1537 run_with_key!($runner, MonotonicString32);
1538 }
1539 };
1540 }
1541
1542 trait TestKey: Storable + Ord + Clone + Builder + std::fmt::Debug {}
1544 impl<T> TestKey for T where T: Storable + Ord + Clone + Builder + std::fmt::Debug {}
1545
1546 trait TestValue: Storable + Clone + Builder + std::fmt::Debug + PartialEq {}
1548 impl<T> TestValue for T where T: Storable + Clone + Builder + std::fmt::Debug + PartialEq {}
1549
1550 fn insert_get_init_preserves_data<K: TestKey, V: TestValue>() {
1551 let (key, value) = (K::build, V::build);
1552 run_btree_test(|mut btree| {
1553 let n = 1_000;
1554 for i in 0..n {
1555 assert_eq!(btree.insert(key(i), value(i)), None);
1556 assert_eq!(btree.get(&key(i)), Some(value(i)));
1557 }
1558
1559 let btree = BTreeMap::<K, V, VectorMemory>::init(btree.into_memory());
1561 for i in 0..n {
1562 assert_eq!(btree.get(&key(i)), Some(value(i)));
1563 }
1564 });
1565 }
1566 btree_test!(
1567 test_insert_get_init_preserves_data,
1568 insert_get_init_preserves_data
1569 );
1570
1571 fn insert_overwrites_previous_value<K: TestKey, V: TestValue>() {
1572 let (key, value) = (K::build, V::build);
1573 run_btree_test(|mut btree| {
1574 let n = 1_000;
1575 for i in 0..n {
1576 assert_eq!(btree.insert(key(i), value(i)), None);
1577 assert_eq!(btree.insert(key(i), value(i + 1)), Some(value(i)));
1578 assert_eq!(btree.get(&key(i)), Some(value(i + 1)));
1579 }
1580 });
1581 }
1582 btree_test!(
1583 test_insert_overwrites_previous_value,
1584 insert_overwrites_previous_value
1585 );
1586
1587 fn insert_same_key_many<K: TestKey, V: TestValue>() {
1588 let (key, value) = (K::build, V::build);
1589 run_btree_test(|mut btree| {
1590 let n = 1_000;
1591 assert_eq!(btree.insert(key(1), value(2)), None);
1592 for i in 2..n {
1593 assert_eq!(btree.insert(key(1), value(i + 1)), Some(value(i)));
1594 }
1595 assert_eq!(btree.get(&key(1)), Some(value(n)));
1596 });
1597 }
1598 btree_test!(test_insert_same_key_many, insert_same_key_many);
1599
1600 fn insert_overwrite_median_key_in_full_child_node<K: TestKey, V: TestValue>() {
1601 let (key, value) = (K::build, V::build);
1602 run_btree_test(|mut btree| {
1603 for i in 1..=17 {
1604 assert_eq!(btree.insert(key(i), value(0)), None);
1605 }
1606
1607 let root = btree.load_node(btree.root_addr);
1613 assert_eq!(root.node_type(), NodeType::Internal);
1614 assert_eq!(
1615 root.entries(btree.memory()),
1616 vec![(key(6), encode(value(0)))]
1617 );
1618 assert_eq!(root.children_len(), 2);
1619
1620 let right_child = btree.load_node(root.child(1));
1622 assert!(right_child.is_full());
1623 let median_index = right_child.entries_len() / 2;
1624 let median_key = key(12);
1625 assert_eq!(right_child.key(median_index, btree.memory()), &median_key);
1626
1627 assert_eq!(btree.insert(median_key.clone(), value(123)), Some(value(0)));
1629 assert_eq!(btree.get(&median_key), Some(value(123)));
1630
1631 let right_child = btree.load_node(root.child(1));
1633 assert_eq!(right_child.node_type(), NodeType::Leaf);
1634 assert!(right_child.is_full());
1635 });
1636 }
1637 btree_test!(
1638 test_insert_overwrite_median_key_in_full_child_node,
1639 insert_overwrite_median_key_in_full_child_node
1640 );
1641
1642 fn insert_overwrite_key_in_full_root_node<K: TestKey, V: TestValue>() {
1643 let (key, value) = (K::build, V::build);
1644 run_btree_test(|mut btree| {
1645 for i in 1..=11 {
1646 assert_eq!(btree.insert(key(i), value(0)), None);
1647 }
1648
1649 let root = btree.load_node(btree.root_addr);
1653 assert!(root.is_full());
1654
1655 assert_eq!(btree.insert(key(6), value(456)), Some(value(0)));
1657
1658 let root = btree.load_node(btree.root_addr);
1659 assert_eq!(root.node_type(), NodeType::Leaf);
1660 assert_eq!(btree.get(&key(6)), Some(value(456)));
1661 assert_eq!(root.entries_len(), 11);
1662 });
1663 }
1664 btree_test!(
1665 test_insert_overwrite_key_in_full_root_node,
1666 insert_overwrite_key_in_full_root_node
1667 );
1668
1669 fn allocations_without_split<K: TestKey, V: TestValue>() {
1670 let key = K::build;
1671 run_btree_test(|mut btree| {
1672 assert_eq!(btree.allocator.num_allocated_chunks(), 0);
1673
1674 assert_eq!(btree.insert(key(1), V::empty()), None);
1675 assert_eq!(btree.allocator.num_allocated_chunks(), 1);
1676
1677 assert_eq!(btree.remove(&key(1)), Some(V::empty()));
1678 assert_eq!(btree.allocator.num_allocated_chunks(), 0);
1679 });
1680 }
1681 btree_test!(test_allocations_without_split, allocations_without_split);
1682
1683 fn allocations_with_split<K: TestKey, V: TestValue>() {
1684 let key = K::build;
1685 run_btree_test(|mut btree| {
1686 let mut i = 0;
1688 loop {
1689 assert_eq!(btree.insert(key(i), V::empty()), None);
1690 let root = btree.load_node(btree.root_addr);
1691 if root.is_full() {
1692 break;
1693 }
1694 i += 1;
1695 }
1696
1697 assert_eq!(btree.allocator.num_allocated_chunks(), 1);
1699
1700 assert_eq!(btree.insert(key(255), V::empty()), None);
1701
1702 assert_eq!(btree.allocator.num_allocated_chunks(), 3);
1704 });
1705 }
1706 btree_test!(test_allocations_with_split, allocations_with_split);
1707
1708 fn insert_split_node<K: TestKey, V: TestValue>() {
1709 let (key, value) = (K::build, V::build);
1710 run_btree_test(|mut btree| {
1711 for i in 1..=11 {
1712 assert_eq!(btree.insert(key(i), value(i)), None);
1713 }
1714
1715 let root = btree.load_node(btree.root_addr);
1717 assert!(root.is_full());
1718 assert_eq!(btree.insert(key(12), value(12)), None);
1719
1720 for i in 1..=12 {
1726 assert_eq!(btree.get(&key(i)), Some(value(i)));
1727 }
1728 });
1729 }
1730 btree_test!(test_insert_split_node, insert_split_node);
1731
1732 fn insert_split_multiple_nodes<K: TestKey, V: TestValue>() {
1733 let key = K::build;
1734 let e = |i: u32| (key(i), encode(V::empty()));
1735 run_btree_test(|mut btree| {
1736 for i in 1..=11 {
1737 assert_eq!(btree.insert(key(i), V::empty()), None);
1738 }
1739 assert_eq!(btree.insert(key(12), V::empty()), None);
1741
1742 let root = btree.load_node(btree.root_addr);
1748 assert_eq!(root.node_type(), NodeType::Internal);
1749 assert_eq!(root.entries(btree.memory()), vec![e(6)]);
1750 assert_eq!(root.children_len(), 2);
1751
1752 let child_0 = btree.load_node(root.child(0));
1753 assert_eq!(child_0.node_type(), NodeType::Leaf);
1754 assert_eq!(
1755 child_0.entries(btree.memory()),
1756 vec![e(1), e(2), e(3), e(4), e(5)]
1757 );
1758
1759 let child_1 = btree.load_node(root.child(1));
1760 assert_eq!(child_1.node_type(), NodeType::Leaf);
1761 assert_eq!(
1762 child_1.entries(btree.memory()),
1763 vec![e(7), e(8), e(9), e(10), e(11), e(12)]
1764 );
1765
1766 for i in 1..=12 {
1767 assert_eq!(btree.get(&key(i)), Some(V::empty()));
1768 }
1769
1770 assert_eq!(btree.insert(key(13), V::empty()), None);
1772 assert_eq!(btree.insert(key(14), V::empty()), None);
1773 assert_eq!(btree.insert(key(15), V::empty()), None);
1774 assert_eq!(btree.insert(key(16), V::empty()), None);
1775 assert_eq!(btree.insert(key(17), V::empty()), None);
1776 assert_eq!(btree.insert(key(18), V::empty()), None);
1778
1779 for i in 1..=18 {
1780 assert_eq!(btree.get(&key(i)), Some(V::empty()));
1781 }
1782
1783 let root = btree.load_node(btree.root_addr);
1784 assert_eq!(root.node_type(), NodeType::Internal);
1785 assert_eq!(root.entries(btree.memory()), vec![e(6), e(12)],);
1786 assert_eq!(root.children_len(), 3);
1787
1788 let child_0 = btree.load_node(root.child(0));
1789 assert_eq!(child_0.node_type(), NodeType::Leaf);
1790 assert_eq!(
1791 child_0.entries(btree.memory()),
1792 vec![e(1), e(2), e(3), e(4), e(5)]
1793 );
1794
1795 let child_1 = btree.load_node(root.child(1));
1796 assert_eq!(child_1.node_type(), NodeType::Leaf);
1797 assert_eq!(
1798 child_1.entries(btree.memory()),
1799 vec![e(7), e(8), e(9), e(10), e(11)]
1800 );
1801
1802 let child_2 = btree.load_node(root.child(2));
1803 assert_eq!(child_2.node_type(), NodeType::Leaf);
1804 assert_eq!(
1805 child_2.entries(btree.memory()),
1806 vec![e(13), e(14), e(15), e(16), e(17), e(18)]
1807 );
1808
1809 assert_eq!(btree.allocator.num_allocated_chunks(), 4);
1810 });
1811 }
1812 btree_test!(
1813 test_insert_split_multiple_nodes,
1814 insert_split_multiple_nodes
1815 );
1816
1817 fn first_key_value<K: TestKey, V: TestValue>() {
1818 let (key, value) = (K::build, V::build);
1819 run_btree_test(|mut btree: BTreeMap<K, V, _>| {
1820 assert_eq!(btree.first_key_value(), None);
1821
1822 let n = 1_000;
1823 for i in (0..n).rev() {
1824 assert_eq!(btree.insert(key(i), value(i)), None);
1825 assert_eq!(btree.first_key_value(), Some((key(i), value(i))));
1826 }
1827
1828 for i in 0..n {
1829 assert_eq!(btree.remove(&key(i)), Some(value(i)));
1830 if !btree.is_empty() {
1831 assert_eq!(btree.first_key_value(), Some((key(i + 1), value(i + 1))));
1832 }
1833 }
1834 assert_eq!(btree.first_key_value(), None);
1835 });
1836 }
1837 btree_test!(test_first_key_value, first_key_value);
1838
1839 fn last_key_value<K: TestKey, V: TestValue>() {
1840 let (key, value) = (K::build, V::build);
1841 run_btree_test(|mut btree: BTreeMap<K, V, _>| {
1842 assert_eq!(btree.last_key_value(), None);
1843
1844 let n = 1_000;
1845 for i in 0..n {
1846 assert_eq!(btree.insert(key(i), value(i)), None);
1847 assert_eq!(btree.last_key_value(), Some((key(i), value(i))));
1848 }
1849
1850 for i in (0..n).rev() {
1851 assert_eq!(btree.remove(&key(i)), Some(value(i)));
1852 if !btree.is_empty() {
1853 assert_eq!(btree.last_key_value(), Some((key(i - 1), value(i - 1))));
1854 }
1855 }
1856 assert_eq!(btree.last_key_value(), None);
1857 });
1858 }
1859 btree_test!(test_last_key_value, last_key_value);
1860
1861 fn pop_first_single_entry<K: TestKey, V: TestValue>() {
1862 let key = K::build;
1863 run_btree_test(|mut btree| {
1864 assert_eq!(btree.allocator.num_allocated_chunks(), 0);
1865
1866 assert_eq!(btree.insert(key(1), V::empty()), None);
1867 assert!(!btree.is_empty());
1868 assert_eq!(btree.allocator.num_allocated_chunks(), 1);
1869
1870 assert_eq!(btree.pop_first(), Some((key(1), V::empty())));
1871 assert!(btree.is_empty());
1872 assert_eq!(btree.allocator.num_allocated_chunks(), 0);
1873 });
1874 }
1875 btree_test!(test_pop_first_single_entry, pop_first_single_entry);
1876
1877 fn pop_last_single_entry<K: TestKey, V: TestValue>() {
1878 let (key, value) = (K::build, V::build);
1879 run_btree_test(|mut btree| {
1880 assert_eq!(btree.allocator.num_allocated_chunks(), 0);
1881
1882 assert_eq!(btree.insert(key(1), value(1)), None);
1883 assert!(!btree.is_empty());
1884 assert_eq!(btree.allocator.num_allocated_chunks(), 1);
1885
1886 assert_eq!(btree.pop_last(), Some((key(1), value(1))));
1887 assert!(btree.is_empty());
1888 assert_eq!(btree.allocator.num_allocated_chunks(), 0);
1889 });
1890 }
1891 btree_test!(test_pop_last_single_entry, pop_last_single_entry);
1892
1893 fn remove_case_2a_and_2c<K: TestKey, V: TestValue>() {
1894 let key = K::build;
1895 let e = |i: u32| (key(i), encode(V::empty()));
1896 run_btree_test(|mut btree| {
1897 for i in 1..=11 {
1898 assert_eq!(btree.insert(key(i), V::empty()), None);
1899 }
1900 assert_eq!(btree.insert(key(0), V::empty()), None);
1902
1903 for i in 0..=11 {
1909 assert_eq!(btree.get(&key(i)), Some(V::empty()));
1910 }
1911
1912 assert_eq!(btree.remove(&key(6)), Some(V::empty()));
1914
1915 let root = btree.load_node(btree.root_addr);
1920 assert_eq!(root.node_type(), NodeType::Internal);
1921 assert_eq!(root.entries(btree.memory()), vec![e(5)]);
1922 assert_eq!(root.children_len(), 2);
1923
1924 let child_0 = btree.load_node(root.child(0));
1925 assert_eq!(child_0.node_type(), NodeType::Leaf);
1926 assert_eq!(
1927 child_0.entries(btree.memory()),
1928 vec![e(0), e(1), e(2), e(3), e(4)]
1929 );
1930
1931 let child_1 = btree.load_node(root.child(1));
1932 assert_eq!(child_1.node_type(), NodeType::Leaf);
1933 assert_eq!(
1934 child_1.entries(btree.memory()),
1935 vec![e(7), e(8), e(9), e(10), e(11)]
1936 );
1937
1938 assert_eq!(btree.allocator.num_allocated_chunks(), 3);
1940
1941 assert_eq!(btree.remove(&key(5)), Some(V::empty()));
1943
1944 let btree = BTreeMap::<K, V, _>::load(btree.into_memory());
1946
1947 let root = btree.load_node(btree.root_addr);
1950 assert_eq!(
1951 root.entries(btree.memory()),
1952 vec![e(0), e(1), e(2), e(3), e(4), e(7), e(8), e(9), e(10), e(11)]
1953 );
1954
1955 assert_eq!(btree.allocator.num_allocated_chunks(), 1);
1957 });
1958 }
1959 btree_test!(test_remove_case_2a_and_2c, remove_case_2a_and_2c);
1960
1961 fn remove_case_2b<K: TestKey, V: TestValue>() {
1962 let key = K::build;
1963 let e = |i: u32| (key(i), encode(V::empty()));
1964 run_btree_test(|mut btree| {
1965 for i in 1..=11 {
1966 assert_eq!(btree.insert(key(i), V::empty()), None);
1967 }
1968 assert_eq!(btree.insert(key(12), V::empty()), None);
1970
1971 for i in 1..=12 {
1977 assert_eq!(btree.get(&key(i)), Some(V::empty()));
1978 }
1979
1980 assert_eq!(btree.remove(&key(6)), Some(V::empty()));
1982
1983 let root = btree.load_node(btree.root_addr);
1988 assert_eq!(root.node_type(), NodeType::Internal);
1989 assert_eq!(root.entries(btree.memory()), vec![e(7)]);
1990 assert_eq!(root.children_len(), 2);
1991
1992 let child_0 = btree.load_node(root.child(0));
1993 assert_eq!(child_0.node_type(), NodeType::Leaf);
1994 assert_eq!(
1995 child_0.entries(btree.memory()),
1996 vec![e(1), e(2), e(3), e(4), e(5)]
1997 );
1998
1999 let child_1 = btree.load_node(root.child(1));
2000 assert_eq!(child_1.node_type(), NodeType::Leaf);
2001 assert_eq!(
2002 child_1.entries(btree.memory()),
2003 vec![e(8), e(9), e(10), e(11), e(12)]
2004 );
2005
2006 assert_eq!(btree.remove(&key(7)), Some(V::empty()));
2008 let root = btree.load_node(btree.root_addr);
2012 assert_eq!(root.node_type(), NodeType::Leaf);
2013 assert_eq!(
2014 root.entries(btree.memory()),
2015 collect([1, 2, 3, 4, 5, 8, 9, 10, 11, 12].into_iter().map(e))
2016 );
2017
2018 assert_eq!(btree.allocator.num_allocated_chunks(), 1);
2019 });
2020 }
2021 btree_test!(test_remove_case_2b, remove_case_2b);
2022
2023 fn remove_case_3a_right<K: TestKey, V: TestValue>() {
2024 let key = K::build;
2025 let e = |i: u32| (key(i), encode(V::empty()));
2026 run_btree_test(|mut btree| {
2027 for i in 1..=11 {
2028 assert_eq!(btree.insert(key(i), V::empty()), None);
2029 }
2030
2031 assert_eq!(btree.insert(key(12), V::empty()), None);
2033
2034 assert_eq!(btree.remove(&key(3)), Some(V::empty()));
2041
2042 let root = btree.load_node(btree.root_addr);
2047 assert_eq!(root.node_type(), NodeType::Internal);
2048 assert_eq!(root.entries(btree.memory()), vec![e(7)]);
2049 assert_eq!(root.children_len(), 2);
2050
2051 let child_0 = btree.load_node(root.child(0));
2052 assert_eq!(child_0.node_type(), NodeType::Leaf);
2053 assert_eq!(
2054 child_0.entries(btree.memory()),
2055 vec![e(1), e(2), e(4), e(5), e(6)]
2056 );
2057
2058 let child_1 = btree.load_node(root.child(1));
2059 assert_eq!(child_1.node_type(), NodeType::Leaf);
2060 assert_eq!(
2061 child_1.entries(btree.memory()),
2062 vec![e(8), e(9), e(10), e(11), e(12)]
2063 );
2064
2065 assert_eq!(btree.allocator.num_allocated_chunks(), 3);
2067 });
2068 }
2069 btree_test!(test_remove_case_3a_right, remove_case_3a_right);
2070
2071 fn remove_case_3a_left<K: TestKey, V: TestValue>() {
2072 let key = K::build;
2073 let e = |i: u32| (key(i), encode(V::empty()));
2074 run_btree_test(|mut btree| {
2075 for i in 1..=11 {
2076 assert_eq!(btree.insert(key(i), V::empty()), None);
2077 }
2078 assert_eq!(btree.insert(key(0), V::empty()), None);
2080
2081 assert_eq!(btree.remove(&key(8)), Some(V::empty()));
2088
2089 let root = btree.load_node(btree.root_addr);
2094 assert_eq!(root.node_type(), NodeType::Internal);
2095 assert_eq!(root.entries(btree.memory()), vec![e(5)]);
2096 assert_eq!(root.children_len(), 2);
2097
2098 let child_0 = btree.load_node(root.child(0));
2099 assert_eq!(child_0.node_type(), NodeType::Leaf);
2100 assert_eq!(
2101 child_0.entries(btree.memory()),
2102 vec![e(0), e(1), e(2), e(3), e(4)]
2103 );
2104
2105 let child_1 = btree.load_node(root.child(1));
2106 assert_eq!(child_1.node_type(), NodeType::Leaf);
2107 assert_eq!(
2108 child_1.entries(btree.memory()),
2109 vec![e(6), e(7), e(9), e(10), e(11)]
2110 );
2111
2112 assert_eq!(btree.allocator.num_allocated_chunks(), 3);
2114 });
2115 }
2116 btree_test!(test_remove_case_3a_left, remove_case_3a_left);
2117
2118 fn remove_case_3b_merge_into_right<K: TestKey, V: TestValue>() {
2119 let key = K::build;
2120 let e = |i: u32| (key(i), encode(V::empty()));
2121 run_btree_test(|mut btree| {
2122 for i in 1..=11 {
2123 assert_eq!(btree.insert(key(i), V::empty()), None);
2124 }
2125 assert_eq!(btree.insert(key(12), V::empty()), None);
2127
2128 for i in 1..=12 {
2134 assert_eq!(btree.get(&key(i)), Some(V::empty()));
2135 }
2136
2137 assert_eq!(btree.remove(&key(6)), Some(V::empty()));
2139 let root = btree.load_node(btree.root_addr);
2144 assert_eq!(root.node_type(), NodeType::Internal);
2145 assert_eq!(root.entries(btree.memory()), vec![e(7)]);
2146 assert_eq!(root.children_len(), 2);
2147
2148 let child_0 = btree.load_node(root.child(0));
2149 assert_eq!(child_0.node_type(), NodeType::Leaf);
2150 assert_eq!(
2151 child_0.entries(btree.memory()),
2152 vec![e(1), e(2), e(3), e(4), e(5)]
2153 );
2154
2155 let child_1 = btree.load_node(root.child(1));
2156 assert_eq!(child_1.node_type(), NodeType::Leaf);
2157 assert_eq!(
2158 child_1.entries(btree.memory()),
2159 vec![e(8), e(9), e(10), e(11), e(12)]
2160 );
2161
2162 assert_eq!(btree.allocator.num_allocated_chunks(), 3);
2164
2165 assert_eq!(btree.remove(&key(3)), Some(V::empty()));
2167
2168 let btree = BTreeMap::<K, V, _>::load(btree.into_memory());
2170
2171 let root = btree.load_node(btree.root_addr);
2175 assert_eq!(root.node_type(), NodeType::Leaf);
2176 assert_eq!(
2177 root.entries(btree.memory()),
2178 collect([1, 2, 4, 5, 7, 8, 9, 10, 11, 12].into_iter().map(e))
2179 );
2180
2181 assert_eq!(btree.allocator.num_allocated_chunks(), 1);
2183 });
2184 }
2185 btree_test!(
2186 test_remove_case_3b_merge_into_right,
2187 remove_case_3b_merge_into_right
2188 );
2189
2190 fn remove_case_3b_merge_into_left<K: TestKey, V: TestValue>() {
2191 let key = K::build;
2192 let e = |i: u32| (key(i), encode(V::empty()));
2193 run_btree_test(|mut btree| {
2194 for i in 1..=11 {
2195 assert_eq!(btree.insert(key(i), V::empty()), None);
2196 }
2197
2198 assert_eq!(btree.insert(key(12), V::empty()), None);
2200
2201 for i in 1..=12 {
2207 assert_eq!(btree.get(&key(i)), Some(V::empty()));
2208 }
2209
2210 assert_eq!(btree.remove(&key(6)), Some(V::empty()));
2212
2213 let root = btree.load_node(btree.root_addr);
2218 assert_eq!(root.node_type(), NodeType::Internal);
2219 assert_eq!(root.entries(btree.memory()), vec![e(7)]);
2220 assert_eq!(root.children_len(), 2);
2221
2222 let child_0 = btree.load_node(root.child(0));
2223 assert_eq!(child_0.node_type(), NodeType::Leaf);
2224 assert_eq!(
2225 child_0.entries(btree.memory()),
2226 vec![e(1), e(2), e(3), e(4), e(5)]
2227 );
2228
2229 let child_1 = btree.load_node(root.child(1));
2230 assert_eq!(child_1.node_type(), NodeType::Leaf);
2231 assert_eq!(
2232 child_1.entries(btree.memory()),
2233 vec![e(8), e(9), e(10), e(11), e(12)]
2234 );
2235
2236 assert_eq!(btree.allocator.num_allocated_chunks(), 3);
2238
2239 assert_eq!(btree.remove(&key(10)), Some(V::empty()));
2241
2242 let btree = BTreeMap::<K, V, _>::load(btree.into_memory());
2244
2245 let root = btree.load_node(btree.root_addr);
2249 assert_eq!(root.node_type(), NodeType::Leaf);
2250 assert_eq!(
2251 root.entries(btree.memory()),
2252 vec![e(1), e(2), e(3), e(4), e(5), e(7), e(8), e(9), e(11), e(12)]
2253 );
2254
2255 assert_eq!(btree.allocator.num_allocated_chunks(), 1);
2257 });
2258 }
2259 btree_test!(
2260 test_remove_case_3b_merge_into_left,
2261 remove_case_3b_merge_into_left
2262 );
2263
2264 fn insert_remove_many<K: TestKey, V: TestValue>() {
2265 let (key, value) = (K::build, V::build);
2266 run_btree_test(|mut btree| {
2267 let n = 10_000;
2268 for i in 0..n {
2269 assert_eq!(btree.insert(key(i), value(i)), None);
2270 assert_eq!(btree.get(&key(i)), Some(value(i)));
2271 }
2272
2273 let mut btree = BTreeMap::<K, V, _>::load(btree.into_memory());
2274 for i in 0..n {
2275 assert_eq!(btree.remove(&key(i)), Some(value(i)));
2276 assert_eq!(btree.get(&key(i)), None);
2277 }
2278
2279 assert_eq!(btree.allocator.num_allocated_chunks(), 0);
2281 });
2282 }
2283 btree_test!(test_insert_remove_many, insert_remove_many);
2284
2285 fn pop_first_many<K: TestKey, V: TestValue>() {
2286 let (key, value) = (K::build, V::build);
2287 run_btree_test(|mut btree| {
2288 let n = 10_000;
2289
2290 assert!(btree.is_empty());
2291 assert_eq!(btree.pop_first(), None);
2292
2293 for i in 0..n {
2294 assert_eq!(btree.insert(key(i), value(i)), None);
2295 }
2296 assert_eq!(btree.len(), n as u64);
2297
2298 let mut btree = BTreeMap::<K, V, _>::load(btree.into_memory());
2299 for i in 0..n {
2300 assert_eq!(btree.pop_first(), Some((key(i), value(i))));
2301 assert_eq!(btree.get(&key(i)), None);
2302 }
2303
2304 assert!(btree.is_empty());
2306 assert_eq!(btree.allocator.num_allocated_chunks(), 0);
2307 });
2308 }
2309 btree_test!(test_pop_first_many, pop_first_many);
2310
2311 fn pop_last_many<K: TestKey, V: TestValue>() {
2312 let (key, value) = (K::build, V::build);
2313 run_btree_test(|mut btree| {
2314 let n = 10_000;
2315
2316 assert!(btree.is_empty());
2317 assert_eq!(btree.pop_last(), None);
2318
2319 for i in 0..n {
2320 assert_eq!(btree.insert(key(i), value(i)), None);
2321 }
2322 assert_eq!(btree.len(), n as u64);
2323
2324 let mut btree = BTreeMap::<K, V, _>::load(btree.into_memory());
2325 for i in (0..n).rev() {
2326 assert_eq!(btree.pop_last(), Some((key(i), value(i))));
2327 assert_eq!(btree.get(&key(i)), None);
2328 }
2329
2330 assert!(btree.is_empty());
2332 assert_eq!(btree.allocator.num_allocated_chunks(), 0);
2333 });
2334 }
2335 btree_test!(test_pop_last_many, pop_last_many);
2336
2337 fn reloading<K: TestKey, V: TestValue>() {
2338 let (key, value) = (K::build, V::build);
2339 run_btree_test(|mut btree| {
2340 let n = 1_000;
2341 for i in 0..n {
2342 assert_eq!(btree.insert(key(i), value(i)), None);
2343 assert_eq!(btree.get(&key(i)), Some(value(i)));
2344 }
2345 assert_eq!(btree.len(), n as u64);
2346
2347 let mut btree = BTreeMap::<K, V, VectorMemory>::load(btree.into_memory());
2349 assert_eq!(btree.len(), n as u64);
2350 for i in 0..n {
2351 assert_eq!(btree.get(&key(i)), Some(value(i)));
2352 }
2353
2354 for i in 0..n {
2356 assert_eq!(btree.remove(&key(i)), Some(value(i)));
2357 }
2358 assert_eq!(btree.len(), 0);
2359
2360 let btree = BTreeMap::<K, V, VectorMemory>::load(btree.into_memory());
2362 assert_eq!(btree.len(), 0);
2363 });
2364 }
2365 btree_test!(test_reloading, reloading);
2366
2367 fn len<K: TestKey, V: TestValue>() {
2368 let (key, value) = (K::build, V::build);
2369 run_btree_test(|mut btree| {
2370 let n = 1_000;
2371 for i in 0..n {
2372 assert_eq!(btree.insert(key(i), value(i)), None);
2373 }
2374
2375 assert_eq!(btree.len(), n as u64);
2376 assert!(!btree.is_empty());
2377
2378 for i in 0..n {
2379 assert_eq!(btree.remove(&key(i)), Some(value(i)));
2380 }
2381
2382 assert_eq!(btree.len(), 0);
2383 assert!(btree.is_empty());
2384 });
2385 }
2386 btree_test!(test_len, len);
2387
2388 fn contains_key<K: TestKey, V: TestValue>() {
2389 let (key, value) = (K::build, V::build);
2390 run_btree_test(|mut btree| {
2391 let n = 1_000;
2392 for i in (0..n).step_by(2) {
2393 assert_eq!(btree.insert(key(i), value(i)), None);
2394 }
2395
2396 for i in 0..n {
2398 assert_eq!(btree.contains_key(&key(i)), i % 2 == 0);
2399 }
2400 });
2401 }
2402 btree_test!(test_contains_key, contains_key);
2403
2404 fn range_empty<K: TestKey, V: TestValue>() {
2405 let key = K::build;
2406 run_btree_test(|btree: BTreeMap<K, V, _>| {
2407 assert_eq!(collect(btree.range(key(0)..)), vec![]);
2409 assert_eq!(collect(btree.range(key(10)..)), vec![]);
2410 assert_eq!(collect(btree.range(key(100)..)), vec![]);
2411 });
2412 }
2413 btree_test!(test_range_empty, range_empty);
2414
2415 fn range_leaf_prefix_greater_than_all_entries<K: TestKey, V: TestValue>() {
2417 let (key, value) = (K::build, V::build);
2418 run_btree_test(|mut btree| {
2419 btree.insert(key(0), value(0));
2420
2421 assert_eq!(collect(btree.range(key(1)..)), vec![]);
2423 });
2424 }
2425 btree_test!(
2426 test_range_leaf_prefix_greater_than_all_entries,
2427 range_leaf_prefix_greater_than_all_entries
2428 );
2429
2430 fn range_internal_prefix_greater_than_all_entries<K: TestKey, V: TestValue>() {
2432 let (key, value) = (K::build, V::build);
2433 run_btree_test(|mut btree| {
2434 for i in 1..=12 {
2435 assert_eq!(btree.insert(key(i), value(i)), None);
2436 }
2437
2438 assert_eq!(
2445 collect(btree.range(key(7)..key(8))),
2446 vec![(key(7), value(7))]
2447 );
2448 });
2449 }
2450 btree_test!(
2451 test_range_internal_prefix_greater_than_all_entries,
2452 range_internal_prefix_greater_than_all_entries
2453 );
2454
2455 fn range_various_prefixes<K: TestKey, V: TestValue>() {
2456 let (key, value) = (K::build, V::build);
2457 run_btree_test(|mut btree| {
2458 btree.insert(key(1), value(100));
2459 btree.insert(key(2), value(200));
2460 btree.insert(key(3), value(300));
2461 btree.insert(key(4), value(400));
2462 btree.insert(key(11), value(500));
2463 btree.insert(key(12), value(600));
2464 btree.insert(key(13), value(700));
2465 btree.insert(key(14), value(800));
2466 btree.insert(key(21), value(900));
2467 btree.insert(key(22), value(1_000));
2468 btree.insert(key(23), value(1_100));
2469 btree.insert(key(24), value(1_200));
2470
2471 let root = btree.load_node(btree.root_addr);
2477 assert_eq!(root.node_type(), NodeType::Internal);
2478 assert_eq!(
2479 root.entries(btree.memory()),
2480 vec![(key(12), encode(value(600)))]
2481 );
2482 assert_eq!(root.children_len(), 2);
2483
2484 assert_eq!(
2486 collect(btree.range(key(0)..key(11))),
2487 vec![
2488 (key(1), value(100)),
2489 (key(2), value(200)),
2490 (key(3), value(300)),
2491 (key(4), value(400)),
2492 ]
2493 );
2494
2495 assert_eq!(
2497 collect(btree.range(key(10)..key(20))),
2498 vec![
2499 (key(11), value(500)),
2500 (key(12), value(600)),
2501 (key(13), value(700)),
2502 (key(14), value(800)),
2503 ]
2504 );
2505
2506 assert_eq!(
2508 collect(btree.range(key(20)..key(30))),
2509 vec![
2510 (key(21), value(900)),
2511 (key(22), value(1_000)),
2512 (key(23), value(1_100)),
2513 (key(24), value(1_200)),
2514 ]
2515 );
2516 });
2517 }
2518 btree_test!(test_range_various_prefixes, range_various_prefixes);
2519
2520 fn range_various_prefixes_2<K: TestKey, V: TestValue>() {
2521 let (key, value) = (K::build, V::build);
2522 run_btree_test(|mut btree| {
2523 btree.insert(key(1), value(100));
2524 btree.insert(key(2), value(200));
2525 btree.insert(key(3), value(300));
2526 btree.insert(key(4), value(400));
2527 btree.insert(key(12), value(500));
2528 btree.insert(key(14), value(600));
2529 btree.insert(key(16), value(700));
2530 btree.insert(key(18), value(800));
2531 btree.insert(key(19), value(900));
2532 btree.insert(key(21), value(1000));
2533 btree.insert(key(22), value(1100));
2534 btree.insert(key(23), value(1200));
2535 btree.insert(key(24), value(1300));
2536 btree.insert(key(25), value(1400));
2537 btree.insert(key(26), value(1500));
2538 btree.insert(key(27), value(1600));
2539 btree.insert(key(28), value(1700));
2540 btree.insert(key(29), value(1800));
2541
2542 let root = btree.load_node(btree.root_addr);
2549 assert_eq!(root.node_type(), NodeType::Internal);
2550 assert_eq!(
2551 root.entries(btree.memory()),
2552 vec![
2553 (key(14), encode(value(600))),
2554 (key(23), encode(value(1200)))
2555 ]
2556 );
2557 assert_eq!(root.children_len(), 3);
2558 let child_0 = btree.load_node(root.child(0));
2559 assert_eq!(child_0.node_type(), NodeType::Leaf);
2560 assert_eq!(
2561 child_0.entries(btree.memory()),
2562 vec![
2563 (key(1), encode(value(100))),
2564 (key(2), encode(value(200))),
2565 (key(3), encode(value(300))),
2566 (key(4), encode(value(400))),
2567 (key(12), encode(value(500))),
2568 ]
2569 );
2570
2571 let child_1 = btree.load_node(root.child(1));
2572 assert_eq!(child_1.node_type(), NodeType::Leaf);
2573 assert_eq!(
2574 child_1.entries(btree.memory()),
2575 vec![
2576 (key(16), encode(value(700))),
2577 (key(18), encode(value(800))),
2578 (key(19), encode(value(900))),
2579 (key(21), encode(value(1000))),
2580 (key(22), encode(value(1100))),
2581 ]
2582 );
2583
2584 let child_2 = btree.load_node(root.child(2));
2585 assert_eq!(
2586 child_2.entries(btree.memory()),
2587 vec![
2588 (key(24), encode(value(1300))),
2589 (key(25), encode(value(1400))),
2590 (key(26), encode(value(1500))),
2591 (key(27), encode(value(1600))),
2592 (key(28), encode(value(1700))),
2593 (key(29), encode(value(1800))),
2594 ]
2595 );
2596
2597 assert_eq!(collect(btree.range(key(15)..key(16))), vec![]);
2599
2600 assert_eq!(
2602 collect(btree.range(key(15)..=key(26))),
2603 vec![
2604 (key(16), value(700)),
2605 (key(18), value(800)),
2606 (key(19), value(900)),
2607 (key(21), value(1000)),
2608 (key(22), value(1100)),
2609 (key(23), value(1200)),
2610 (key(24), value(1300)),
2611 (key(25), value(1400)),
2612 (key(26), value(1500)),
2613 ]
2614 );
2615
2616 assert_eq!(
2618 collect(btree.range(key(10)..key(20))),
2619 vec![
2620 (key(12), value(500)),
2621 (key(14), value(600)),
2622 (key(16), value(700)),
2623 (key(18), value(800)),
2624 (key(19), value(900)),
2625 ]
2626 );
2627
2628 assert_eq!(
2631 collect(btree.range(key(20)..key(30))),
2632 vec![
2633 (key(21), value(1000)),
2634 (key(22), value(1100)),
2635 (key(23), value(1200)),
2636 (key(24), value(1300)),
2637 (key(25), value(1400)),
2638 (key(26), value(1500)),
2639 (key(27), value(1600)),
2640 (key(28), value(1700)),
2641 (key(29), value(1800)),
2642 ]
2643 );
2644 });
2645 }
2646 btree_test!(test_range_various_prefixes_2, range_various_prefixes_2);
2647
2648 fn range_large<K: TestKey, V: TestValue>() {
2649 let (key, value) = (K::build, V::build);
2650 run_btree_test(|mut btree| {
2651 const TOTAL: u32 = 2_000;
2652 const MID: u32 = TOTAL / 2;
2653
2654 for i in 0..TOTAL {
2655 assert_eq!(btree.insert(key(i), value(i)), None);
2656 }
2657
2658 let keys: Vec<_> = btree.range(key(0)..key(MID)).collect();
2660 assert_eq!(keys.len(), MID as usize);
2661 for (i, (k, v)) in keys.into_iter().enumerate() {
2662 let j = i as u32;
2663 assert_eq!(k, key(j));
2664 assert_eq!(v, value(j));
2665 }
2666
2667 let keys: Vec<_> = btree.range(key(MID)..).collect();
2669 assert_eq!(keys.len(), (TOTAL - MID) as usize);
2670 for (i, (k, v)) in keys.into_iter().enumerate() {
2671 let j = MID + i as u32;
2672 assert_eq!(k, key(j));
2673 assert_eq!(v, value(j));
2674 }
2675 });
2676 }
2677 btree_test!(test_range_large, range_large);
2678
2679 fn range_various_prefixes_with_offset<K: TestKey, V: TestValue>() {
2680 let (key, value) = (K::build, V::build);
2681 run_btree_test(|mut btree| {
2682 btree.insert(key(1), value(100));
2683 btree.insert(key(2), value(200));
2684 btree.insert(key(3), value(300));
2685 btree.insert(key(4), value(400));
2686 btree.insert(key(11), value(500));
2687 btree.insert(key(12), value(600));
2688 btree.insert(key(13), value(700));
2689 btree.insert(key(14), value(800));
2690 btree.insert(key(21), value(900));
2691 btree.insert(key(22), value(1000));
2692 btree.insert(key(23), value(1100));
2693 btree.insert(key(24), value(1200));
2694
2695 let root = btree.load_node(btree.root_addr);
2701 assert_eq!(root.node_type(), NodeType::Internal);
2702 assert_eq!(
2703 root.entries(btree.memory()),
2704 vec![(key(12), encode(value(600)))]
2705 );
2706 assert_eq!(root.children_len(), 2);
2707
2708 assert_eq!(
2709 collect(btree.range(key(0)..key(10))),
2710 vec![
2711 (key(1), value(100)),
2712 (key(2), value(200)),
2713 (key(3), value(300)),
2714 (key(4), value(400)),
2715 ]
2716 );
2717
2718 assert_eq!(
2720 collect(btree.range(key(13)..key(20))),
2721 vec![(key(13), value(700)), (key(14), value(800))]
2722 );
2723
2724 assert_eq!(collect(btree.range(key(25)..)), vec![]);
2726 });
2727 }
2728 btree_test!(
2729 test_range_various_prefixes_with_offset,
2730 range_various_prefixes_with_offset
2731 );
2732
2733 fn range_various_prefixes_with_offset_2<K: TestKey, V: TestValue>() {
2734 let (key, value) = (K::build, V::build);
2735 run_btree_test(|mut btree| {
2736 btree.insert(key(1), value(0));
2737 btree.insert(key(2), value(0));
2738 btree.insert(key(3), value(0));
2739 btree.insert(key(4), value(0));
2740 btree.insert(key(12), value(0));
2741 btree.insert(key(14), value(0));
2742 btree.insert(key(16), value(0));
2743 btree.insert(key(18), value(0));
2744 btree.insert(key(19), value(0));
2745 btree.insert(key(21), value(0));
2746 btree.insert(key(22), value(0));
2747 btree.insert(key(23), value(0));
2748 btree.insert(key(24), value(0));
2749 btree.insert(key(25), value(0));
2750 btree.insert(key(26), value(0));
2751 btree.insert(key(27), value(0));
2752 btree.insert(key(28), value(0));
2753 btree.insert(key(29), value(0));
2754
2755 let root = btree.load_node(btree.root_addr);
2762 assert_eq!(root.node_type(), NodeType::Internal);
2763 assert_eq!(
2764 root.entries(btree.memory()),
2765 vec![(key(14), encode(value(0))), (key(23), encode(value(0)))]
2766 );
2767 assert_eq!(root.children_len(), 3);
2768
2769 let child_0 = btree.load_node(root.child(0));
2770 assert_eq!(child_0.node_type(), NodeType::Leaf);
2771 assert_eq!(
2772 child_0.entries(btree.memory()),
2773 vec![
2774 (key(1), encode(value(0))),
2775 (key(2), encode(value(0))),
2776 (key(3), encode(value(0))),
2777 (key(4), encode(value(0))),
2778 (key(12), encode(value(0))),
2779 ]
2780 );
2781
2782 let child_1 = btree.load_node(root.child(1));
2783 assert_eq!(child_1.node_type(), NodeType::Leaf);
2784 assert_eq!(
2785 child_1.entries(btree.memory()),
2786 vec![
2787 (key(16), encode(value(0))),
2788 (key(18), encode(value(0))),
2789 (key(19), encode(value(0))),
2790 (key(21), encode(value(0))),
2791 (key(22), encode(value(0))),
2792 ]
2793 );
2794
2795 let child_2 = btree.load_node(root.child(2));
2796 assert_eq!(
2797 child_2.entries(btree.memory()),
2798 vec![
2799 (key(24), encode(value(0))),
2800 (key(25), encode(value(0))),
2801 (key(26), encode(value(0))),
2802 (key(27), encode(value(0))),
2803 (key(28), encode(value(0))),
2804 (key(29), encode(value(0))),
2805 ]
2806 );
2807
2808 assert_eq!(
2810 collect(btree.range(key(14)..key(20))),
2811 vec![
2812 (key(14), value(0)),
2813 (key(16), value(0)),
2814 (key(18), value(0)),
2815 (key(19), value(0)),
2816 ]
2817 );
2818
2819 assert_eq!(
2822 collect(btree.range(key(22)..key(30))),
2823 vec![
2824 (key(22), value(0)),
2825 (key(23), value(0)),
2826 (key(24), value(0)),
2827 (key(25), value(0)),
2828 (key(26), value(0)),
2829 (key(27), value(0)),
2830 (key(28), value(0)),
2831 (key(29), value(0)),
2832 ]
2833 );
2834 });
2835 }
2836 btree_test!(
2837 test_range_various_prefixes_with_offset_2,
2838 range_various_prefixes_with_offset_2
2839 );
2840
2841 #[test]
2842 #[should_panic(expected = "max_key_size must be <= 4")]
2843 fn v1_rejects_increases_in_max_key_size() {
2844 let mem = make_memory();
2845 let btree: BTreeMap<Blob<4>, Blob<3>, _> = BTreeMap::init_v1(mem);
2846 let _btree: BTreeMap<Blob<5>, Blob<3>, _> = BTreeMap::init_v1(btree.into_memory());
2847 }
2848
2849 #[test]
2850 fn v2_handles_increases_in_max_key_size_and_max_value_size() {
2851 let mem = make_memory();
2852 let mut btree: BTreeMap<Blob<4>, Blob<4>, _> = BTreeMap::init(mem);
2853 btree.insert(
2854 [1u8; 4].as_slice().try_into().unwrap(),
2855 [1u8; 4].as_slice().try_into().unwrap(),
2856 );
2857
2858 let mut btree: BTreeMap<Blob<5>, Blob<5>, _> = BTreeMap::init(btree.into_memory());
2860 btree.insert(
2861 [2u8; 5].as_slice().try_into().unwrap(),
2862 [2u8; 5].as_slice().try_into().unwrap(),
2863 );
2864
2865 assert_eq!(
2867 btree.get(&([1u8; 4].as_slice().try_into().unwrap())),
2868 Some([1u8; 4].as_slice().try_into().unwrap())
2869 );
2870
2871 assert_eq!(
2872 btree.get(&([2u8; 5].as_slice().try_into().unwrap())),
2873 Some([2u8; 5].as_slice().try_into().unwrap())
2874 );
2875 }
2876
2877 #[test]
2878 fn accepts_small_or_equal_key_sizes() {
2879 let mem = make_memory();
2880 let btree: BTreeMap<Blob<4>, Blob<3>, _> = BTreeMap::init(mem);
2881 let btree: BTreeMap<Blob<3>, Blob<3>, _> = BTreeMap::init(btree.into_memory());
2883 let _btree: BTreeMap<Blob<4>, Blob<3>, _> = BTreeMap::init(btree.into_memory());
2885 }
2886
2887 #[test]
2888 #[should_panic(expected = "max_value_size must be <= 3")]
2889 fn v1_rejects_larger_value_sizes() {
2890 let mem = make_memory();
2891 let btree: BTreeMap<Blob<4>, Blob<3>, _> = BTreeMap::init_v1(mem);
2892 let _btree: BTreeMap<Blob<4>, Blob<4>, _> = BTreeMap::init_v1(btree.into_memory());
2893 }
2894
2895 #[test]
2896 fn accepts_small_or_equal_value_sizes() {
2897 let mem = make_memory();
2898 let btree: BTreeMap<Blob<4>, Blob<3>, _> = BTreeMap::init(mem);
2899 let btree: BTreeMap<Blob<4>, Blob<2>, _> = BTreeMap::init(btree.into_memory());
2901 let _btree: BTreeMap<Blob<4>, Blob<3>, _> = BTreeMap::init(btree.into_memory());
2903 }
2904
2905 fn bruteforce_range_search<K: TestKey, V: TestValue>() {
2906 let (key, value) = (K::build, V::build);
2907 run_btree_test(|mut stable_map| {
2908 use std::collections::BTreeMap;
2909 const NKEYS: u32 = 60;
2910 let mut std_map = BTreeMap::new();
2911
2912 for i in 0..NKEYS {
2913 std_map.insert(key(i), value(i));
2914 stable_map.insert(key(i), value(i));
2915 }
2916
2917 assert_eq!(collect_kv(std_map.range(..)), collect(stable_map.range(..)));
2918
2919 for l in 0..=NKEYS {
2920 assert_eq!(
2921 collect_kv(std_map.range(key(l)..)),
2922 collect(stable_map.range(key(l)..))
2923 );
2924
2925 assert_eq!(
2926 collect_kv(std_map.range(..key(l))),
2927 collect(stable_map.range(..key(l)))
2928 );
2929
2930 assert_eq!(
2931 collect_kv(std_map.range(..=key(l))),
2932 collect(stable_map.range(..=key(l)))
2933 );
2934
2935 for r in l + 1..=NKEYS {
2936 for lbound in [Bound::Included(key(l)), Bound::Excluded(key(l))] {
2937 for rbound in [Bound::Included(key(r)), Bound::Excluded(key(r))] {
2938 let range = (lbound.clone(), rbound);
2939 assert_eq!(
2940 collect_kv(std_map.range(range.clone())),
2941 collect(stable_map.range(range.clone())),
2942 "range: {range:?}"
2943 );
2944 }
2945 }
2946 }
2947 }
2948 });
2949 }
2950 btree_test!(test_bruteforce_range_search, bruteforce_range_search);
2951
2952 fn test_iter_upper_bound<K: TestKey, V: TestValue>() {
2953 let (key, value) = (K::build, V::build);
2954 run_btree_test(|mut btree| {
2955 for j in 0..100 {
2956 btree.insert(key(j), value(j));
2957 for i in 0..=j {
2958 assert_eq!(
2959 btree.iter_upper_bound(&key(i + 1)).next(),
2960 Some((key(i), value(i))),
2961 "failed to get an upper bound for key({})",
2962 i + 1
2963 );
2964 }
2965 assert_eq!(
2966 btree.iter_upper_bound(&key(0)).next(),
2967 None,
2968 "key(0) must not have an upper bound"
2969 );
2970 }
2971 });
2972 }
2973 btree_test!(test_test_iter_upper_bound, test_iter_upper_bound);
2974
2975 #[derive(Clone, Ord, PartialOrd, Eq, PartialEq)]
2977 struct BuggyStruct;
2978 impl crate::Storable for BuggyStruct {
2979 fn to_bytes(&self) -> Cow<[u8]> {
2980 Cow::Borrowed(&[1, 2, 3, 4])
2981 }
2982
2983 fn from_bytes(_: Cow<[u8]>) -> Self {
2984 unimplemented!();
2985 }
2986
2987 const BOUND: StorableBound = StorableBound::Bounded {
2988 max_size: 1,
2989 is_fixed_size: false,
2990 };
2991 }
2992
2993 #[test]
2994 #[should_panic(expected = "expected an element with length <= 1 bytes, but found 4")]
2995 fn v1_panics_if_key_is_bigger_than_max_size() {
2996 let mut btree = BTreeMap::init_v1(make_memory());
2997 btree.insert(BuggyStruct, ());
2998 }
2999
3000 #[test]
3001 #[should_panic(expected = "expected an element with length <= 1 bytes, but found 4")]
3002 fn v2_panics_if_key_is_bigger_than_max_size() {
3003 let mut btree = BTreeMap::init(make_memory());
3004 btree.insert(BuggyStruct, ());
3005 }
3006
3007 #[test]
3008 #[should_panic(expected = "expected an element with length <= 1 bytes, but found 4")]
3009 fn v1_panics_if_value_is_bigger_than_max_size() {
3010 let mut btree = BTreeMap::init(make_memory());
3011 btree.insert((), BuggyStruct);
3012 }
3013
3014 #[test]
3015 #[should_panic(expected = "expected an element with length <= 1 bytes, but found 4")]
3016 fn v2_panics_if_value_is_bigger_than_max_size() {
3017 let mut btree = BTreeMap::init(make_memory());
3018 btree.insert((), BuggyStruct);
3019 }
3020
3021 #[test]
3024 #[ignore]
3025 fn create_btreemap_dump_file() {
3026 let mem = make_memory();
3027 let mut btree = BTreeMap::init_v1(mem.clone());
3028
3029 let key = b(&[1, 2, 3]);
3030 let value = b(&[4, 5, 6]);
3031 assert_eq!(btree.insert(key, value), None);
3032 assert_eq!(btree.get(&key), Some(value));
3033
3034 use std::io::prelude::*;
3035 let mut file =
3036 std::fs::File::create(format!("dumps/btreemap_v{LAYOUT_VERSION}.dump")).unwrap();
3037 file.write_all(&mem.borrow()).unwrap();
3038 }
3039
3040 #[test]
3041 fn produces_layout_identical_to_layout_version_1_with_packed_headers() {
3042 let mem = make_memory();
3043 let mut btree = BTreeMap::init_v1(mem.clone());
3044
3045 let key = b(&[1, 2, 3]);
3046 let value = b(&[4, 5, 6]);
3047 assert_eq!(btree.insert(key, value), None);
3048 assert_eq!(btree.get(&key), Some(value));
3049
3050 let btreemap_v1 = include_bytes!("../dumps/btreemap_v1_packed_headers.dump");
3051 assert_eq!(*mem.borrow(), btreemap_v1);
3052 }
3053
3054 #[test]
3055 fn read_write_header_is_identical_to_read_write_struct() {
3056 #[repr(C, packed)]
3057 struct BTreePackedHeader {
3058 magic: [u8; 3],
3059 version: u8,
3060 max_key_size: u32,
3061 max_value_size: u32,
3062 root_addr: Address,
3063 length: u64,
3064 _buffer: [u8; 24],
3065 }
3066 let packed_header = BTreePackedHeader {
3067 magic: *MAGIC,
3068 version: LAYOUT_VERSION,
3069 root_addr: Address::from(0xDEADBEEF),
3070 max_key_size: 0x12345678,
3071 max_value_size: 0x87654321,
3072 length: 0xA1B2D3C4,
3073 _buffer: [0; 24],
3074 };
3075
3076 let packed_mem = make_memory();
3077 crate::write_struct(&packed_header, Address::from(0), &packed_mem);
3078
3079 let v1_header = BTreeHeader {
3080 version: Version::V1(DerivedPageSize {
3081 max_key_size: 0x12345678,
3082 max_value_size: 0x87654321,
3083 }),
3084 root_addr: Address::from(0xDEADBEEF),
3085 length: 0xA1B2D3C4,
3086 };
3087
3088 let v1_mem = make_memory();
3089 BTreeMap::<Vec<_>, Vec<_>, RefCell<Vec<_>>>::write_header(&v1_header, &v1_mem);
3090
3091 assert_eq!(packed_mem, v1_mem);
3092
3093 let packed_header: BTreePackedHeader = crate::read_struct(Address::from(0), &v1_mem);
3094 let v1_header = BTreeMap::<Vec<_>, Vec<_>, RefCell<Vec<_>>>::read_header(&v1_mem);
3095 assert!(packed_header.magic == *MAGIC);
3096 assert!(packed_header.version == LAYOUT_VERSION);
3097 match v1_header.version {
3098 Version::V1(DerivedPageSize {
3099 max_key_size,
3100 max_value_size,
3101 }) => {
3102 assert!(packed_header.max_key_size == max_key_size);
3103 assert!(packed_header.max_value_size == max_value_size);
3104 }
3105 _ => unreachable!("version must be v1"),
3106 };
3107
3108 assert!(packed_header.root_addr == v1_header.root_addr);
3109 assert!(packed_header.length == v1_header.length);
3110 }
3111
3112 #[test]
3113 fn migrate_from_bounded_to_unbounded_and_back() {
3114 #[derive(PartialOrd, Ord, Clone, Eq, PartialEq, Debug)]
3116 struct T;
3117 impl Storable for T {
3118 fn to_bytes(&self) -> Cow<[u8]> {
3119 Cow::Owned(vec![1, 2, 3])
3120 }
3121
3122 fn from_bytes(bytes: Cow<[u8]>) -> Self {
3123 assert_eq!(bytes.to_vec(), vec![1, 2, 3]);
3124 T
3125 }
3126
3127 const BOUND: StorableBound = StorableBound::Bounded {
3128 max_size: 3,
3129 is_fixed_size: true,
3130 };
3131 }
3132
3133 #[derive(PartialOrd, Ord, Clone, Eq, PartialEq, Debug)]
3135 struct T2;
3136 impl Storable for T2 {
3137 fn to_bytes(&self) -> Cow<[u8]> {
3138 Cow::Owned(vec![1, 2, 3])
3139 }
3140
3141 fn from_bytes(bytes: Cow<[u8]>) -> Self {
3142 assert_eq!(bytes.to_vec(), vec![1, 2, 3]);
3143 T2
3144 }
3145
3146 const BOUND: StorableBound = StorableBound::Unbounded;
3147 }
3148
3149 let mem = make_memory();
3151 let mut btree: BTreeMap<T, T, _> = BTreeMap::new_v1(mem);
3152 btree.insert(T, T);
3153
3154 let btree: BTreeMap<T2, T2, _> = BTreeMap::init(btree.into_memory());
3156 btree.save_header();
3157
3158 let btree: BTreeMap<T2, T2, _> = BTreeMap::init(btree.into_memory());
3160 assert_eq!(btree.get(&T2), Some(T2));
3161
3162 let btree: BTreeMap<T, T, _> = BTreeMap::init(btree.into_memory());
3164 assert_eq!(btree.get(&T), Some(T));
3165 }
3166
3167 #[test]
3168 fn test_clear_new_bounded_type() {
3169 let mem = make_memory();
3170 let mut btree: BTreeMap<Blob<4>, Blob<4>, _> = BTreeMap::new(mem.clone());
3171
3172 btree.insert(
3173 [1u8; 4].as_slice().try_into().unwrap(),
3174 [1u8; 4].as_slice().try_into().unwrap(),
3175 );
3176
3177 assert_ne!(btree.len(), 0);
3178 assert_ne!(btree.allocator.num_allocated_chunks(), 0);
3179 assert_ne!(btree.root_addr, NULL);
3180
3181 btree.clear_new();
3182
3183 let header_actual = BTreeMap::<Blob<4>, Blob<4>, _>::read_header(&mem);
3184
3185 BTreeMap::<Blob<4>, Blob<4>, _>::new(mem.clone());
3186
3187 let header_expected = BTreeMap::<Blob<4>, Blob<4>, _>::read_header(&mem);
3188
3189 assert_eq!(header_actual, header_expected);
3190 }
3191
3192 #[test]
3193 fn test_clear_new_unbounded_type() {
3194 let mem = make_memory();
3195 let mut btree: BTreeMap<String, String, _> = BTreeMap::new(mem.clone());
3196 btree.insert("asd".into(), "bce".into());
3197
3198 assert_ne!(btree.len(), 0);
3199 assert_ne!(btree.allocator.num_allocated_chunks(), 0);
3200 assert_ne!(btree.root_addr, NULL);
3201
3202 btree.clear_new();
3203
3204 let header_actual = BTreeMap::<String, String, _>::read_header(&mem);
3205
3206 BTreeMap::<String, String, _>::new(mem.clone());
3207
3208 let header_expected = BTreeMap::<String, String, _>::read_header(&mem);
3209
3210 assert_eq!(header_actual, header_expected);
3211 }
3212
3213 #[test]
3214 fn deallocating_node_with_overflows() {
3215 let mem = make_memory();
3216 let mut btree: BTreeMap<Vec<u8>, Vec<u8>, _> = BTreeMap::new(mem.clone());
3217
3218 assert_eq!(btree.allocator.num_allocated_chunks(), 0);
3220
3221 btree.insert(vec![0; 10_000], vec![]);
3223
3224 assert!(btree.allocator.num_allocated_chunks() >= 2);
3227 btree.remove(&vec![0; 10_000]);
3228
3229 assert_eq!(btree.allocator.num_allocated_chunks(), 0);
3231 }
3232
3233 #[test]
3234 fn repeatedly_deallocating_nodes_with_overflows() {
3235 let mem = make_memory();
3236 let mut btree: BTreeMap<Vec<u8>, Vec<u8>, _> = BTreeMap::new(mem.clone());
3237
3238 assert_eq!(btree.allocator.num_allocated_chunks(), 0);
3240
3241 for _ in 0..100 {
3242 for i in 0..100 {
3243 btree.insert(vec![i; 10_000], vec![]);
3244 }
3245
3246 for i in 0..100 {
3247 btree.remove(&vec![i; 10_000]);
3248 }
3249 }
3250
3251 assert_eq!(btree.allocator.num_allocated_chunks(), 0);
3253 }
3254
3255 #[test]
3256 fn deallocating_root_does_not_leak_memory() {
3257 let mem = make_memory();
3258 let mut btree: BTreeMap<Vec<u8>, _, _> = BTreeMap::new(mem.clone());
3259
3260 for i in 1..=11 {
3261 assert_eq!(btree.insert(vec![i; 10_000], ()), None);
3263 }
3264
3265 assert_eq!(btree.insert(vec![0; 10_000], ()), None);
3267
3268 let root = btree.load_node(btree.root_addr);
3273 assert_eq!(root.node_type(), NodeType::Internal);
3274 assert_eq!(root.keys(btree.memory()), vec![&[6; 10_000]]);
3275 assert_eq!(root.children_len(), 2);
3276
3277 btree.remove(&vec![6; 10_000]);
3279
3280 let root = btree.load_node(btree.root_addr);
3285 assert_eq!(root.node_type(), NodeType::Internal);
3286 assert_eq!(root.keys(btree.memory()), vec![&[5; 10_000]]);
3287 assert_eq!(root.children_len(), 2);
3288
3289 btree.remove(&vec![5; 10_000]);
3292
3293 let root = btree.load_node(btree.root_addr);
3296 assert_eq!(root.node_type(), NodeType::Leaf);
3297 assert_eq!(
3298 root.keys(btree.memory()),
3299 vec![
3300 &[0; 10_000],
3301 &[1; 10_000],
3302 &[2; 10_000],
3303 &[3; 10_000],
3304 &[4; 10_000],
3305 &[7; 10_000],
3306 &[8; 10_000],
3307 &[9; 10_000],
3308 &[10; 10_000],
3309 &[11; 10_000],
3310 ]
3311 );
3312
3313 for i in 0..=11 {
3315 btree.remove(&vec![i; 10_000]);
3316 }
3317
3318 assert_eq!(btree.allocator.num_allocated_chunks(), 0);
3319 }
3320}