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>
234where
235 K: Storable + Ord + Clone,
236 V: Storable,
237 M: Memory,
238{
239 root_addr: Address,
242
243 version: Version,
244
245 allocator: Allocator<M>,
247
248 length: u64,
250
251 _phantom: PhantomData<(K, V)>,
253}
254
255#[derive(PartialEq, Debug)]
256struct BTreeHeader {
258 version: Version,
259 root_addr: Address,
260 length: u64,
261 }
263
264impl<K, V, M> BTreeMap<K, V, M>
265where
266 K: Storable + Ord + Clone,
267 V: Storable,
268 M: Memory,
269{
270 pub fn init(memory: M) -> Self {
275 if memory.size() == 0 {
276 return BTreeMap::new(memory);
278 }
279
280 let mut dst = vec![0; 3];
282 memory.read(0, &mut dst);
283 if dst != MAGIC {
284 BTreeMap::new(memory)
286 } else {
287 BTreeMap::load(memory)
289 }
290 }
291
292 #[cfg(test)]
296 pub fn init_v1(memory: M) -> Self {
297 if memory.size() == 0 {
298 return BTreeMap::new_v1(memory);
300 }
301
302 let mut dst = vec![0; 3];
304 memory.read(0, &mut dst);
305 if dst != MAGIC {
306 BTreeMap::new_v1(memory)
308 } else {
309 BTreeMap::load_helper(memory, false)
312 }
313 }
314
315 pub fn new(memory: M) -> Self {
327 let page_size = match (K::BOUND, V::BOUND) {
328 (
330 StorableBound::Bounded {
331 max_size: max_key_size,
332 ..
333 },
334 StorableBound::Bounded {
335 max_size: max_value_size,
336 ..
337 },
338 ) => {
339 let max_node_size = Node::<K>::max_size(max_key_size, max_value_size).get();
341
342 PageSize::Value((max_node_size * 3 / 4) as u32)
347 }
348 _ => PageSize::Value(DEFAULT_PAGE_SIZE),
350 };
351
352 let btree = Self {
353 root_addr: NULL,
354 allocator: Allocator::new(
355 memory,
356 Address::from(ALLOCATOR_OFFSET as u64),
357 page_size.get().into(),
358 ),
359 version: Version::V2(page_size),
360 length: 0,
361 _phantom: PhantomData,
362 };
363
364 btree.save_header();
365 btree
366 }
367
368 #[cfg(test)]
372 pub fn new_v1(memory: M) -> Self {
373 let max_key_size = K::BOUND.max_size();
374 let max_value_size = V::BOUND.max_size();
375
376 let btree = Self {
377 root_addr: NULL,
378 allocator: Allocator::new(
379 memory,
380 Address::from(ALLOCATOR_OFFSET as u64),
381 Node::<K>::max_size(max_key_size, max_value_size),
382 ),
383 version: Version::V1(DerivedPageSize {
384 max_key_size,
385 max_value_size,
386 }),
387 length: 0,
388 _phantom: PhantomData,
389 };
390
391 btree.save_header();
392 btree
393 }
394
395 pub fn load(memory: M) -> Self {
397 Self::load_helper(memory, true)
398 }
399
400 fn load_helper(memory: M, migrate_to_v2: bool) -> Self {
402 let header = Self::read_header(&memory);
404
405 let version = match header.version {
406 Version::V1(derived_page_size) => {
407 if migrate_to_v2 {
409 Version::V2(PageSize::Derived(derived_page_size))
410 } else {
411 let max_key_size = derived_page_size.max_key_size;
413 let max_value_size = derived_page_size.max_value_size;
414
415 assert!(
416 K::BOUND.max_size() <= max_key_size,
417 "max_key_size must be <= {max_key_size}",
418 );
419
420 assert!(
421 V::BOUND.max_size() <= max_value_size,
422 "max_value_size must be <= {max_value_size}"
423 );
424
425 Version::V1(derived_page_size)
426 }
427 }
428 other => other,
429 };
430
431 let allocator_addr = Address::from(ALLOCATOR_OFFSET as u64);
432 Self {
433 root_addr: header.root_addr,
434 allocator: Allocator::load(memory, allocator_addr),
435 version,
436 length: header.length,
437 _phantom: PhantomData,
438 }
439 }
440
441 fn read_header(memory: &M) -> BTreeHeader {
443 let mut buf = [0; PACKED_HEADER_SIZE];
445 memory.read(0, &mut buf);
446
447 assert_eq!(&buf[0..3], MAGIC, "Bad magic.");
448
449 match buf[3] {
450 LAYOUT_VERSION => {
451 BTreeHeader {
453 version: Version::V1(DerivedPageSize {
454 max_key_size: u32::from_le_bytes(buf[4..8].try_into().unwrap()),
455 max_value_size: u32::from_le_bytes(buf[8..12].try_into().unwrap()),
456 }),
457 root_addr: Address::from(u64::from_le_bytes(buf[12..20].try_into().unwrap())),
458 length: u64::from_le_bytes(buf[20..28].try_into().unwrap()),
459 }
460 }
461 LAYOUT_VERSION_2 => {
462 let page_size = {
464 let a = u32::from_le_bytes(buf[4..8].try_into().unwrap());
466 let b = u32::from_le_bytes(buf[8..12].try_into().unwrap());
467
468 if b == PAGE_SIZE_VALUE_MARKER {
469 PageSize::Value(a)
471 } else {
472 PageSize::Derived(DerivedPageSize {
474 max_key_size: a,
475 max_value_size: b,
476 })
477 }
478 };
479
480 BTreeHeader {
482 version: Version::V2(page_size),
483 root_addr: Address::from(u64::from_le_bytes(buf[12..20].try_into().unwrap())),
484 length: u64::from_le_bytes(buf[20..28].try_into().unwrap()),
485 }
486 }
487 version => {
488 panic!("Unsupported version: {version}.");
489 }
490 }
491 }
492
493 pub fn insert(&mut self, key: K, value: V) -> Option<V> {
501 let value = value.into_bytes_checked();
502
503 let root = if self.root_addr == NULL {
504 let node = self.allocate_node(NodeType::Leaf);
506 self.root_addr = node.address();
507 self.save_header();
508 node
509 } else {
510 let mut root = self.load_node(self.root_addr);
512
513 if let Ok(idx) = root.search(&key, self.memory()) {
515 return Some(V::from_bytes(Cow::Owned(
517 self.update_value(&mut root, idx, value),
518 )));
519 }
520
521 if root.is_full() {
527 let mut new_root = self.allocate_node(NodeType::Internal);
529
530 new_root.push_child(self.root_addr);
532
533 self.root_addr = new_root.address();
535 self.save_header();
536
537 self.split_child(&mut new_root, 0);
539
540 new_root
541 } else {
542 root
543 }
544 };
545
546 self.insert_nonfull(root, key, value)
547 .map(Cow::Owned)
548 .map(V::from_bytes)
549 }
550
551 fn insert_nonfull(&mut self, mut node: Node<K>, key: K, value: Vec<u8>) -> Option<Vec<u8>> {
553 assert!(!node.is_full());
555
556 match node.search(&key, self.memory()) {
558 Ok(idx) => {
559 Some(self.update_value(&mut node, idx, value))
561 }
562 Err(idx) => {
563 match node.node_type() {
566 NodeType::Leaf => {
567 node.insert_entry(idx, (key, value));
570 self.save_node(&mut node);
571
572 self.length += 1;
574 self.save_header();
575
576 None
578 }
579 NodeType::Internal => {
580 let mut child = self.load_node(node.child(idx));
583
584 if child.is_full() {
585 if let Ok(idx) = child.search(&key, self.memory()) {
587 return Some(self.update_value(&mut child, idx, value));
589 }
590
591 self.split_child(&mut node, idx);
593
594 let idx = node.search(&key, self.memory()).unwrap_or_else(|idx| idx);
597 child = self.load_node(node.child(idx));
598 }
599
600 assert!(!child.is_full());
602
603 self.insert_nonfull(child, key, value)
604 }
605 }
606 }
607 }
608 }
609
610 fn split_child(&mut self, node: &mut Node<K>, full_child_idx: usize) {
628 assert!(!node.is_full());
630
631 let mut full_child = self.load_node(node.child(full_child_idx));
633 assert!(full_child.is_full());
634
635 let mut sibling = self.allocate_node(full_child.node_type());
637 assert_eq!(sibling.node_type(), full_child.node_type());
638
639 node.insert_child(full_child_idx + 1, sibling.address());
641
642 let (median_key, median_value) = full_child.split(&mut sibling, self.memory());
643
644 node.insert_entry(full_child_idx, (median_key, median_value));
645
646 self.save_node(&mut sibling);
647 self.save_node(&mut full_child);
648 self.save_node(node);
649 }
650
651 pub fn get(&self, key: &K) -> Option<V> {
653 if self.root_addr == NULL {
654 return None;
655 }
656 self.traverse(self.root_addr, key, |node, idx| {
657 node.extract_entry_at(idx, self.memory()).1 })
659 .map(Cow::Owned)
660 .map(V::from_bytes)
661 }
662
663 pub fn contains_key(&self, key: &K) -> bool {
665 self.root_addr != NULL && self.traverse(self.root_addr, key, |_, _| ()).is_some()
667 }
668
669 fn traverse<F, R>(&self, node_addr: Address, key: &K, f: F) -> Option<R>
671 where
672 F: Fn(&mut Node<K>, usize) -> R,
673 {
674 let mut node = self.load_node(node_addr);
675 match node.search(key, self.memory()) {
677 Ok(idx) => Some(f(&mut node, idx)), Err(idx) => match node.node_type() {
679 NodeType::Leaf => None, NodeType::Internal => self.traverse(node.child(idx), key, f), },
682 }
683 }
684
685 pub fn is_empty(&self) -> bool {
687 self.length == 0
688 }
689
690 pub fn len(&self) -> u64 {
692 self.length
693 }
694
695 pub fn into_memory(self) -> M {
697 self.allocator.into_memory()
698 }
699
700 #[deprecated(since = "0.6.3", note = "please use `clear_new` instead")]
702 pub fn clear(self) -> Self {
703 let mem = self.allocator.into_memory();
704 Self::new(mem)
705 }
706
707 pub fn clear_new(&mut self) {
709 self.root_addr = NULL;
710 self.length = 0;
711 self.allocator.clear();
712 self.save_header();
713 }
714
715 pub fn first_key_value(&self) -> Option<(K, V)> {
718 if self.root_addr == NULL {
719 return None;
720 }
721 let root = self.load_node(self.root_addr);
722 let (k, encoded_v) = root.get_min(self.memory());
723 Some((k, V::from_bytes(Cow::Owned(encoded_v))))
724 }
725
726 pub fn last_key_value(&self) -> Option<(K, V)> {
729 if self.root_addr == NULL {
730 return None;
731 }
732 let root = self.load_node(self.root_addr);
733 let (k, encoded_v) = root.get_max(self.memory());
734 Some((k, V::from_bytes(Cow::Owned(encoded_v))))
735 }
736
737 fn memory(&self) -> &M {
738 self.allocator.memory()
739 }
740
741 fn allocator_mut(&mut self) -> &mut Allocator<M> {
742 &mut self.allocator
743 }
744
745 pub fn remove(&mut self, key: &K) -> Option<V> {
747 if self.root_addr == NULL {
748 return None;
749 }
750
751 let root_node = self.load_node(self.root_addr);
752 self.remove_helper(root_node, key)
753 .map(Cow::Owned)
754 .map(V::from_bytes)
755 }
756
757 pub fn pop_last(&mut self) -> Option<(K, V)> {
759 if self.root_addr == NULL {
760 return None;
761 }
762
763 let root = self.load_node(self.root_addr);
764 let (max_key, _) = root.get_max(self.memory());
765 self.remove_helper(root, &max_key)
766 .map(|v| (max_key, V::from_bytes(Cow::Owned(v))))
767 }
768
769 pub fn pop_first(&mut self) -> Option<(K, V)> {
771 if self.root_addr == NULL {
772 return None;
773 }
774
775 let root = self.load_node(self.root_addr);
776 let (min_key, _) = root.get_min(self.memory());
777 self.remove_helper(root, &min_key)
778 .map(|v| (min_key, V::from_bytes(Cow::Owned(v))))
779 }
780
781 fn remove_helper(&mut self, mut node: Node<K>, key: &K) -> Option<Vec<u8>> {
783 if node.address() != self.root_addr {
784 assert!(node.can_remove_entry_without_merging());
789 }
790
791 match node.node_type() {
792 NodeType::Leaf => {
793 match node.search(key, self.memory()) {
794 Ok(idx) => {
795 let value = node.remove_entry(idx, self.memory()).1;
798 self.length -= 1;
799
800 if node.entries_len() == 0 {
801 assert_eq!(
802 node.address(), self.root_addr,
803 "Removal can only result in an empty leaf node if that node is the root"
804 );
805
806 self.deallocate_node(node);
808 self.root_addr = NULL;
809 } else {
810 self.save_node(&mut node);
811 }
812
813 self.save_header();
814 Some(value)
815 }
816 _ => None, }
818 }
819 NodeType::Internal => {
820 match node.search(key, self.memory()) {
821 Ok(idx) => {
822 let left_child = self.load_node(node.child(idx));
825 if left_child.can_remove_entry_without_merging() {
826 let predecessor = left_child.get_max(self.memory());
849 self.remove_helper(left_child, &predecessor.0)?;
850
851 let (_, old_value) = node.swap_entry(idx, predecessor, self.memory());
853
854 self.save_node(&mut node);
856 return Some(old_value);
857 }
858
859 let right_child = self.load_node(node.child(idx + 1));
860 if right_child.can_remove_entry_without_merging() {
861 let successor = right_child.get_min(self.memory());
884 self.remove_helper(right_child, &successor.0)?;
885
886 let (_, old_value) = node.swap_entry(idx, successor, self.memory());
888
889 self.save_node(&mut node);
891 return Some(old_value);
892 }
893
894 assert!(left_child.at_minimum());
914 assert!(right_child.at_minimum());
915
916 let mut new_child = self.merge(
918 right_child,
919 left_child,
920 node.remove_entry(idx, self.memory()),
921 );
922
923 node.remove_child(idx + 1);
925
926 if node.entries_len() == 0 {
927 assert_eq!(node.address(), self.root_addr);
929 assert_eq!(node.child(0), new_child.address());
930 assert_eq!(node.children_len(), 1);
931
932 self.root_addr = new_child.address();
933
934 self.deallocate_node(node);
936 self.save_header();
937 } else {
938 self.save_node(&mut node);
939 }
940
941 self.save_node(&mut new_child);
942
943 self.remove_helper(new_child, key)
945 }
946 Err(idx) => {
947 let mut child = self.load_node(node.child(idx));
952
953 if child.can_remove_entry_without_merging() {
954 return self.remove_helper(child, key);
957 }
958
959 let mut left_sibling = if idx > 0 {
962 Some(self.load_node(node.child(idx - 1)))
963 } else {
964 None
965 };
966
967 let mut right_sibling = if idx + 1 < node.children_len() {
968 Some(self.load_node(node.child(idx + 1)))
969 } else {
970 None
971 };
972
973 if let Some(ref mut left_sibling) = left_sibling {
974 if left_sibling.can_remove_entry_without_merging() {
975 let (left_sibling_key, left_sibling_value) =
998 left_sibling.pop_entry(self.memory()).unwrap();
999
1000 let (parent_key, parent_value) = node.swap_entry(
1002 idx - 1,
1003 (left_sibling_key, left_sibling_value),
1004 self.memory(),
1005 );
1006
1007 child.insert_entry(0, (parent_key, parent_value));
1009
1010 if let Some(last_child) = left_sibling.pop_child() {
1012 assert_eq!(left_sibling.node_type(), NodeType::Internal);
1013 assert_eq!(child.node_type(), NodeType::Internal);
1014
1015 child.insert_child(0, last_child);
1016 } else {
1017 assert_eq!(left_sibling.node_type(), NodeType::Leaf);
1018 assert_eq!(child.node_type(), NodeType::Leaf);
1019 }
1020
1021 self.save_node(left_sibling);
1022 self.save_node(&mut child);
1023 self.save_node(&mut node);
1024 return self.remove_helper(child, key);
1025 }
1026 }
1027
1028 if let Some(right_sibling) = &mut right_sibling {
1029 if right_sibling.can_remove_entry_without_merging() {
1030 let (right_sibling_key, right_sibling_value) =
1053 right_sibling.remove_entry(0, self.memory());
1054
1055 let parent_entry = node.swap_entry(
1057 idx,
1058 (right_sibling_key, right_sibling_value),
1059 self.memory(),
1060 );
1061
1062 child.push_entry(parent_entry);
1064
1065 match right_sibling.node_type() {
1067 NodeType::Internal => {
1068 assert_eq!(child.node_type(), NodeType::Internal);
1069 child.push_child(right_sibling.remove_child(0));
1070 }
1071 NodeType::Leaf => {
1072 assert_eq!(child.node_type(), NodeType::Leaf);
1073 }
1074 }
1075
1076 self.save_node(right_sibling);
1077 self.save_node(&mut child);
1078 self.save_node(&mut node);
1079 return self.remove_helper(child, key);
1080 }
1081 }
1082
1083 if let Some(left_sibling) = left_sibling {
1086 assert!(left_sibling.at_minimum());
1089 let left_sibling = self.merge(
1090 child,
1091 left_sibling,
1092 node.remove_entry(idx - 1, self.memory()),
1093 );
1094 node.remove_child(idx);
1096
1097 if node.entries_len() == 0 {
1098 let node_address = node.address();
1099 self.deallocate_node(node);
1100
1101 if node_address == self.root_addr {
1102 self.root_addr = left_sibling.address();
1104 self.save_header();
1105 }
1106 } else {
1107 self.save_node(&mut node);
1108 }
1109
1110 return self.remove_helper(left_sibling, key);
1111 }
1112
1113 if let Some(right_sibling) = right_sibling {
1114 assert!(right_sibling.at_minimum());
1117 let right_sibling = self.merge(
1118 child,
1119 right_sibling,
1120 node.remove_entry(idx, self.memory()),
1121 );
1122
1123 node.remove_child(idx);
1125
1126 if node.entries_len() == 0 {
1127 let node_address = node.address();
1128 self.deallocate_node(node);
1129
1130 if node_address == self.root_addr {
1131 self.root_addr = right_sibling.address();
1133 self.save_header();
1134 }
1135 } else {
1136 self.save_node(&mut node);
1137 }
1138
1139 return self.remove_helper(right_sibling, key);
1140 }
1141
1142 unreachable!("At least one of the siblings must exist.");
1143 }
1144 }
1145 }
1146 }
1147 }
1148
1149 pub fn iter(&self) -> Iter<'_, K, V, M> {
1165 self.iter_internal().into()
1166 }
1167
1168 pub fn range(&self, key_range: impl RangeBounds<K>) -> Iter<'_, K, V, M> {
1188 self.range_internal(key_range).into()
1189 }
1190
1191 pub fn iter_from_prev_key(&self, bound: &K) -> Iter<'_, K, V, M> {
1198 if let Some(entry) = self.range(..bound).next_back() {
1199 let start_key = entry.key().clone();
1200 IterInternal::new_in_range(self, (Bound::Included(start_key), Bound::Unbounded)).into()
1201 } else {
1202 IterInternal::null(self).into()
1203 }
1204 }
1205
1206 #[deprecated(note = "use `iter_from_prev_key` instead")]
1215 pub fn iter_upper_bound(&self, bound: &K) -> Iter<'_, K, V, M> {
1216 self.iter_from_prev_key(bound)
1217 }
1218
1219 pub fn keys(&self) -> KeysIter<'_, K, V, M> {
1221 self.iter_internal().into()
1222 }
1223
1224 pub fn keys_range(&self, key_range: impl RangeBounds<K>) -> KeysIter<'_, K, V, M> {
1226 self.range_internal(key_range).into()
1227 }
1228
1229 pub fn values(&self) -> ValuesIter<'_, K, V, M> {
1231 self.iter_internal().into()
1232 }
1233
1234 pub fn values_range(&self, key_range: impl RangeBounds<K>) -> ValuesIter<'_, K, V, M> {
1237 self.range_internal(key_range).into()
1238 }
1239
1240 fn iter_internal(&self) -> IterInternal<'_, K, V, M> {
1241 IterInternal::new(self)
1242 }
1243
1244 fn range_internal(&self, key_range: impl RangeBounds<K>) -> IterInternal<'_, K, V, M> {
1245 if self.root_addr == NULL {
1246 return IterInternal::null(self);
1248 }
1249
1250 let range = (
1251 key_range.start_bound().cloned(),
1252 key_range.end_bound().cloned(),
1253 );
1254
1255 IterInternal::new_in_range(self, range)
1256 }
1257
1258 fn merge(&mut self, source: Node<K>, mut into: Node<K>, median: Entry<K>) -> Node<K> {
1271 into.merge(source, median, &mut self.allocator);
1272 into
1273 }
1274
1275 fn allocate_node(&mut self, node_type: NodeType) -> Node<K> {
1277 match self.version {
1278 Version::V1(page_size) => Node::new_v1(self.allocator.allocate(), node_type, page_size),
1279 Version::V2(page_size) => Node::new_v2(self.allocator.allocate(), node_type, page_size),
1280 }
1281 }
1282
1283 #[inline]
1285 fn deallocate_node(&mut self, node: Node<K>) {
1286 node.deallocate(self.allocator_mut());
1287 }
1288
1289 #[inline]
1291 fn load_node(&self, address: Address) -> Node<K> {
1292 Node::load(address, self.version.page_size(), self.memory())
1293 }
1294
1295 #[inline]
1297 fn save_node(&mut self, node: &mut Node<K>) {
1298 node.save(self.allocator_mut());
1299 }
1300
1301 fn update_value(&mut self, node: &mut Node<K>, idx: usize, new_value: Vec<u8>) -> Vec<u8> {
1303 let old_value = node.swap_value(idx, new_value, self.memory());
1304 self.save_node(node);
1305 old_value
1306 }
1307
1308 fn save_header(&self) {
1310 let header = BTreeHeader {
1311 version: self.version,
1312 root_addr: self.root_addr,
1313 length: self.length,
1314 };
1315
1316 Self::write_header(&header, self.memory());
1317 }
1318
1319 fn write_header(header: &BTreeHeader, memory: &M) {
1321 let mut buf = [0; PACKED_HEADER_SIZE];
1323 buf[0..3].copy_from_slice(MAGIC.as_slice());
1324 match header.version {
1325 Version::V1(DerivedPageSize {
1326 max_key_size,
1327 max_value_size,
1328 })
1329 | Version::V2(PageSize::Derived(DerivedPageSize {
1330 max_key_size,
1331 max_value_size,
1332 })) => {
1333 buf[3] = LAYOUT_VERSION;
1334 buf[4..8].copy_from_slice(&max_key_size.to_le_bytes());
1335 buf[8..12].copy_from_slice(&max_value_size.to_le_bytes());
1336 }
1337 Version::V2(PageSize::Value(page_size)) => {
1338 buf[3] = LAYOUT_VERSION_2;
1339 buf[4..8].copy_from_slice(&page_size.to_le_bytes());
1340 buf[8..12].copy_from_slice(&PAGE_SIZE_VALUE_MARKER.to_le_bytes());
1341 }
1342 };
1343 buf[12..20].copy_from_slice(&header.root_addr.get().to_le_bytes());
1344 buf[20..28].copy_from_slice(&header.length.to_le_bytes());
1345 crate::write(memory, 0, &buf);
1347 }
1348}
1349
1350#[cfg(test)]
1351mod test {
1352 use super::*;
1353 use crate::{
1354 btreemap::iter::LazyEntry,
1355 storable::{Blob, Bound as StorableBound},
1356 VectorMemory,
1357 };
1358 use std::cell::RefCell;
1359 use std::convert::TryFrom;
1360 use std::rc::Rc;
1361
1362 fn collect_kv<'a, K: Clone + 'a, V: Clone + 'a>(
1364 iter: impl Iterator<Item = (&'a K, &'a V)>,
1365 ) -> Vec<(K, V)> {
1366 iter.map(|(k, v)| (k.clone(), v.clone())).collect()
1367 }
1368
1369 fn collect<K: Clone, V: Clone>(it: impl Iterator<Item = (K, V)>) -> Vec<(K, V)> {
1371 it.collect()
1372 }
1373
1374 fn collect_entry<'a, K, V, M>(it: impl Iterator<Item = LazyEntry<'a, K, V, M>>) -> Vec<(K, V)>
1376 where
1377 K: Storable + Ord + Clone + 'a,
1378 V: Storable + 'a,
1379 M: Memory + 'a,
1380 {
1381 it.map(|entry| entry.into_pair()).collect()
1382 }
1383
1384 fn monotonic_buffer<const N: usize>(i: u32) -> [u8; N] {
1386 let mut buf = [0u8; N];
1387 let bytes = i.to_be_bytes();
1388 buf[N - bytes.len()..].copy_from_slice(&bytes);
1389 buf
1390 }
1391
1392 #[test]
1393 fn test_monotonic_buffer() {
1394 let cases: &[(u32, [u8; 4])] = &[
1395 (1, [0, 0, 0, 1]),
1396 (2, [0, 0, 0, 2]),
1397 ((1 << 8) - 1, [0, 0, 0, 255]),
1398 ((1 << 8), [0, 0, 1, 0]),
1399 ((1 << 16) - 1, [0, 0, 255, 255]),
1400 (1 << 16, [0, 1, 0, 0]),
1401 ((1 << 24) - 1, [0, 255, 255, 255]),
1402 (1 << 24, [1, 0, 0, 0]),
1403 ];
1404
1405 for &(input, expected) in cases {
1406 let output = monotonic_buffer::<4>(input);
1407 assert_eq!(
1408 output, expected,
1409 "monotonic_buffer::<4>({}) returned {:?}, expected {:?}",
1410 input, output, expected
1411 );
1412 }
1413 }
1414
1415 trait Builder {
1417 fn build(i: u32) -> Self;
1418 fn empty() -> Self;
1419 }
1420
1421 impl Builder for () {
1422 fn build(_i: u32) -> Self {}
1423 fn empty() -> Self {}
1424 }
1425
1426 impl Builder for u32 {
1427 fn build(i: u32) -> Self {
1428 i
1429 }
1430 fn empty() -> Self {
1431 0
1432 }
1433 }
1434
1435 impl<const N: usize> Builder for Blob<N> {
1436 fn build(i: u32) -> Self {
1437 Blob::try_from(&monotonic_buffer::<N>(i)[..]).unwrap()
1438 }
1439 fn empty() -> Self {
1440 Blob::try_from(&[][..]).unwrap()
1441 }
1442 }
1443
1444 type MonotonicVec32 = Vec<u8>;
1445 impl Builder for MonotonicVec32 {
1446 fn build(i: u32) -> Self {
1447 monotonic_buffer::<32>(i).to_vec()
1448 }
1449 fn empty() -> Self {
1450 Vec::new()
1451 }
1452 }
1453
1454 type MonotonicString32 = String;
1455 impl Builder for MonotonicString32 {
1456 fn build(i: u32) -> Self {
1457 format!("{i:0>32}")
1458 }
1459 fn empty() -> Self {
1460 String::new()
1461 }
1462 }
1463
1464 fn encode<T: Storable>(object: T) -> Vec<u8> {
1466 object.into_bytes_checked()
1467 }
1468
1469 pub(crate) fn b(x: &[u8]) -> Blob<10> {
1471 Blob::<10>::try_from(x).unwrap()
1472 }
1473
1474 pub(crate) fn make_memory() -> Rc<RefCell<Vec<u8>>> {
1476 Rc::new(RefCell::new(Vec::new()))
1477 }
1478
1479 pub fn run_btree_test<K, V, R, F>(f: F)
1481 where
1482 K: Storable + Ord + Clone,
1483 V: Storable,
1484 F: Fn(BTreeMap<K, V, VectorMemory>) -> R,
1485 {
1486 if K::BOUND != StorableBound::Unbounded && V::BOUND != StorableBound::Unbounded {
1488 let mem = make_memory();
1490 let tree_v1 = BTreeMap::new_v1(mem);
1491 f(tree_v1);
1492
1493 let mem = make_memory();
1495 let tree_v1: BTreeMap<K, V, _> = BTreeMap::new_v1(mem);
1496 let tree_v2_migrated = BTreeMap::load_helper(tree_v1.into_memory(), true);
1497 f(tree_v2_migrated);
1498 }
1499
1500 let mem = make_memory();
1502 let tree_v2 = BTreeMap::new(mem);
1503 f(tree_v2);
1504 }
1505
1506 fn verify_monotonic<T: Builder + PartialOrd>() {
1509 for shift_bits in [8, 16, 24] {
1510 let i = (1 << shift_bits) - 1;
1511 assert!(
1512 T::build(i) < T::build(i + 1),
1513 "Monotonicity failed at i: {i}",
1514 );
1515 }
1516 }
1517
1518 macro_rules! assert_bounded {
1520 ($ty:ty) => {
1521 assert_ne!(<$ty>::BOUND, StorableBound::Unbounded, "Must be Bounded");
1522 };
1523 }
1524
1525 macro_rules! assert_unbounded {
1527 ($ty:ty) => {
1528 assert_eq!(<$ty>::BOUND, StorableBound::Unbounded, "Must be Unbounded");
1529 };
1530 }
1531
1532 macro_rules! run_with_key {
1533 ($runner:ident, $K:ty) => {{
1534 verify_monotonic::<$K>();
1535
1536 $runner::<$K, ()>();
1538
1539 assert_bounded!(u32);
1541 $runner::<$K, u32>();
1542
1543 assert_bounded!(Blob<20>);
1544 $runner::<$K, Blob<20>>();
1545
1546 assert_unbounded!(MonotonicVec32);
1548 $runner::<$K, MonotonicVec32>();
1549
1550 assert_unbounded!(MonotonicString32);
1551 $runner::<$K, MonotonicString32>();
1552 }};
1553 }
1554
1555 macro_rules! btree_test {
1557 ($name:ident, $runner:ident) => {
1558 #[test]
1559 fn $name() {
1560 assert_bounded!(u32);
1562 run_with_key!($runner, u32);
1563
1564 assert_bounded!(Blob<10>);
1565 run_with_key!($runner, Blob<10>);
1566
1567 assert_unbounded!(MonotonicVec32);
1569 run_with_key!($runner, MonotonicVec32);
1570
1571 assert_unbounded!(MonotonicString32);
1572 run_with_key!($runner, MonotonicString32);
1573 }
1574 };
1575 }
1576
1577 trait TestKey: Storable + Ord + Clone + Builder + std::fmt::Debug {}
1579 impl<T> TestKey for T where T: Storable + Ord + Clone + Builder + std::fmt::Debug {}
1580
1581 trait TestValue: Storable + Clone + Builder + std::fmt::Debug + PartialEq {}
1583 impl<T> TestValue for T where T: Storable + Clone + Builder + std::fmt::Debug + PartialEq {}
1584
1585 fn insert_get_init_preserves_data<K: TestKey, V: TestValue>() {
1586 let (key, value) = (K::build, V::build);
1587 run_btree_test(|mut btree| {
1588 let n = 1_000;
1589 for i in 0..n {
1590 assert_eq!(btree.insert(key(i), value(i)), None);
1591 assert_eq!(btree.get(&key(i)), Some(value(i)));
1592 }
1593
1594 let btree = BTreeMap::<K, V, VectorMemory>::init(btree.into_memory());
1596 for i in 0..n {
1597 assert_eq!(btree.get(&key(i)), Some(value(i)));
1598 }
1599 });
1600 }
1601 btree_test!(
1602 test_insert_get_init_preserves_data,
1603 insert_get_init_preserves_data
1604 );
1605
1606 fn insert_overwrites_previous_value<K: TestKey, V: TestValue>() {
1607 let (key, value) = (K::build, V::build);
1608 run_btree_test(|mut btree| {
1609 let n = 1_000;
1610 for i in 0..n {
1611 assert_eq!(btree.insert(key(i), value(i)), None);
1612 assert_eq!(btree.insert(key(i), value(i + 1)), Some(value(i)));
1613 assert_eq!(btree.get(&key(i)), Some(value(i + 1)));
1614 }
1615 });
1616 }
1617 btree_test!(
1618 test_insert_overwrites_previous_value,
1619 insert_overwrites_previous_value
1620 );
1621
1622 fn insert_same_key_many<K: TestKey, V: TestValue>() {
1623 let (key, value) = (K::build, V::build);
1624 run_btree_test(|mut btree| {
1625 let n = 1_000;
1626 assert_eq!(btree.insert(key(1), value(2)), None);
1627 for i in 2..n {
1628 assert_eq!(btree.insert(key(1), value(i + 1)), Some(value(i)));
1629 }
1630 assert_eq!(btree.get(&key(1)), Some(value(n)));
1631 });
1632 }
1633 btree_test!(test_insert_same_key_many, insert_same_key_many);
1634
1635 fn insert_overwrite_median_key_in_full_child_node<K: TestKey, V: TestValue>() {
1636 let (key, value) = (K::build, V::build);
1637 run_btree_test(|mut btree| {
1638 for i in 1..=17 {
1639 assert_eq!(btree.insert(key(i), value(0)), None);
1640 }
1641
1642 let root = btree.load_node(btree.root_addr);
1648 assert_eq!(root.node_type(), NodeType::Internal);
1649 assert_eq!(
1650 root.entries(btree.memory()),
1651 vec![(key(6), encode(value(0)))]
1652 );
1653 assert_eq!(root.children_len(), 2);
1654
1655 let right_child = btree.load_node(root.child(1));
1657 assert!(right_child.is_full());
1658 let median_index = right_child.entries_len() / 2;
1659 let median_key = key(12);
1660 assert_eq!(right_child.key(median_index, btree.memory()), &median_key);
1661
1662 assert_eq!(btree.insert(median_key.clone(), value(123)), Some(value(0)));
1664 assert_eq!(btree.get(&median_key), Some(value(123)));
1665
1666 let right_child = btree.load_node(root.child(1));
1668 assert_eq!(right_child.node_type(), NodeType::Leaf);
1669 assert!(right_child.is_full());
1670 });
1671 }
1672 btree_test!(
1673 test_insert_overwrite_median_key_in_full_child_node,
1674 insert_overwrite_median_key_in_full_child_node
1675 );
1676
1677 fn insert_overwrite_key_in_full_root_node<K: TestKey, V: TestValue>() {
1678 let (key, value) = (K::build, V::build);
1679 run_btree_test(|mut btree| {
1680 for i in 1..=11 {
1681 assert_eq!(btree.insert(key(i), value(0)), None);
1682 }
1683
1684 let root = btree.load_node(btree.root_addr);
1688 assert!(root.is_full());
1689
1690 assert_eq!(btree.insert(key(6), value(456)), Some(value(0)));
1692
1693 let root = btree.load_node(btree.root_addr);
1694 assert_eq!(root.node_type(), NodeType::Leaf);
1695 assert_eq!(btree.get(&key(6)), Some(value(456)));
1696 assert_eq!(root.entries_len(), 11);
1697 });
1698 }
1699 btree_test!(
1700 test_insert_overwrite_key_in_full_root_node,
1701 insert_overwrite_key_in_full_root_node
1702 );
1703
1704 fn allocations_without_split<K: TestKey, V: TestValue>() {
1705 let key = K::build;
1706 run_btree_test(|mut btree| {
1707 assert_eq!(btree.allocator.num_allocated_chunks(), 0);
1708
1709 assert_eq!(btree.insert(key(1), V::empty()), None);
1710 assert_eq!(btree.allocator.num_allocated_chunks(), 1);
1711
1712 assert_eq!(btree.remove(&key(1)), Some(V::empty()));
1713 assert_eq!(btree.allocator.num_allocated_chunks(), 0);
1714 });
1715 }
1716 btree_test!(test_allocations_without_split, allocations_without_split);
1717
1718 fn allocations_with_split<K: TestKey, V: TestValue>() {
1719 let key = K::build;
1720 run_btree_test(|mut btree| {
1721 let mut i = 0;
1723 loop {
1724 assert_eq!(btree.insert(key(i), V::empty()), None);
1725 let root = btree.load_node(btree.root_addr);
1726 if root.is_full() {
1727 break;
1728 }
1729 i += 1;
1730 }
1731
1732 assert_eq!(btree.allocator.num_allocated_chunks(), 1);
1734
1735 assert_eq!(btree.insert(key(255), V::empty()), None);
1736
1737 assert_eq!(btree.allocator.num_allocated_chunks(), 3);
1739 });
1740 }
1741 btree_test!(test_allocations_with_split, allocations_with_split);
1742
1743 fn insert_split_node<K: TestKey, V: TestValue>() {
1744 let (key, value) = (K::build, V::build);
1745 run_btree_test(|mut btree| {
1746 for i in 1..=11 {
1747 assert_eq!(btree.insert(key(i), value(i)), None);
1748 }
1749
1750 let root = btree.load_node(btree.root_addr);
1752 assert!(root.is_full());
1753 assert_eq!(btree.insert(key(12), value(12)), None);
1754
1755 for i in 1..=12 {
1761 assert_eq!(btree.get(&key(i)), Some(value(i)));
1762 }
1763 });
1764 }
1765 btree_test!(test_insert_split_node, insert_split_node);
1766
1767 fn insert_split_multiple_nodes<K: TestKey, V: TestValue>() {
1768 let key = K::build;
1769 let e = |i: u32| (key(i), encode(V::empty()));
1770 run_btree_test(|mut btree| {
1771 for i in 1..=11 {
1772 assert_eq!(btree.insert(key(i), V::empty()), None);
1773 }
1774 assert_eq!(btree.insert(key(12), V::empty()), None);
1776
1777 let root = btree.load_node(btree.root_addr);
1783 assert_eq!(root.node_type(), NodeType::Internal);
1784 assert_eq!(root.entries(btree.memory()), vec![e(6)]);
1785 assert_eq!(root.children_len(), 2);
1786
1787 let child_0 = btree.load_node(root.child(0));
1788 assert_eq!(child_0.node_type(), NodeType::Leaf);
1789 assert_eq!(
1790 child_0.entries(btree.memory()),
1791 vec![e(1), e(2), e(3), e(4), e(5)]
1792 );
1793
1794 let child_1 = btree.load_node(root.child(1));
1795 assert_eq!(child_1.node_type(), NodeType::Leaf);
1796 assert_eq!(
1797 child_1.entries(btree.memory()),
1798 vec![e(7), e(8), e(9), e(10), e(11), e(12)]
1799 );
1800
1801 for i in 1..=12 {
1802 assert_eq!(btree.get(&key(i)), Some(V::empty()));
1803 }
1804
1805 assert_eq!(btree.insert(key(13), V::empty()), None);
1807 assert_eq!(btree.insert(key(14), V::empty()), None);
1808 assert_eq!(btree.insert(key(15), V::empty()), None);
1809 assert_eq!(btree.insert(key(16), V::empty()), None);
1810 assert_eq!(btree.insert(key(17), V::empty()), None);
1811 assert_eq!(btree.insert(key(18), V::empty()), None);
1813
1814 for i in 1..=18 {
1815 assert_eq!(btree.get(&key(i)), Some(V::empty()));
1816 }
1817
1818 let root = btree.load_node(btree.root_addr);
1819 assert_eq!(root.node_type(), NodeType::Internal);
1820 assert_eq!(root.entries(btree.memory()), vec![e(6), e(12)],);
1821 assert_eq!(root.children_len(), 3);
1822
1823 let child_0 = btree.load_node(root.child(0));
1824 assert_eq!(child_0.node_type(), NodeType::Leaf);
1825 assert_eq!(
1826 child_0.entries(btree.memory()),
1827 vec![e(1), e(2), e(3), e(4), e(5)]
1828 );
1829
1830 let child_1 = btree.load_node(root.child(1));
1831 assert_eq!(child_1.node_type(), NodeType::Leaf);
1832 assert_eq!(
1833 child_1.entries(btree.memory()),
1834 vec![e(7), e(8), e(9), e(10), e(11)]
1835 );
1836
1837 let child_2 = btree.load_node(root.child(2));
1838 assert_eq!(child_2.node_type(), NodeType::Leaf);
1839 assert_eq!(
1840 child_2.entries(btree.memory()),
1841 vec![e(13), e(14), e(15), e(16), e(17), e(18)]
1842 );
1843
1844 assert_eq!(btree.allocator.num_allocated_chunks(), 4);
1845 });
1846 }
1847 btree_test!(
1848 test_insert_split_multiple_nodes,
1849 insert_split_multiple_nodes
1850 );
1851
1852 fn first_key_value<K: TestKey, V: TestValue>() {
1853 let (key, value) = (K::build, V::build);
1854 run_btree_test(|mut btree: BTreeMap<K, V, _>| {
1855 assert_eq!(btree.first_key_value(), None);
1856
1857 let n = 1_000;
1858 for i in (0..n).rev() {
1859 assert_eq!(btree.insert(key(i), value(i)), None);
1860 assert_eq!(btree.first_key_value(), Some((key(i), value(i))));
1861 }
1862
1863 for i in 0..n {
1864 assert_eq!(btree.remove(&key(i)), Some(value(i)));
1865 if !btree.is_empty() {
1866 assert_eq!(btree.first_key_value(), Some((key(i + 1), value(i + 1))));
1867 }
1868 }
1869 assert_eq!(btree.first_key_value(), None);
1870 });
1871 }
1872 btree_test!(test_first_key_value, first_key_value);
1873
1874 fn last_key_value<K: TestKey, V: TestValue>() {
1875 let (key, value) = (K::build, V::build);
1876 run_btree_test(|mut btree: BTreeMap<K, V, _>| {
1877 assert_eq!(btree.last_key_value(), None);
1878
1879 let n = 1_000;
1880 for i in 0..n {
1881 assert_eq!(btree.insert(key(i), value(i)), None);
1882 assert_eq!(btree.last_key_value(), Some((key(i), value(i))));
1883 }
1884
1885 for i in (0..n).rev() {
1886 assert_eq!(btree.remove(&key(i)), Some(value(i)));
1887 if !btree.is_empty() {
1888 assert_eq!(btree.last_key_value(), Some((key(i - 1), value(i - 1))));
1889 }
1890 }
1891 assert_eq!(btree.last_key_value(), None);
1892 });
1893 }
1894 btree_test!(test_last_key_value, last_key_value);
1895
1896 fn pop_first_single_entry<K: TestKey, V: TestValue>() {
1897 let key = K::build;
1898 run_btree_test(|mut btree| {
1899 assert_eq!(btree.allocator.num_allocated_chunks(), 0);
1900
1901 assert_eq!(btree.insert(key(1), V::empty()), None);
1902 assert!(!btree.is_empty());
1903 assert_eq!(btree.allocator.num_allocated_chunks(), 1);
1904
1905 assert_eq!(btree.pop_first(), Some((key(1), V::empty())));
1906 assert!(btree.is_empty());
1907 assert_eq!(btree.allocator.num_allocated_chunks(), 0);
1908 });
1909 }
1910 btree_test!(test_pop_first_single_entry, pop_first_single_entry);
1911
1912 fn pop_last_single_entry<K: TestKey, V: TestValue>() {
1913 let (key, value) = (K::build, V::build);
1914 run_btree_test(|mut btree| {
1915 assert_eq!(btree.allocator.num_allocated_chunks(), 0);
1916
1917 assert_eq!(btree.insert(key(1), value(1)), None);
1918 assert!(!btree.is_empty());
1919 assert_eq!(btree.allocator.num_allocated_chunks(), 1);
1920
1921 assert_eq!(btree.pop_last(), Some((key(1), value(1))));
1922 assert!(btree.is_empty());
1923 assert_eq!(btree.allocator.num_allocated_chunks(), 0);
1924 });
1925 }
1926 btree_test!(test_pop_last_single_entry, pop_last_single_entry);
1927
1928 fn remove_case_2a_and_2c<K: TestKey, V: TestValue>() {
1929 let key = K::build;
1930 let e = |i: u32| (key(i), encode(V::empty()));
1931 run_btree_test(|mut btree| {
1932 for i in 1..=11 {
1933 assert_eq!(btree.insert(key(i), V::empty()), None);
1934 }
1935 assert_eq!(btree.insert(key(0), V::empty()), None);
1937
1938 for i in 0..=11 {
1944 assert_eq!(btree.get(&key(i)), Some(V::empty()));
1945 }
1946
1947 assert_eq!(btree.remove(&key(6)), Some(V::empty()));
1949
1950 let root = btree.load_node(btree.root_addr);
1955 assert_eq!(root.node_type(), NodeType::Internal);
1956 assert_eq!(root.entries(btree.memory()), vec![e(5)]);
1957 assert_eq!(root.children_len(), 2);
1958
1959 let child_0 = btree.load_node(root.child(0));
1960 assert_eq!(child_0.node_type(), NodeType::Leaf);
1961 assert_eq!(
1962 child_0.entries(btree.memory()),
1963 vec![e(0), e(1), e(2), e(3), e(4)]
1964 );
1965
1966 let child_1 = btree.load_node(root.child(1));
1967 assert_eq!(child_1.node_type(), NodeType::Leaf);
1968 assert_eq!(
1969 child_1.entries(btree.memory()),
1970 vec![e(7), e(8), e(9), e(10), e(11)]
1971 );
1972
1973 assert_eq!(btree.allocator.num_allocated_chunks(), 3);
1975
1976 assert_eq!(btree.remove(&key(5)), Some(V::empty()));
1978
1979 let btree = BTreeMap::<K, V, _>::load(btree.into_memory());
1981
1982 let root = btree.load_node(btree.root_addr);
1985 assert_eq!(
1986 root.entries(btree.memory()),
1987 vec![e(0), e(1), e(2), e(3), e(4), e(7), e(8), e(9), e(10), e(11)]
1988 );
1989
1990 assert_eq!(btree.allocator.num_allocated_chunks(), 1);
1992 });
1993 }
1994 btree_test!(test_remove_case_2a_and_2c, remove_case_2a_and_2c);
1995
1996 fn remove_case_2b<K: TestKey, V: TestValue>() {
1997 let key = K::build;
1998 let e = |i: u32| (key(i), encode(V::empty()));
1999 run_btree_test(|mut btree| {
2000 for i in 1..=11 {
2001 assert_eq!(btree.insert(key(i), V::empty()), None);
2002 }
2003 assert_eq!(btree.insert(key(12), V::empty()), None);
2005
2006 for i in 1..=12 {
2012 assert_eq!(btree.get(&key(i)), Some(V::empty()));
2013 }
2014
2015 assert_eq!(btree.remove(&key(6)), Some(V::empty()));
2017
2018 let root = btree.load_node(btree.root_addr);
2023 assert_eq!(root.node_type(), NodeType::Internal);
2024 assert_eq!(root.entries(btree.memory()), vec![e(7)]);
2025 assert_eq!(root.children_len(), 2);
2026
2027 let child_0 = btree.load_node(root.child(0));
2028 assert_eq!(child_0.node_type(), NodeType::Leaf);
2029 assert_eq!(
2030 child_0.entries(btree.memory()),
2031 vec![e(1), e(2), e(3), e(4), e(5)]
2032 );
2033
2034 let child_1 = btree.load_node(root.child(1));
2035 assert_eq!(child_1.node_type(), NodeType::Leaf);
2036 assert_eq!(
2037 child_1.entries(btree.memory()),
2038 vec![e(8), e(9), e(10), e(11), e(12)]
2039 );
2040
2041 assert_eq!(btree.remove(&key(7)), Some(V::empty()));
2043 let root = btree.load_node(btree.root_addr);
2047 assert_eq!(root.node_type(), NodeType::Leaf);
2048 assert_eq!(
2049 root.entries(btree.memory()),
2050 collect([1, 2, 3, 4, 5, 8, 9, 10, 11, 12].into_iter().map(e))
2051 );
2052
2053 assert_eq!(btree.allocator.num_allocated_chunks(), 1);
2054 });
2055 }
2056 btree_test!(test_remove_case_2b, remove_case_2b);
2057
2058 fn remove_case_3a_right<K: TestKey, V: TestValue>() {
2059 let key = K::build;
2060 let e = |i: u32| (key(i), encode(V::empty()));
2061 run_btree_test(|mut btree| {
2062 for i in 1..=11 {
2063 assert_eq!(btree.insert(key(i), V::empty()), None);
2064 }
2065
2066 assert_eq!(btree.insert(key(12), V::empty()), None);
2068
2069 assert_eq!(btree.remove(&key(3)), Some(V::empty()));
2076
2077 let root = btree.load_node(btree.root_addr);
2082 assert_eq!(root.node_type(), NodeType::Internal);
2083 assert_eq!(root.entries(btree.memory()), vec![e(7)]);
2084 assert_eq!(root.children_len(), 2);
2085
2086 let child_0 = btree.load_node(root.child(0));
2087 assert_eq!(child_0.node_type(), NodeType::Leaf);
2088 assert_eq!(
2089 child_0.entries(btree.memory()),
2090 vec![e(1), e(2), e(4), e(5), e(6)]
2091 );
2092
2093 let child_1 = btree.load_node(root.child(1));
2094 assert_eq!(child_1.node_type(), NodeType::Leaf);
2095 assert_eq!(
2096 child_1.entries(btree.memory()),
2097 vec![e(8), e(9), e(10), e(11), e(12)]
2098 );
2099
2100 assert_eq!(btree.allocator.num_allocated_chunks(), 3);
2102 });
2103 }
2104 btree_test!(test_remove_case_3a_right, remove_case_3a_right);
2105
2106 fn remove_case_3a_left<K: TestKey, V: TestValue>() {
2107 let key = K::build;
2108 let e = |i: u32| (key(i), encode(V::empty()));
2109 run_btree_test(|mut btree| {
2110 for i in 1..=11 {
2111 assert_eq!(btree.insert(key(i), V::empty()), None);
2112 }
2113 assert_eq!(btree.insert(key(0), V::empty()), None);
2115
2116 assert_eq!(btree.remove(&key(8)), Some(V::empty()));
2123
2124 let root = btree.load_node(btree.root_addr);
2129 assert_eq!(root.node_type(), NodeType::Internal);
2130 assert_eq!(root.entries(btree.memory()), vec![e(5)]);
2131 assert_eq!(root.children_len(), 2);
2132
2133 let child_0 = btree.load_node(root.child(0));
2134 assert_eq!(child_0.node_type(), NodeType::Leaf);
2135 assert_eq!(
2136 child_0.entries(btree.memory()),
2137 vec![e(0), e(1), e(2), e(3), e(4)]
2138 );
2139
2140 let child_1 = btree.load_node(root.child(1));
2141 assert_eq!(child_1.node_type(), NodeType::Leaf);
2142 assert_eq!(
2143 child_1.entries(btree.memory()),
2144 vec![e(6), e(7), e(9), e(10), e(11)]
2145 );
2146
2147 assert_eq!(btree.allocator.num_allocated_chunks(), 3);
2149 });
2150 }
2151 btree_test!(test_remove_case_3a_left, remove_case_3a_left);
2152
2153 fn remove_case_3b_merge_into_right<K: TestKey, V: TestValue>() {
2154 let key = K::build;
2155 let e = |i: u32| (key(i), encode(V::empty()));
2156 run_btree_test(|mut btree| {
2157 for i in 1..=11 {
2158 assert_eq!(btree.insert(key(i), V::empty()), None);
2159 }
2160 assert_eq!(btree.insert(key(12), V::empty()), None);
2162
2163 for i in 1..=12 {
2169 assert_eq!(btree.get(&key(i)), Some(V::empty()));
2170 }
2171
2172 assert_eq!(btree.remove(&key(6)), Some(V::empty()));
2174 let root = btree.load_node(btree.root_addr);
2179 assert_eq!(root.node_type(), NodeType::Internal);
2180 assert_eq!(root.entries(btree.memory()), vec![e(7)]);
2181 assert_eq!(root.children_len(), 2);
2182
2183 let child_0 = btree.load_node(root.child(0));
2184 assert_eq!(child_0.node_type(), NodeType::Leaf);
2185 assert_eq!(
2186 child_0.entries(btree.memory()),
2187 vec![e(1), e(2), e(3), e(4), e(5)]
2188 );
2189
2190 let child_1 = btree.load_node(root.child(1));
2191 assert_eq!(child_1.node_type(), NodeType::Leaf);
2192 assert_eq!(
2193 child_1.entries(btree.memory()),
2194 vec![e(8), e(9), e(10), e(11), e(12)]
2195 );
2196
2197 assert_eq!(btree.allocator.num_allocated_chunks(), 3);
2199
2200 assert_eq!(btree.remove(&key(3)), Some(V::empty()));
2202
2203 let btree = BTreeMap::<K, V, _>::load(btree.into_memory());
2205
2206 let root = btree.load_node(btree.root_addr);
2210 assert_eq!(root.node_type(), NodeType::Leaf);
2211 assert_eq!(
2212 root.entries(btree.memory()),
2213 collect([1, 2, 4, 5, 7, 8, 9, 10, 11, 12].into_iter().map(e))
2214 );
2215
2216 assert_eq!(btree.allocator.num_allocated_chunks(), 1);
2218 });
2219 }
2220 btree_test!(
2221 test_remove_case_3b_merge_into_right,
2222 remove_case_3b_merge_into_right
2223 );
2224
2225 fn remove_case_3b_merge_into_left<K: TestKey, V: TestValue>() {
2226 let key = K::build;
2227 let e = |i: u32| (key(i), encode(V::empty()));
2228 run_btree_test(|mut btree| {
2229 for i in 1..=11 {
2230 assert_eq!(btree.insert(key(i), V::empty()), None);
2231 }
2232
2233 assert_eq!(btree.insert(key(12), V::empty()), None);
2235
2236 for i in 1..=12 {
2242 assert_eq!(btree.get(&key(i)), Some(V::empty()));
2243 }
2244
2245 assert_eq!(btree.remove(&key(6)), Some(V::empty()));
2247
2248 let root = btree.load_node(btree.root_addr);
2253 assert_eq!(root.node_type(), NodeType::Internal);
2254 assert_eq!(root.entries(btree.memory()), vec![e(7)]);
2255 assert_eq!(root.children_len(), 2);
2256
2257 let child_0 = btree.load_node(root.child(0));
2258 assert_eq!(child_0.node_type(), NodeType::Leaf);
2259 assert_eq!(
2260 child_0.entries(btree.memory()),
2261 vec![e(1), e(2), e(3), e(4), e(5)]
2262 );
2263
2264 let child_1 = btree.load_node(root.child(1));
2265 assert_eq!(child_1.node_type(), NodeType::Leaf);
2266 assert_eq!(
2267 child_1.entries(btree.memory()),
2268 vec![e(8), e(9), e(10), e(11), e(12)]
2269 );
2270
2271 assert_eq!(btree.allocator.num_allocated_chunks(), 3);
2273
2274 assert_eq!(btree.remove(&key(10)), Some(V::empty()));
2276
2277 let btree = BTreeMap::<K, V, _>::load(btree.into_memory());
2279
2280 let root = btree.load_node(btree.root_addr);
2284 assert_eq!(root.node_type(), NodeType::Leaf);
2285 assert_eq!(
2286 root.entries(btree.memory()),
2287 vec![e(1), e(2), e(3), e(4), e(5), e(7), e(8), e(9), e(11), e(12)]
2288 );
2289
2290 assert_eq!(btree.allocator.num_allocated_chunks(), 1);
2292 });
2293 }
2294 btree_test!(
2295 test_remove_case_3b_merge_into_left,
2296 remove_case_3b_merge_into_left
2297 );
2298
2299 fn insert_remove_many<K: TestKey, V: TestValue>() {
2300 let (key, value) = (K::build, V::build);
2301 run_btree_test(|mut btree| {
2302 let n = 10_000;
2303 for i in 0..n {
2304 assert_eq!(btree.insert(key(i), value(i)), None);
2305 assert_eq!(btree.get(&key(i)), Some(value(i)));
2306 }
2307
2308 let mut btree = BTreeMap::<K, V, _>::load(btree.into_memory());
2309 for i in 0..n {
2310 assert_eq!(btree.remove(&key(i)), Some(value(i)));
2311 assert_eq!(btree.get(&key(i)), None);
2312 }
2313
2314 assert_eq!(btree.allocator.num_allocated_chunks(), 0);
2316 });
2317 }
2318 btree_test!(test_insert_remove_many, insert_remove_many);
2319
2320 fn pop_first_many<K: TestKey, V: TestValue>() {
2321 let (key, value) = (K::build, V::build);
2322 run_btree_test(|mut btree| {
2323 let n = 10_000;
2324
2325 assert!(btree.is_empty());
2326 assert_eq!(btree.pop_first(), None);
2327
2328 for i in 0..n {
2329 assert_eq!(btree.insert(key(i), value(i)), None);
2330 }
2331 assert_eq!(btree.len(), n as u64);
2332
2333 let mut btree = BTreeMap::<K, V, _>::load(btree.into_memory());
2334 for i in 0..n {
2335 assert_eq!(btree.pop_first(), Some((key(i), value(i))));
2336 assert_eq!(btree.get(&key(i)), None);
2337 }
2338
2339 assert!(btree.is_empty());
2341 assert_eq!(btree.allocator.num_allocated_chunks(), 0);
2342 });
2343 }
2344 btree_test!(test_pop_first_many, pop_first_many);
2345
2346 fn pop_last_many<K: TestKey, V: TestValue>() {
2347 let (key, value) = (K::build, V::build);
2348 run_btree_test(|mut btree| {
2349 let n = 10_000;
2350
2351 assert!(btree.is_empty());
2352 assert_eq!(btree.pop_last(), None);
2353
2354 for i in 0..n {
2355 assert_eq!(btree.insert(key(i), value(i)), None);
2356 }
2357 assert_eq!(btree.len(), n as u64);
2358
2359 let mut btree = BTreeMap::<K, V, _>::load(btree.into_memory());
2360 for i in (0..n).rev() {
2361 assert_eq!(btree.pop_last(), Some((key(i), value(i))));
2362 assert_eq!(btree.get(&key(i)), None);
2363 }
2364
2365 assert!(btree.is_empty());
2367 assert_eq!(btree.allocator.num_allocated_chunks(), 0);
2368 });
2369 }
2370 btree_test!(test_pop_last_many, pop_last_many);
2371
2372 fn reloading<K: TestKey, V: TestValue>() {
2373 let (key, value) = (K::build, V::build);
2374 run_btree_test(|mut btree| {
2375 let n = 1_000;
2376 for i in 0..n {
2377 assert_eq!(btree.insert(key(i), value(i)), None);
2378 assert_eq!(btree.get(&key(i)), Some(value(i)));
2379 }
2380 assert_eq!(btree.len(), n as u64);
2381
2382 let mut btree = BTreeMap::<K, V, VectorMemory>::load(btree.into_memory());
2384 assert_eq!(btree.len(), n as u64);
2385 for i in 0..n {
2386 assert_eq!(btree.get(&key(i)), Some(value(i)));
2387 }
2388
2389 for i in 0..n {
2391 assert_eq!(btree.remove(&key(i)), Some(value(i)));
2392 }
2393 assert_eq!(btree.len(), 0);
2394
2395 let btree = BTreeMap::<K, V, VectorMemory>::load(btree.into_memory());
2397 assert_eq!(btree.len(), 0);
2398 });
2399 }
2400 btree_test!(test_reloading, reloading);
2401
2402 fn len<K: TestKey, V: TestValue>() {
2403 let (key, value) = (K::build, V::build);
2404 run_btree_test(|mut btree| {
2405 let n = 1_000;
2406 for i in 0..n {
2407 assert_eq!(btree.insert(key(i), value(i)), None);
2408 }
2409
2410 assert_eq!(btree.len(), n as u64);
2411 assert!(!btree.is_empty());
2412
2413 for i in 0..n {
2414 assert_eq!(btree.remove(&key(i)), Some(value(i)));
2415 }
2416
2417 assert_eq!(btree.len(), 0);
2418 assert!(btree.is_empty());
2419 });
2420 }
2421 btree_test!(test_len, len);
2422
2423 fn contains_key<K: TestKey, V: TestValue>() {
2424 let (key, value) = (K::build, V::build);
2425 run_btree_test(|mut btree| {
2426 let n = 1_000;
2427 for i in (0..n).step_by(2) {
2428 assert_eq!(btree.insert(key(i), value(i)), None);
2429 }
2430
2431 for i in 0..n {
2433 assert_eq!(btree.contains_key(&key(i)), i % 2 == 0);
2434 }
2435 });
2436 }
2437 btree_test!(test_contains_key, contains_key);
2438
2439 fn range_empty<K: TestKey, V: TestValue>() {
2440 let key = K::build;
2441 run_btree_test(|btree: BTreeMap<K, V, _>| {
2442 assert_eq!(collect_entry(btree.range(key(0)..)), vec![]);
2444 assert_eq!(collect_entry(btree.range(key(10)..)), vec![]);
2445 assert_eq!(collect_entry(btree.range(key(100)..)), vec![]);
2446 });
2447 }
2448 btree_test!(test_range_empty, range_empty);
2449
2450 fn range_leaf_prefix_greater_than_all_entries<K: TestKey, V: TestValue>() {
2452 let (key, value) = (K::build, V::build);
2453 run_btree_test(|mut btree| {
2454 btree.insert(key(0), value(0));
2455
2456 assert_eq!(collect_entry(btree.range(key(1)..)), vec![]);
2458 });
2459 }
2460 btree_test!(
2461 test_range_leaf_prefix_greater_than_all_entries,
2462 range_leaf_prefix_greater_than_all_entries
2463 );
2464
2465 fn range_internal_prefix_greater_than_all_entries<K: TestKey, V: TestValue>() {
2467 let (key, value) = (K::build, V::build);
2468 run_btree_test(|mut btree| {
2469 for i in 1..=12 {
2470 assert_eq!(btree.insert(key(i), value(i)), None);
2471 }
2472
2473 assert_eq!(
2480 collect_entry(btree.range(key(7)..key(8))),
2481 vec![(key(7), value(7))]
2482 );
2483 });
2484 }
2485 btree_test!(
2486 test_range_internal_prefix_greater_than_all_entries,
2487 range_internal_prefix_greater_than_all_entries
2488 );
2489
2490 fn range_various_prefixes<K: TestKey, V: TestValue>() {
2491 let (key, value) = (K::build, V::build);
2492 run_btree_test(|mut btree| {
2493 btree.insert(key(1), value(100));
2494 btree.insert(key(2), value(200));
2495 btree.insert(key(3), value(300));
2496 btree.insert(key(4), value(400));
2497 btree.insert(key(11), value(500));
2498 btree.insert(key(12), value(600));
2499 btree.insert(key(13), value(700));
2500 btree.insert(key(14), value(800));
2501 btree.insert(key(21), value(900));
2502 btree.insert(key(22), value(1_000));
2503 btree.insert(key(23), value(1_100));
2504 btree.insert(key(24), value(1_200));
2505
2506 let root = btree.load_node(btree.root_addr);
2512 assert_eq!(root.node_type(), NodeType::Internal);
2513 assert_eq!(
2514 root.entries(btree.memory()),
2515 vec![(key(12), encode(value(600)))]
2516 );
2517 assert_eq!(root.children_len(), 2);
2518
2519 assert_eq!(
2521 collect_entry(btree.range(key(0)..key(11))),
2522 vec![
2523 (key(1), value(100)),
2524 (key(2), value(200)),
2525 (key(3), value(300)),
2526 (key(4), value(400)),
2527 ]
2528 );
2529
2530 assert_eq!(
2532 collect_entry(btree.range(key(10)..key(20))),
2533 vec![
2534 (key(11), value(500)),
2535 (key(12), value(600)),
2536 (key(13), value(700)),
2537 (key(14), value(800)),
2538 ]
2539 );
2540
2541 assert_eq!(
2543 collect_entry(btree.range(key(20)..key(30))),
2544 vec![
2545 (key(21), value(900)),
2546 (key(22), value(1_000)),
2547 (key(23), value(1_100)),
2548 (key(24), value(1_200)),
2549 ]
2550 );
2551 });
2552 }
2553 btree_test!(test_range_various_prefixes, range_various_prefixes);
2554
2555 fn range_various_prefixes_2<K: TestKey, V: TestValue>() {
2556 let (key, value) = (K::build, V::build);
2557 run_btree_test(|mut btree| {
2558 btree.insert(key(1), value(100));
2559 btree.insert(key(2), value(200));
2560 btree.insert(key(3), value(300));
2561 btree.insert(key(4), value(400));
2562 btree.insert(key(12), value(500));
2563 btree.insert(key(14), value(600));
2564 btree.insert(key(16), value(700));
2565 btree.insert(key(18), value(800));
2566 btree.insert(key(19), value(900));
2567 btree.insert(key(21), value(1000));
2568 btree.insert(key(22), value(1100));
2569 btree.insert(key(23), value(1200));
2570 btree.insert(key(24), value(1300));
2571 btree.insert(key(25), value(1400));
2572 btree.insert(key(26), value(1500));
2573 btree.insert(key(27), value(1600));
2574 btree.insert(key(28), value(1700));
2575 btree.insert(key(29), value(1800));
2576
2577 let root = btree.load_node(btree.root_addr);
2584 assert_eq!(root.node_type(), NodeType::Internal);
2585 assert_eq!(
2586 root.entries(btree.memory()),
2587 vec![
2588 (key(14), encode(value(600))),
2589 (key(23), encode(value(1200)))
2590 ]
2591 );
2592 assert_eq!(root.children_len(), 3);
2593 let child_0 = btree.load_node(root.child(0));
2594 assert_eq!(child_0.node_type(), NodeType::Leaf);
2595 assert_eq!(
2596 child_0.entries(btree.memory()),
2597 vec![
2598 (key(1), encode(value(100))),
2599 (key(2), encode(value(200))),
2600 (key(3), encode(value(300))),
2601 (key(4), encode(value(400))),
2602 (key(12), encode(value(500))),
2603 ]
2604 );
2605
2606 let child_1 = btree.load_node(root.child(1));
2607 assert_eq!(child_1.node_type(), NodeType::Leaf);
2608 assert_eq!(
2609 child_1.entries(btree.memory()),
2610 vec![
2611 (key(16), encode(value(700))),
2612 (key(18), encode(value(800))),
2613 (key(19), encode(value(900))),
2614 (key(21), encode(value(1000))),
2615 (key(22), encode(value(1100))),
2616 ]
2617 );
2618
2619 let child_2 = btree.load_node(root.child(2));
2620 assert_eq!(
2621 child_2.entries(btree.memory()),
2622 vec![
2623 (key(24), encode(value(1300))),
2624 (key(25), encode(value(1400))),
2625 (key(26), encode(value(1500))),
2626 (key(27), encode(value(1600))),
2627 (key(28), encode(value(1700))),
2628 (key(29), encode(value(1800))),
2629 ]
2630 );
2631
2632 assert_eq!(collect_entry(btree.range(key(15)..key(16))), vec![]);
2634
2635 assert_eq!(
2637 collect_entry(btree.range(key(15)..=key(26))),
2638 vec![
2639 (key(16), value(700)),
2640 (key(18), value(800)),
2641 (key(19), value(900)),
2642 (key(21), value(1000)),
2643 (key(22), value(1100)),
2644 (key(23), value(1200)),
2645 (key(24), value(1300)),
2646 (key(25), value(1400)),
2647 (key(26), value(1500)),
2648 ]
2649 );
2650
2651 assert_eq!(
2653 collect_entry(btree.range(key(10)..key(20))),
2654 vec![
2655 (key(12), value(500)),
2656 (key(14), value(600)),
2657 (key(16), value(700)),
2658 (key(18), value(800)),
2659 (key(19), value(900)),
2660 ]
2661 );
2662
2663 assert_eq!(
2666 collect_entry(btree.range(key(20)..key(30))),
2667 vec![
2668 (key(21), value(1000)),
2669 (key(22), value(1100)),
2670 (key(23), value(1200)),
2671 (key(24), value(1300)),
2672 (key(25), value(1400)),
2673 (key(26), value(1500)),
2674 (key(27), value(1600)),
2675 (key(28), value(1700)),
2676 (key(29), value(1800)),
2677 ]
2678 );
2679 });
2680 }
2681 btree_test!(test_range_various_prefixes_2, range_various_prefixes_2);
2682
2683 fn range_large<K: TestKey, V: TestValue>() {
2684 let (key, value) = (K::build, V::build);
2685 run_btree_test(|mut btree| {
2686 const TOTAL: u32 = 2_000;
2687 const MID: u32 = TOTAL / 2;
2688
2689 for i in 0..TOTAL {
2690 assert_eq!(btree.insert(key(i), value(i)), None);
2691 }
2692
2693 let keys: Vec<_> = btree.range(key(0)..key(MID)).collect();
2695 assert_eq!(keys.len(), MID as usize);
2696 for (i, entry) in keys.into_iter().enumerate() {
2697 let j = i as u32;
2698 assert_eq!(*entry.key(), key(j));
2699 assert_eq!(entry.value(), value(j));
2700 }
2701
2702 let keys: Vec<_> = btree.range(key(MID)..).collect();
2704 assert_eq!(keys.len(), (TOTAL - MID) as usize);
2705 for (i, entry) in keys.into_iter().enumerate() {
2706 let j = MID + i as u32;
2707 assert_eq!(*entry.key(), key(j));
2708 assert_eq!(entry.value(), value(j));
2709 }
2710 });
2711 }
2712 btree_test!(test_range_large, range_large);
2713
2714 fn range_various_prefixes_with_offset<K: TestKey, V: TestValue>() {
2715 let (key, value) = (K::build, V::build);
2716 run_btree_test(|mut btree| {
2717 btree.insert(key(1), value(100));
2718 btree.insert(key(2), value(200));
2719 btree.insert(key(3), value(300));
2720 btree.insert(key(4), value(400));
2721 btree.insert(key(11), value(500));
2722 btree.insert(key(12), value(600));
2723 btree.insert(key(13), value(700));
2724 btree.insert(key(14), value(800));
2725 btree.insert(key(21), value(900));
2726 btree.insert(key(22), value(1000));
2727 btree.insert(key(23), value(1100));
2728 btree.insert(key(24), value(1200));
2729
2730 let root = btree.load_node(btree.root_addr);
2736 assert_eq!(root.node_type(), NodeType::Internal);
2737 assert_eq!(
2738 root.entries(btree.memory()),
2739 vec![(key(12), encode(value(600)))]
2740 );
2741 assert_eq!(root.children_len(), 2);
2742
2743 assert_eq!(
2744 collect_entry(btree.range(key(0)..key(10))),
2745 vec![
2746 (key(1), value(100)),
2747 (key(2), value(200)),
2748 (key(3), value(300)),
2749 (key(4), value(400)),
2750 ]
2751 );
2752
2753 assert_eq!(
2755 collect_entry(btree.range(key(13)..key(20))),
2756 vec![(key(13), value(700)), (key(14), value(800))]
2757 );
2758
2759 assert_eq!(collect_entry(btree.range(key(25)..)), vec![]);
2761 });
2762 }
2763 btree_test!(
2764 test_range_various_prefixes_with_offset,
2765 range_various_prefixes_with_offset
2766 );
2767
2768 fn range_various_prefixes_with_offset_2<K: TestKey, V: TestValue>() {
2769 let (key, value) = (K::build, V::build);
2770 run_btree_test(|mut btree| {
2771 btree.insert(key(1), value(0));
2772 btree.insert(key(2), value(0));
2773 btree.insert(key(3), value(0));
2774 btree.insert(key(4), value(0));
2775 btree.insert(key(12), value(0));
2776 btree.insert(key(14), value(0));
2777 btree.insert(key(16), value(0));
2778 btree.insert(key(18), value(0));
2779 btree.insert(key(19), value(0));
2780 btree.insert(key(21), value(0));
2781 btree.insert(key(22), value(0));
2782 btree.insert(key(23), value(0));
2783 btree.insert(key(24), value(0));
2784 btree.insert(key(25), value(0));
2785 btree.insert(key(26), value(0));
2786 btree.insert(key(27), value(0));
2787 btree.insert(key(28), value(0));
2788 btree.insert(key(29), value(0));
2789
2790 let root = btree.load_node(btree.root_addr);
2797 assert_eq!(root.node_type(), NodeType::Internal);
2798 assert_eq!(
2799 root.entries(btree.memory()),
2800 vec![(key(14), encode(value(0))), (key(23), encode(value(0)))]
2801 );
2802 assert_eq!(root.children_len(), 3);
2803
2804 let child_0 = btree.load_node(root.child(0));
2805 assert_eq!(child_0.node_type(), NodeType::Leaf);
2806 assert_eq!(
2807 child_0.entries(btree.memory()),
2808 vec![
2809 (key(1), encode(value(0))),
2810 (key(2), encode(value(0))),
2811 (key(3), encode(value(0))),
2812 (key(4), encode(value(0))),
2813 (key(12), encode(value(0))),
2814 ]
2815 );
2816
2817 let child_1 = btree.load_node(root.child(1));
2818 assert_eq!(child_1.node_type(), NodeType::Leaf);
2819 assert_eq!(
2820 child_1.entries(btree.memory()),
2821 vec![
2822 (key(16), encode(value(0))),
2823 (key(18), encode(value(0))),
2824 (key(19), encode(value(0))),
2825 (key(21), encode(value(0))),
2826 (key(22), encode(value(0))),
2827 ]
2828 );
2829
2830 let child_2 = btree.load_node(root.child(2));
2831 assert_eq!(
2832 child_2.entries(btree.memory()),
2833 vec![
2834 (key(24), encode(value(0))),
2835 (key(25), encode(value(0))),
2836 (key(26), encode(value(0))),
2837 (key(27), encode(value(0))),
2838 (key(28), encode(value(0))),
2839 (key(29), encode(value(0))),
2840 ]
2841 );
2842
2843 assert_eq!(
2845 collect_entry(btree.range(key(14)..key(20))),
2846 vec![
2847 (key(14), value(0)),
2848 (key(16), value(0)),
2849 (key(18), value(0)),
2850 (key(19), value(0)),
2851 ]
2852 );
2853
2854 assert_eq!(
2857 collect_entry(btree.range(key(22)..key(30))),
2858 vec![
2859 (key(22), value(0)),
2860 (key(23), value(0)),
2861 (key(24), value(0)),
2862 (key(25), value(0)),
2863 (key(26), value(0)),
2864 (key(27), value(0)),
2865 (key(28), value(0)),
2866 (key(29), value(0)),
2867 ]
2868 );
2869 });
2870 }
2871 btree_test!(
2872 test_range_various_prefixes_with_offset_2,
2873 range_various_prefixes_with_offset_2
2874 );
2875
2876 #[test]
2877 #[should_panic(expected = "max_key_size must be <= 4")]
2878 fn v1_rejects_increases_in_max_key_size() {
2879 let mem = make_memory();
2880 let btree: BTreeMap<Blob<4>, Blob<3>, _> = BTreeMap::init_v1(mem);
2881 let _btree: BTreeMap<Blob<5>, Blob<3>, _> = BTreeMap::init_v1(btree.into_memory());
2882 }
2883
2884 #[test]
2885 fn v2_handles_increases_in_max_key_size_and_max_value_size() {
2886 let mem = make_memory();
2887 let mut btree: BTreeMap<Blob<4>, Blob<4>, _> = BTreeMap::init(mem);
2888 btree.insert(
2889 [1u8; 4].as_slice().try_into().unwrap(),
2890 [1u8; 4].as_slice().try_into().unwrap(),
2891 );
2892
2893 let mut btree: BTreeMap<Blob<5>, Blob<5>, _> = BTreeMap::init(btree.into_memory());
2895 btree.insert(
2896 [2u8; 5].as_slice().try_into().unwrap(),
2897 [2u8; 5].as_slice().try_into().unwrap(),
2898 );
2899
2900 assert_eq!(
2902 btree.get(&([1u8; 4].as_slice().try_into().unwrap())),
2903 Some([1u8; 4].as_slice().try_into().unwrap())
2904 );
2905
2906 assert_eq!(
2907 btree.get(&([2u8; 5].as_slice().try_into().unwrap())),
2908 Some([2u8; 5].as_slice().try_into().unwrap())
2909 );
2910 }
2911
2912 #[test]
2913 fn accepts_small_or_equal_key_sizes() {
2914 let mem = make_memory();
2915 let btree: BTreeMap<Blob<4>, Blob<3>, _> = BTreeMap::init(mem);
2916 let btree: BTreeMap<Blob<3>, Blob<3>, _> = BTreeMap::init(btree.into_memory());
2918 let _btree: BTreeMap<Blob<4>, Blob<3>, _> = BTreeMap::init(btree.into_memory());
2920 }
2921
2922 #[test]
2923 #[should_panic(expected = "max_value_size must be <= 3")]
2924 fn v1_rejects_larger_value_sizes() {
2925 let mem = make_memory();
2926 let btree: BTreeMap<Blob<4>, Blob<3>, _> = BTreeMap::init_v1(mem);
2927 let _btree: BTreeMap<Blob<4>, Blob<4>, _> = BTreeMap::init_v1(btree.into_memory());
2928 }
2929
2930 #[test]
2931 fn accepts_small_or_equal_value_sizes() {
2932 let mem = make_memory();
2933 let btree: BTreeMap<Blob<4>, Blob<3>, _> = BTreeMap::init(mem);
2934 let btree: BTreeMap<Blob<4>, Blob<2>, _> = BTreeMap::init(btree.into_memory());
2936 let _btree: BTreeMap<Blob<4>, Blob<3>, _> = BTreeMap::init(btree.into_memory());
2938 }
2939
2940 fn bruteforce_range_search<K: TestKey, V: TestValue>() {
2941 let (key, value) = (K::build, V::build);
2942 run_btree_test(|mut stable_map| {
2943 use std::collections::BTreeMap;
2944 const NKEYS: u32 = 60;
2945 let mut std_map = BTreeMap::new();
2946
2947 for i in 0..NKEYS {
2948 std_map.insert(key(i), value(i));
2949 stable_map.insert(key(i), value(i));
2950 }
2951
2952 assert_eq!(
2953 collect_kv(std_map.range(..)),
2954 collect_entry(stable_map.range(..))
2955 );
2956
2957 for l in 0..=NKEYS {
2958 assert_eq!(
2959 collect_kv(std_map.range(key(l)..)),
2960 collect_entry(stable_map.range(key(l)..))
2961 );
2962
2963 assert_eq!(
2964 collect_kv(std_map.range(..key(l))),
2965 collect_entry(stable_map.range(..key(l)))
2966 );
2967
2968 assert_eq!(
2969 collect_kv(std_map.range(..=key(l))),
2970 collect_entry(stable_map.range(..=key(l)))
2971 );
2972
2973 for r in l + 1..=NKEYS {
2974 for lbound in [Bound::Included(key(l)), Bound::Excluded(key(l))] {
2975 for rbound in [Bound::Included(key(r)), Bound::Excluded(key(r))] {
2976 let range = (lbound.clone(), rbound);
2977 assert_eq!(
2978 collect_kv(std_map.range(range.clone())),
2979 collect_entry(stable_map.range(range.clone())),
2980 "range: {range:?}"
2981 );
2982 }
2983 }
2984 }
2985 }
2986 });
2987 }
2988 btree_test!(test_bruteforce_range_search, bruteforce_range_search);
2989
2990 fn test_iter_from_prev_key<K: TestKey, V: TestValue>() {
2991 let (key, value) = (K::build, V::build);
2992 run_btree_test(|mut btree| {
2993 for j in 0..100 {
2994 btree.insert(key(j), value(j));
2995 for i in 0..=j {
2996 assert_eq!(
2997 btree
2998 .iter_from_prev_key(&key(i + 1))
2999 .next()
3000 .map(|entry| entry.into_pair()),
3001 Some((key(i), value(i))),
3002 "failed to get an upper bound for key({})",
3003 i + 1
3004 );
3005 }
3006 assert_eq!(
3007 btree
3008 .iter_from_prev_key(&key(0))
3009 .next()
3010 .map(|entry| entry.into_pair()),
3011 None,
3012 "key(0) must not have an upper bound"
3013 );
3014 }
3015 });
3016 }
3017 btree_test!(test_test_iter_from_prev_key, test_iter_from_prev_key);
3018
3019 #[derive(Clone, Ord, PartialOrd, Eq, PartialEq)]
3021 struct BuggyStruct;
3022 impl crate::Storable for BuggyStruct {
3023 fn to_bytes(&self) -> Cow<'_, [u8]> {
3024 Cow::Borrowed(&[1, 2, 3, 4])
3025 }
3026
3027 fn into_bytes(self) -> Vec<u8> {
3028 self.to_bytes().into_owned()
3029 }
3030
3031 fn from_bytes(_: Cow<[u8]>) -> Self {
3032 unimplemented!();
3033 }
3034
3035 const BOUND: StorableBound = StorableBound::Bounded {
3036 max_size: 1,
3037 is_fixed_size: false,
3038 };
3039 }
3040
3041 #[test]
3042 #[should_panic(expected = "expected an element with length <= 1 bytes, but found 4")]
3043 fn v1_panics_if_key_is_bigger_than_max_size() {
3044 let mut btree = BTreeMap::init_v1(make_memory());
3045 btree.insert(BuggyStruct, ());
3046 }
3047
3048 #[test]
3049 #[should_panic(expected = "expected an element with length <= 1 bytes, but found 4")]
3050 fn v2_panics_if_key_is_bigger_than_max_size() {
3051 let mut btree = BTreeMap::init(make_memory());
3052 btree.insert(BuggyStruct, ());
3053 }
3054
3055 #[test]
3056 #[should_panic(expected = "expected an element with length <= 1 bytes, but found 4")]
3057 fn v1_panics_if_value_is_bigger_than_max_size() {
3058 let mut btree = BTreeMap::init(make_memory());
3059 btree.insert((), BuggyStruct);
3060 }
3061
3062 #[test]
3063 #[should_panic(expected = "expected an element with length <= 1 bytes, but found 4")]
3064 fn v2_panics_if_value_is_bigger_than_max_size() {
3065 let mut btree = BTreeMap::init(make_memory());
3066 btree.insert((), BuggyStruct);
3067 }
3068
3069 #[test]
3072 #[ignore]
3073 fn create_btreemap_dump_file() {
3074 let mem = make_memory();
3075 let mut btree = BTreeMap::init_v1(mem.clone());
3076
3077 let key = b(&[1, 2, 3]);
3078 let value = b(&[4, 5, 6]);
3079 assert_eq!(btree.insert(key, value), None);
3080 assert_eq!(btree.get(&key), Some(value));
3081
3082 use std::io::prelude::*;
3083 let mut file =
3084 std::fs::File::create(format!("dumps/btreemap_v{LAYOUT_VERSION}.dump")).unwrap();
3085 file.write_all(&mem.borrow()).unwrap();
3086 }
3087
3088 #[test]
3089 fn produces_layout_identical_to_layout_version_1_with_packed_headers() {
3090 let mem = make_memory();
3091 let mut btree = BTreeMap::init_v1(mem.clone());
3092
3093 let key = b(&[1, 2, 3]);
3094 let value = b(&[4, 5, 6]);
3095 assert_eq!(btree.insert(key, value), None);
3096 assert_eq!(btree.get(&key), Some(value));
3097
3098 let btreemap_v1 = include_bytes!("../dumps/btreemap_v1_packed_headers.dump");
3099 assert_eq!(*mem.borrow(), btreemap_v1);
3100 }
3101
3102 #[test]
3103 fn read_write_header_is_identical_to_read_write_struct() {
3104 #[repr(C, packed)]
3105 struct BTreePackedHeader {
3106 magic: [u8; 3],
3107 version: u8,
3108 max_key_size: u32,
3109 max_value_size: u32,
3110 root_addr: Address,
3111 length: u64,
3112 _buffer: [u8; 24],
3113 }
3114 let packed_header = BTreePackedHeader {
3115 magic: *MAGIC,
3116 version: LAYOUT_VERSION,
3117 root_addr: Address::from(0xDEADBEEF),
3118 max_key_size: 0x12345678,
3119 max_value_size: 0x87654321,
3120 length: 0xA1B2D3C4,
3121 _buffer: [0; 24],
3122 };
3123
3124 let packed_mem = make_memory();
3125 crate::write_struct(&packed_header, Address::from(0), &packed_mem);
3126
3127 let v1_header = BTreeHeader {
3128 version: Version::V1(DerivedPageSize {
3129 max_key_size: 0x12345678,
3130 max_value_size: 0x87654321,
3131 }),
3132 root_addr: Address::from(0xDEADBEEF),
3133 length: 0xA1B2D3C4,
3134 };
3135
3136 let v1_mem = make_memory();
3137 BTreeMap::<Vec<_>, Vec<_>, RefCell<Vec<_>>>::write_header(&v1_header, &v1_mem);
3138
3139 assert_eq!(packed_mem, v1_mem);
3140
3141 let packed_header: BTreePackedHeader = crate::read_struct(Address::from(0), &v1_mem);
3142 let v1_header = BTreeMap::<Vec<_>, Vec<_>, RefCell<Vec<_>>>::read_header(&v1_mem);
3143 assert!(packed_header.magic == *MAGIC);
3144 assert!(packed_header.version == LAYOUT_VERSION);
3145 match v1_header.version {
3146 Version::V1(DerivedPageSize {
3147 max_key_size,
3148 max_value_size,
3149 }) => {
3150 assert!(packed_header.max_key_size == max_key_size);
3151 assert!(packed_header.max_value_size == max_value_size);
3152 }
3153 _ => unreachable!("version must be v1"),
3154 };
3155
3156 assert!(packed_header.root_addr == v1_header.root_addr);
3157 assert!(packed_header.length == v1_header.length);
3158 }
3159
3160 #[test]
3161 fn migrate_from_bounded_to_unbounded_and_back() {
3162 #[derive(PartialOrd, Ord, Clone, Eq, PartialEq, Debug)]
3164 struct T;
3165 impl Storable for T {
3166 fn to_bytes(&self) -> Cow<'_, [u8]> {
3167 Cow::Borrowed(&[1, 2, 3])
3168 }
3169
3170 fn into_bytes(self) -> Vec<u8> {
3171 self.to_bytes().into_owned()
3172 }
3173
3174 fn from_bytes(bytes: Cow<[u8]>) -> Self {
3175 assert_eq!(bytes.to_vec(), vec![1, 2, 3]);
3176 T
3177 }
3178
3179 const BOUND: StorableBound = StorableBound::Bounded {
3180 max_size: 3,
3181 is_fixed_size: true,
3182 };
3183 }
3184
3185 #[derive(PartialOrd, Ord, Clone, Eq, PartialEq, Debug)]
3187 struct T2;
3188 impl Storable for T2 {
3189 fn to_bytes(&self) -> Cow<'_, [u8]> {
3190 Cow::Owned(vec![1, 2, 3])
3191 }
3192
3193 fn into_bytes(self) -> Vec<u8> {
3194 self.to_bytes().into_owned()
3195 }
3196
3197 fn from_bytes(bytes: Cow<[u8]>) -> Self {
3198 assert_eq!(bytes.to_vec(), vec![1, 2, 3]);
3199 T2
3200 }
3201
3202 const BOUND: StorableBound = StorableBound::Unbounded;
3203 }
3204
3205 let mem = make_memory();
3207 let mut btree: BTreeMap<T, T, _> = BTreeMap::new_v1(mem);
3208 btree.insert(T, T);
3209
3210 let btree: BTreeMap<T2, T2, _> = BTreeMap::init(btree.into_memory());
3212 btree.save_header();
3213
3214 let btree: BTreeMap<T2, T2, _> = BTreeMap::init(btree.into_memory());
3216 assert_eq!(btree.get(&T2), Some(T2));
3217
3218 let btree: BTreeMap<T, T, _> = BTreeMap::init(btree.into_memory());
3220 assert_eq!(btree.get(&T), Some(T));
3221 }
3222
3223 #[test]
3224 fn test_clear_new_bounded_type() {
3225 let mem = make_memory();
3226 let mut btree: BTreeMap<Blob<4>, Blob<4>, _> = BTreeMap::new(mem.clone());
3227
3228 btree.insert(
3229 [1u8; 4].as_slice().try_into().unwrap(),
3230 [1u8; 4].as_slice().try_into().unwrap(),
3231 );
3232
3233 assert_ne!(btree.len(), 0);
3234 assert_ne!(btree.allocator.num_allocated_chunks(), 0);
3235 assert_ne!(btree.root_addr, NULL);
3236
3237 btree.clear_new();
3238
3239 let header_actual = BTreeMap::<Blob<4>, Blob<4>, _>::read_header(&mem);
3240
3241 BTreeMap::<Blob<4>, Blob<4>, _>::new(mem.clone());
3242
3243 let header_expected = BTreeMap::<Blob<4>, Blob<4>, _>::read_header(&mem);
3244
3245 assert_eq!(header_actual, header_expected);
3246 }
3247
3248 #[test]
3249 fn test_clear_new_unbounded_type() {
3250 let mem = make_memory();
3251 let mut btree: BTreeMap<String, String, _> = BTreeMap::new(mem.clone());
3252 btree.insert("asd".into(), "bce".into());
3253
3254 assert_ne!(btree.len(), 0);
3255 assert_ne!(btree.allocator.num_allocated_chunks(), 0);
3256 assert_ne!(btree.root_addr, NULL);
3257
3258 btree.clear_new();
3259
3260 let header_actual = BTreeMap::<String, String, _>::read_header(&mem);
3261
3262 BTreeMap::<String, String, _>::new(mem.clone());
3263
3264 let header_expected = BTreeMap::<String, String, _>::read_header(&mem);
3265
3266 assert_eq!(header_actual, header_expected);
3267 }
3268
3269 #[test]
3270 fn deallocating_node_with_overflows() {
3271 let mem = make_memory();
3272 let mut btree: BTreeMap<Vec<u8>, Vec<u8>, _> = BTreeMap::new(mem.clone());
3273
3274 assert_eq!(btree.allocator.num_allocated_chunks(), 0);
3276
3277 btree.insert(vec![0; 10_000], vec![]);
3279
3280 assert!(btree.allocator.num_allocated_chunks() >= 2);
3283 btree.remove(&vec![0; 10_000]);
3284
3285 assert_eq!(btree.allocator.num_allocated_chunks(), 0);
3287 }
3288
3289 #[test]
3290 fn repeatedly_deallocating_nodes_with_overflows() {
3291 let mem = make_memory();
3292 let mut btree: BTreeMap<Vec<u8>, Vec<u8>, _> = BTreeMap::new(mem.clone());
3293
3294 assert_eq!(btree.allocator.num_allocated_chunks(), 0);
3296
3297 for _ in 0..100 {
3298 for i in 0..100 {
3299 btree.insert(vec![i; 10_000], vec![]);
3300 }
3301
3302 for i in 0..100 {
3303 btree.remove(&vec![i; 10_000]);
3304 }
3305 }
3306
3307 assert_eq!(btree.allocator.num_allocated_chunks(), 0);
3309 }
3310
3311 #[test]
3312 fn deallocating_root_does_not_leak_memory() {
3313 let mem = make_memory();
3314 let mut btree: BTreeMap<Vec<u8>, _, _> = BTreeMap::new(mem.clone());
3315
3316 for i in 1..=11 {
3317 assert_eq!(btree.insert(vec![i; 10_000], ()), None);
3319 }
3320
3321 assert_eq!(btree.insert(vec![0; 10_000], ()), None);
3323
3324 let root = btree.load_node(btree.root_addr);
3329 assert_eq!(root.node_type(), NodeType::Internal);
3330 assert_eq!(root.keys(btree.memory()), vec![&[6; 10_000]]);
3331 assert_eq!(root.children_len(), 2);
3332
3333 btree.remove(&vec![6; 10_000]);
3335
3336 let root = btree.load_node(btree.root_addr);
3341 assert_eq!(root.node_type(), NodeType::Internal);
3342 assert_eq!(root.keys(btree.memory()), vec![&[5; 10_000]]);
3343 assert_eq!(root.children_len(), 2);
3344
3345 btree.remove(&vec![5; 10_000]);
3348
3349 let root = btree.load_node(btree.root_addr);
3352 assert_eq!(root.node_type(), NodeType::Leaf);
3353 assert_eq!(
3354 root.keys(btree.memory()),
3355 vec![
3356 &[0; 10_000],
3357 &[1; 10_000],
3358 &[2; 10_000],
3359 &[3; 10_000],
3360 &[4; 10_000],
3361 &[7; 10_000],
3362 &[8; 10_000],
3363 &[9; 10_000],
3364 &[10; 10_000],
3365 &[11; 10_000],
3366 ]
3367 );
3368
3369 for i in 0..=11 {
3371 btree.remove(&vec![i; 10_000]);
3372 }
3373
3374 assert_eq!(btree.allocator.num_allocated_chunks(), 0);
3375 }
3376}