ic_stable_structures/
btreemap.rs

1//! This module implements a key/value store based on a B-Tree
2//! in stable memory.
3//!
4//! # V2 layout
5//!
6//! ```text
7//! ---------------------------------------- <- Address 0
8//! Magic "BTR"                 ↕ 3 bytes
9//! ----------------------------------------
10//! Layout version              ↕ 1 byte
11//! ----------------------------------------
12//! Max key size                ↕ 4 bytes             Page size                   ↕ 4 bytes
13//! ----------------------------------------   OR   ----------------------------------------
14//! Max value size              ↕ 4 bytes             PAGE_SIZE_VALUE_MARKER      ↕ 4 bytes
15//! ----------------------------------------
16//! Root node address           ↕ 8 bytes
17//! ----------------------------------------
18//! Length (number of elements) ↕ 8 bytes
19//! ---------------------------------------- <- Address 28 (PACKED_HEADER_SIZE)
20//! Reserved space              ↕ 24 bytes
21//! ---------------------------------------- <- Address 52 (ALLOCATOR_OFFSET)
22//! Allocator
23//! ----------------------------------------
24//! ... free memory for nodes
25//! ----------------------------------------
26//! ```
27//!
28//! # V1 layout
29//!
30//! ```text
31//! ---------------------------------------- <- Address 0
32//! Magic "BTR"                 ↕ 3 bytes
33//! ----------------------------------------
34//! Layout version              ↕ 1 byte
35//! ----------------------------------------
36//! Max key size                ↕ 4 bytes
37//! ----------------------------------------
38//! Max value size              ↕ 4 bytes
39//! ----------------------------------------
40//! Root node address           ↕ 8 bytes
41//! ----------------------------------------
42//! Length (number of elements) ↕ 8 bytes
43//! ---------------------------------------- <- Address 28 (PACKED_HEADER_SIZE)
44//! Reserved space              ↕ 24 bytes
45//! ---------------------------------------- <- Address 52 (ALLOCATOR_OFFSET)
46//! Allocator
47//! ----------------------------------------
48//! ... free memory for nodes
49//! ----------------------------------------
50//! ```
51mod 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;
73// The sum of all the header fields, i.e. size of a packed header.
74const PACKED_HEADER_SIZE: usize = 28;
75// The offset where the allocator begins.
76const ALLOCATOR_OFFSET: usize = 52;
77
78// The default page size to use in BTreeMap V2 in bytes.
79const DEFAULT_PAGE_SIZE: u32 = 1024;
80
81// A marker to indicate that the `PageSize` stored in the header is a `PageSize::Value`.
82const PAGE_SIZE_VALUE_MARKER: u32 = u32::MAX;
83
84/// A B-Tree map implementation that stores its data into a designated memory.
85///
86/// # Memory Implementations
87///
88/// `BTreeMap` works with any memory implementation that satisfies the [`Memory`] trait:
89///
90/// - [`Ic0StableMemory`](crate::Ic0StableMemory): Stores data in the Internet Computer's stable memory
91/// - [`VectorMemory`](crate::VectorMemory): In-memory implementation backed by a Rust `Vec<u8>`
92/// - [`FileMemory`](crate::FileMemory): Persists data to disk using a file
93/// - [`DefaultMemoryImpl`](crate::DefaultMemoryImpl): Default implementation that automatically selects the
94///   appropriate memory backend based on the environment:
95///   - Uses `Ic0StableMemory` when running in an Internet Computer canister (wasm32 target)
96///   - Falls back to `VectorMemory` in other environments (like tests or non-IC contexts)
97///
98/// For almost all use cases, [`DefaultMemoryImpl`](crate::DefaultMemoryImpl) is recommended as it provides
99/// the right implementation based on the runtime context.
100///
101/// # Examples
102///
103/// ## Basic Usage with a Single BTreeMap
104///
105/// ```rust
106/// use ic_stable_structures::{BTreeMap, DefaultMemoryImpl};
107/// let mut map: BTreeMap<u64, String, _> = BTreeMap::init(DefaultMemoryImpl::default());
108///
109/// map.insert(1, "hello".to_string());
110/// # assert_eq!(map.get(&1), Some("hello".to_string()));
111/// ```
112///
113/// ## Multiple BTreeMaps and Memory Management
114///
115/// **Important**: Each stable structure requires its own designated memory region. Attempting to
116/// initialize multiple structures with the same memory will lead to data corruption.
117///
118/// ### What NOT to do:
119///
120/// ```rust,no_run
121/// use ic_stable_structures::{BTreeMap, DefaultMemoryImpl};
122///
123/// // ERROR: Using the same memory for multiple BTreeMaps will corrupt data
124/// let mut map_1: BTreeMap<u64, String, _> = BTreeMap::init(DefaultMemoryImpl::default());
125/// let mut map_2: BTreeMap<u64, String, _> = BTreeMap::init(DefaultMemoryImpl::default());
126///
127/// map_1.insert(1, "two".to_string());
128/// map_2.insert(1, "three".to_string());
129/// // This assertion would fail: changes to map_2 corrupt map_1's data
130/// assert_eq!(map_1.get(&1), Some("two".to_string()));
131/// ```
132///
133/// ### Correct approach using MemoryManager:
134///
135/// ```rust
136/// use ic_stable_structures::{
137///    memory_manager::{MemoryId, MemoryManager},
138///    BTreeMap, DefaultMemoryImpl,
139/// };
140///
141/// // Initialize the memory manager with a single memory
142/// let memory_manager = MemoryManager::init(DefaultMemoryImpl::default());
143///
144/// // Get separate virtual memories for each BTreeMap
145/// let mut map_1: BTreeMap<u64, String, _> = BTreeMap::init(memory_manager.get(MemoryId::new(0)));
146/// let mut map_2: BTreeMap<u64, String, _> = BTreeMap::init(memory_manager.get(MemoryId::new(1)));
147///
148/// map_1.insert(1, "two".to_string());
149/// map_2.insert(1, "three".to_string());
150/// // Now this works as expected
151/// assert_eq!(map_1.get(&1), Some("two".to_string()));
152/// ```
153///
154/// The [`MemoryManager`](crate::memory_manager::MemoryManager) creates up to 255 virtual memories
155/// from a single contiguous memory, allowing multiple stable structures to safely coexist.
156///
157/// For complete examples of using multiple stable structures in a production environment, see the
158/// [Quick Start example](https://github.com/dfinity/stable-structures/tree/main/examples/src/quick_start).
159///
160/// ## Custom Types
161///
162/// You can store custom structs in a `BTreeMap` by implementing the `Storable` trait:
163///
164/// ```rust
165/// use ic_stable_structures::{BTreeMap, DefaultMemoryImpl, Storable, storable::Bound};
166/// use std::borrow::Cow;
167///
168/// #[derive(Debug, PartialEq)]
169/// struct User {
170///     id: u64,
171///     name: String,
172/// }
173///
174/// impl Storable for User {
175///     fn to_bytes(&self) -> Cow<'_, [u8]> {
176///         let mut bytes = Vec::new();
177///         // TODO: Convert your struct to bytes...
178///         Cow::Owned(bytes)
179///     }
180///
181///     fn into_bytes(self) -> Vec<u8> {
182///         let mut bytes = Vec::new();
183///         // TODO: Convert your struct to bytes...
184///         bytes
185///     }
186///
187///     fn from_bytes(bytes: Cow<[u8]>) -> Self {
188///         // TODO: Convert bytes back to your struct
189///         let (id, name) = (0, "".to_string());
190///         User { id, name }
191///     }
192///
193///     // Types can be bounded or unbounded:
194///     // - Use Bound::Unbounded if the size can vary or isn't known in advance (recommended for most cases)
195///     // - Use Bound::Bounded if you know the maximum size and want memory optimization
196///     const BOUND: Bound = Bound::Unbounded;
197/// }
198///
199/// // Now you can use your custom type in a BTreeMap
200/// let mut map: BTreeMap<u64, User, _> = BTreeMap::init(DefaultMemoryImpl::default());
201///
202/// let user = User { id: 1, name: "Alice".to_string() };
203/// map.insert(1, user);
204///
205/// // Retrieving the custom type
206/// if let Some(user) = map.get(&1) {
207///     println!("Found user: {}", user.name);
208/// }
209/// ```
210///
211/// ### Bounded vs Unbounded Types
212///
213/// When implementing `Storable`, you must specify whether your type is bounded or unbounded:
214///
215/// - **Unbounded (`Bound::Unbounded`)**:
216///   - Use when your type's serialized size can vary or has no fixed maximum
217///   - Recommended for most custom types, especially those containing Strings or Vecs
218///   - Example: `const BOUND: Bound = Bound::Unbounded;`
219///
220/// - **Bounded (`Bound::Bounded{ max_size, is_fixed_size }`)**:
221///   - Use when you know the maximum serialized size of your type
222///   - Enables memory optimizations in the BTreeMap
223///   - Example: `const BOUND: Bound = Bound::Bounded { max_size: 100, is_fixed_size: false };`
224///   - For types with truly fixed size (like primitive types), set `is_fixed_size: true`
225///
226/// If unsure, use `Bound::Unbounded` as it's the safer choice.
227///
228/// # Warning
229///
230/// Once you've deployed with a bounded type, you cannot increase its `max_size` in
231/// future versions without risking data corruption. You can, however, migrate from a bounded type
232/// to an unbounded type if needed. For evolving data structures, prefer `Bound::Unbounded`.
233pub struct BTreeMap<K, V, M>
234where
235    K: Storable + Ord + Clone,
236    V: Storable,
237    M: Memory,
238{
239    // The address of the root node. If a root node doesn't exist, the address
240    // is set to NULL.
241    root_addr: Address,
242
243    version: Version,
244
245    // An allocator used for managing memory and allocating nodes.
246    allocator: Allocator<M>,
247
248    // The number of elements in the map.
249    length: u64,
250
251    // A marker to communicate to the Rust compiler that we own these types.
252    _phantom: PhantomData<(K, V)>,
253}
254
255#[derive(PartialEq, Debug)]
256/// The packed header size must be <= ALLOCATOR_OFFSET.
257struct BTreeHeader {
258    version: Version,
259    root_addr: Address,
260    length: u64,
261    // Reserved bytes for future extensions
262}
263
264impl<K, V, M> BTreeMap<K, V, M>
265where
266    K: Storable + Ord + Clone,
267    V: Storable,
268    M: Memory,
269{
270    /// Initializes a `BTreeMap`.
271    ///
272    /// If the memory provided already contains a `BTreeMap`, then that
273    /// map is loaded. Otherwise, a new `BTreeMap` instance is created.
274    pub fn init(memory: M) -> Self {
275        if memory.size() == 0 {
276            // Memory is empty. Create a new map.
277            return BTreeMap::new(memory);
278        }
279
280        // Check if the magic in the memory corresponds to a BTreeMap.
281        let mut dst = vec![0; 3];
282        memory.read(0, &mut dst);
283        if dst != MAGIC {
284            // No BTreeMap found. Create a new instance.
285            BTreeMap::new(memory)
286        } else {
287            // The memory already contains a BTreeMap. Load it.
288            BTreeMap::load(memory)
289        }
290    }
291
292    /// Initializes a v1 `BTreeMap`.
293    ///
294    /// This is exposed only in testing.
295    #[cfg(test)]
296    pub fn init_v1(memory: M) -> Self {
297        if memory.size() == 0 {
298            // Memory is empty. Create a new map.
299            return BTreeMap::new_v1(memory);
300        }
301
302        // Check if the magic in the memory corresponds to a BTreeMap.
303        let mut dst = vec![0; 3];
304        memory.read(0, &mut dst);
305        if dst != MAGIC {
306            // No BTreeMap found. Create a new instance.
307            BTreeMap::new_v1(memory)
308        } else {
309            // The memory already contains a BTreeMap. Load it, making sure
310            // we don't migrate the BTreeMap to v2.
311            BTreeMap::load_helper(memory, false)
312        }
313    }
314
315    /// Creates a new instance a `BTreeMap`.
316    ///
317    /// The given `memory` is assumed to be exclusively reserved for this data
318    /// structure and that it starts at address zero. Typically `memory` will
319    /// be an instance of `RestrictedMemory`.
320    ///
321    /// When initialized, the data structure has the following memory layout:
322    ///
323    ///    |  BTreeHeader  |  Allocator | ... free memory for nodes |
324    ///
325    /// See `Allocator` for more details on its own memory layout.
326    pub fn new(memory: M) -> Self {
327        let page_size = match (K::BOUND, V::BOUND) {
328            // The keys and values are both bounded.
329            (
330                StorableBound::Bounded {
331                    max_size: max_key_size,
332                    ..
333                },
334                StorableBound::Bounded {
335                    max_size: max_value_size,
336                    ..
337                },
338            ) => {
339                // Get the maximum possible node size.
340                let max_node_size = Node::<K>::max_size(max_key_size, max_value_size).get();
341
342                // A node can have at most 11 entries, and an analysis has shown that ~70% of all
343                // nodes have <= 8 entries. We can therefore use a page size that's 8/11 the
344                // maximum size with little performance difference but with a significant storage
345                // saving. We round the 8/11 to be 3/4.
346                PageSize::Value((max_node_size * 3 / 4) as u32)
347            }
348            // Use a default page size.
349            _ => 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    /// Create a v1 instance of the BTree.
369    ///
370    /// This is only exposed for testing.
371    #[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    /// Loads the map from memory.
396    pub fn load(memory: M) -> Self {
397        Self::load_helper(memory, true)
398    }
399
400    /// Loads the map from memory, potentially migrating the map from V1 to V2.
401    fn load_helper(memory: M, migrate_to_v2: bool) -> Self {
402        // Read the header from memory.
403        let header = Self::read_header(&memory);
404
405        let version = match header.version {
406            Version::V1(derived_page_size) => {
407                // Migrate to V2 if flag is enabled.
408                if migrate_to_v2 {
409                    Version::V2(PageSize::Derived(derived_page_size))
410                } else {
411                    // Assert that the bounds are correct.
412                    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    /// Reads the header from the specified memory.
442    fn read_header(memory: &M) -> BTreeHeader {
443        // Read the header
444        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                // Deserialize the fields
452                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                // Load the page size.
463                let page_size = {
464                    // Page sizes can be stored either as a direct value or as max/value sizes.
465                    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                        // Page size is stored as a direct value
470                        PageSize::Value(a)
471                    } else {
472                        // Page size is stored as a derived value.
473                        PageSize::Derived(DerivedPageSize {
474                            max_key_size: a,
475                            max_value_size: b,
476                        })
477                    }
478                };
479
480                // Deserialize the fields
481                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    /// Inserts a key-value pair into the map.
494    ///
495    /// The previous value of the key, if present, is returned.
496    ///
497    /// PRECONDITION:
498    ///   key.to_bytes().len() <= max_size(Key)
499    ///   value.to_bytes().len() <= max_size(Value)
500    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            // No root present. Allocate one.
505            let node = self.allocate_node(NodeType::Leaf);
506            self.root_addr = node.address();
507            self.save_header();
508            node
509        } else {
510            // Load the root from memory.
511            let mut root = self.load_node(self.root_addr);
512
513            // Check if the key already exists in the root.
514            if let Ok(idx) = root.search(&key, self.memory()) {
515                // Key found, replace its value and return the old one.
516                return Some(V::from_bytes(Cow::Owned(
517                    self.update_value(&mut root, idx, value),
518                )));
519            }
520
521            // If the root is full, we need to introduce a new node as the root.
522            //
523            // NOTE: In the case where we are overwriting an existing key, then introducing
524            // a new root node isn't strictly necessary. However, that's a micro-optimization
525            // that adds more complexity than it's worth.
526            if root.is_full() {
527                // The root is full. Allocate a new node that will be used as the new root.
528                let mut new_root = self.allocate_node(NodeType::Internal);
529
530                // The new root has the old root as its only child.
531                new_root.push_child(self.root_addr);
532
533                // Update the root address.
534                self.root_addr = new_root.address();
535                self.save_header();
536
537                // Split the old (full) root.
538                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    /// Inserts an entry into a node that is *not full*.
552    fn insert_nonfull(&mut self, mut node: Node<K>, key: K, value: Vec<u8>) -> Option<Vec<u8>> {
553        // We're guaranteed by the caller that the provided node is not full.
554        assert!(!node.is_full());
555
556        // Look for the key in the node.
557        match node.search(&key, self.memory()) {
558            Ok(idx) => {
559                // Key found, replace its value and return the old one.
560                Some(self.update_value(&mut node, idx, value))
561            }
562            Err(idx) => {
563                // The key isn't in the node. `idx` is where that key should be inserted.
564
565                match node.node_type() {
566                    NodeType::Leaf => {
567                        // The node is a non-full leaf.
568                        // Insert the entry at the proper location.
569                        node.insert_entry(idx, (key, value));
570                        self.save_node(&mut node);
571
572                        // Update the length.
573                        self.length += 1;
574                        self.save_header();
575
576                        // No previous value to return.
577                        None
578                    }
579                    NodeType::Internal => {
580                        // The node is an internal node.
581                        // Load the child that we should add the entry to.
582                        let mut child = self.load_node(node.child(idx));
583
584                        if child.is_full() {
585                            // Check if the key already exists in the child.
586                            if let Ok(idx) = child.search(&key, self.memory()) {
587                                // Key found, replace its value and return the old one.
588                                return Some(self.update_value(&mut child, idx, value));
589                            }
590
591                            // The child is full. Split the child.
592                            self.split_child(&mut node, idx);
593
594                            // The children have now changed. Search again for
595                            // the child where we need to store the entry in.
596                            let idx = node.search(&key, self.memory()).unwrap_or_else(|idx| idx);
597                            child = self.load_node(node.child(idx));
598                        }
599
600                        // The child should now be not full.
601                        assert!(!child.is_full());
602
603                        self.insert_nonfull(child, key, value)
604                    }
605                }
606            }
607        }
608    }
609
610    /// Takes as input a nonfull internal `node` and index to its full child, then
611    /// splits this child into two, adding an additional child to `node`.
612    ///
613    /// Example:
614    /// ```ignore
615    ///                          [ ... M   Y ... ]
616    ///                                  |
617    ///                 [ N  O  P  Q  R  S  T  U  V  W  X ]
618    /// ```
619    ///
620    /// After splitting becomes:
621    /// ```ignore
622    ///                         [ ... M  S  Y ... ]
623    ///                                 / \
624    ///                [ N  O  P  Q  R ]   [ T  U  V  W  X ]
625    /// ```
626    ///
627    fn split_child(&mut self, node: &mut Node<K>, full_child_idx: usize) {
628        // The node must not be full.
629        assert!(!node.is_full());
630
631        // The node's child must be full.
632        let mut full_child = self.load_node(node.child(full_child_idx));
633        assert!(full_child.is_full());
634
635        // Create a sibling to this full child (which has to be the same type).
636        let mut sibling = self.allocate_node(full_child.node_type());
637        assert_eq!(sibling.node_type(), full_child.node_type());
638
639        // Add sibling as a new child in the node.
640        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    /// Returns the value for the given key, if it exists.
652    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 // Extract value.
658        })
659        .map(Cow::Owned)
660        .map(V::from_bytes)
661    }
662
663    /// Returns true if the key exists.
664    pub fn contains_key(&self, key: &K) -> bool {
665        // An empty closure returns Some(()) if the key is found.
666        self.root_addr != NULL && self.traverse(self.root_addr, key, |_, _| ()).is_some()
667    }
668
669    /// Recursively traverses from `node_addr`, invoking `f` if `key` is found. Stops at a leaf if not.
670    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        // Look for the key in the current node.
676        match node.search(key, self.memory()) {
677            Ok(idx) => Some(f(&mut node, idx)), // Key found: apply `f`.
678            Err(idx) => match node.node_type() {
679                NodeType::Leaf => None, // At a leaf: key not present.
680                NodeType::Internal => self.traverse(node.child(idx), key, f), // Continue search in child.
681            },
682        }
683    }
684
685    /// Returns `true` if the map contains no elements.
686    pub fn is_empty(&self) -> bool {
687        self.length == 0
688    }
689
690    /// Returns the number of elements in the map.
691    pub fn len(&self) -> u64 {
692        self.length
693    }
694
695    /// Returns the underlying memory.
696    pub fn into_memory(self) -> M {
697        self.allocator.into_memory()
698    }
699
700    /// Removes all elements from the map.
701    #[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    /// Removes all elements from the map.
708    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    /// Returns the first key-value pair in the map. The key in this
716    /// pair is the minimum key in the map.
717    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    /// Returns the last key-value pair in the map. The key in this
727    /// pair is the maximum key in the map.
728    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    /// Removes a key from the map, returning the previous value at the key if it exists.
746    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    /// Removes and returns the last element in the map. The key of this element is the maximum key that was in the map
758    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    /// Removes and returns the first element in the map. The key of this element is the minimum key that was in the map
770    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    /// A helper method for recursively removing a key from the B-tree.
782    fn remove_helper(&mut self, mut node: Node<K>, key: &K) -> Option<Vec<u8>> {
783        if node.address() != self.root_addr {
784            // We're guaranteed that whenever this method is called an entry can be
785            // removed from the node without it needing to be merged into a sibling.
786            // This strengthened condition allows us to delete an entry in a single
787            // pass most of the time without having to back up.
788            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                        // Case 1: The node is a leaf node and the key exists in it.
796                        // This is the simplest case. The key is removed from the leaf.
797                        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                            // Deallocate the empty node.
807                            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, // Key not found.
817                }
818            }
819            NodeType::Internal => {
820                match node.search(key, self.memory()) {
821                    Ok(idx) => {
822                        // Case 2: The node is an internal node and the key exists in it.
823
824                        let left_child = self.load_node(node.child(idx));
825                        if left_child.can_remove_entry_without_merging() {
826                            // Case 2.a: A key can be removed from the left child without merging.
827                            //
828                            //                       parent
829                            //                  [..., key, ...]
830                            //                       /   \
831                            //            [left child]   [...]
832                            //           /            \
833                            //        [...]         [..., key predecessor]
834                            //
835                            // In this case, we replace `key` with the key's predecessor from the
836                            // left child's subtree, then we recursively delete the key's
837                            // predecessor for the following end result:
838                            //
839                            //                       parent
840                            //            [..., key predecessor, ...]
841                            //                       /   \
842                            //            [left child]   [...]
843                            //           /            \
844                            //        [...]          [...]
845
846                            // Recursively delete the predecessor.
847                            // TODO(EXC-1034): Do this in a single pass.
848                            let predecessor = left_child.get_max(self.memory());
849                            self.remove_helper(left_child, &predecessor.0)?;
850
851                            // Replace the `key` with its predecessor.
852                            let (_, old_value) = node.swap_entry(idx, predecessor, self.memory());
853
854                            // Save the parent node.
855                            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                            // Case 2.b: A key can be removed from the right child without merging.
862                            //
863                            //                       parent
864                            //                  [..., key, ...]
865                            //                       /   \
866                            //                   [...]   [right child]
867                            //                          /             \
868                            //              [key successor, ...]     [...]
869                            //
870                            // In this case, we replace `key` with the key's successor from the
871                            // right child's subtree, then we recursively delete the key's
872                            // successor for the following end result:
873                            //
874                            //                       parent
875                            //            [..., key successor, ...]
876                            //                       /   \
877                            //                  [...]   [right child]
878                            //                           /            \
879                            //                        [...]          [...]
880
881                            // Recursively delete the successor.
882                            // TODO(EXC-1034): Do this in a single pass.
883                            let successor = right_child.get_min(self.memory());
884                            self.remove_helper(right_child, &successor.0)?;
885
886                            // Replace the `key` with its successor.
887                            let (_, old_value) = node.swap_entry(idx, successor, self.memory());
888
889                            // Save the parent node.
890                            self.save_node(&mut node);
891                            return Some(old_value);
892                        }
893
894                        // Case 2.c: Both the left and right child are at their minimum sizes.
895                        //
896                        //                       parent
897                        //                  [..., key, ...]
898                        //                       /   \
899                        //            [left child]   [right child]
900                        //
901                        // In this case, we merge (left child, key, right child) into a single
902                        // node. The result will look like this:
903                        //
904                        //                       parent
905                        //                     [...  ...]
906                        //                         |
907                        //          [left child, `key`, right child] <= new child
908                        //
909                        // We then recurse on this new child to delete `key`.
910                        //
911                        // If `parent` becomes empty (which can only happen if it's the root),
912                        // then `parent` is deleted and `new_child` becomes the new root.
913                        assert!(left_child.at_minimum());
914                        assert!(right_child.at_minimum());
915
916                        // Merge the right child into the left child.
917                        let mut new_child = self.merge(
918                            right_child,
919                            left_child,
920                            node.remove_entry(idx, self.memory()),
921                        );
922
923                        // Remove the right child from the parent node.
924                        node.remove_child(idx + 1);
925
926                        if node.entries_len() == 0 {
927                            // Can only happen if this node is root.
928                            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                            // Deallocate the root node.
935                            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                        // Recursively delete the key.
944                        self.remove_helper(new_child, key)
945                    }
946                    Err(idx) => {
947                        // Case 3: The node is an internal node and the key does NOT exist in it.
948
949                        // If the key does exist in the tree, it will exist in the subtree at index
950                        // `idx`.
951                        let mut child = self.load_node(node.child(idx));
952
953                        if child.can_remove_entry_without_merging() {
954                            // The child has enough nodes. Recurse to delete the `key` from the
955                            // `child`.
956                            return self.remove_helper(child, key);
957                        }
958
959                        // An entry can't be removed from the child without merging.
960                        // See if it has a sibling where an entry can be removed without merging.
961                        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                                // Case 3.a (left):
976                                // A key can be removed from the left child without merging.
977                                //
978                                //                            [d] (parent)
979                                //                           /   \
980                                //  (left sibling) [a, b, c]     [e, f] (child)
981                                //                         \
982                                //                         [c']
983                                //
984                                // In this case, we move a key down from the parent into the child
985                                // and move a key from the left sibling up into the parent
986                                // resulting in the following tree:
987                                //
988                                //                            [c] (parent)
989                                //                           /   \
990                                //       (left sibling) [a, b]   [d, e, f] (child)
991                                //                              /
992                                //                            [c']
993                                //
994                                // We then recurse to delete the key from the child.
995
996                                // Remove the last entry from the left sibling.
997                                let (left_sibling_key, left_sibling_value) =
998                                    left_sibling.pop_entry(self.memory()).unwrap();
999
1000                                // Replace the parent's entry with the one from the left sibling.
1001                                let (parent_key, parent_value) = node.swap_entry(
1002                                    idx - 1,
1003                                    (left_sibling_key, left_sibling_value),
1004                                    self.memory(),
1005                                );
1006
1007                                // Move the entry from the parent into the child.
1008                                child.insert_entry(0, (parent_key, parent_value));
1009
1010                                // Move the last child from left sibling into child.
1011                                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                                // Case 3.a (right):
1031                                // A key can be removed from the right child without merging.
1032                                //
1033                                //                            [c] (parent)
1034                                //                           /   \
1035                                //             (child) [a, b]     [d, e, f] (right sibling)
1036                                //                               /
1037                                //                            [d']
1038                                //
1039                                // In this case, we move a key down from the parent into the child
1040                                // and move a key from the right sibling up into the parent
1041                                // resulting in the following tree:
1042                                //
1043                                //                            [d] (parent)
1044                                //                           /   \
1045                                //          (child) [a, b, c]     [e, f] (right sibling)
1046                                //                          \
1047                                //                           [d']
1048                                //
1049                                // We then recurse to delete the key from the child.
1050
1051                                // Remove the first entry from the right sibling.
1052                                let (right_sibling_key, right_sibling_value) =
1053                                    right_sibling.remove_entry(0, self.memory());
1054
1055                                // Replace the parent's entry with the one from the right sibling.
1056                                let parent_entry = node.swap_entry(
1057                                    idx,
1058                                    (right_sibling_key, right_sibling_value),
1059                                    self.memory(),
1060                                );
1061
1062                                // Move the entry from the parent into the child.
1063                                child.push_entry(parent_entry);
1064
1065                                // Move the first child of right_sibling into `child`.
1066                                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                        // Case 3.b: Both the left and right siblings are at their minimum sizes.
1084
1085                        if let Some(left_sibling) = left_sibling {
1086                            // Merge child into left sibling if it exists.
1087
1088                            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                            // Removing child from parent.
1095                            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                                    // Update the root.
1103                                    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                            // Merge child into right sibling.
1115
1116                            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                            // Removing child from parent.
1124                            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                                    // Update the root.
1132                                    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    /// Returns an iterator over the entries of the map, sorted by key.
1150    ///
1151    /// # Example
1152    ///
1153    /// ```rust
1154    /// use ic_stable_structures::{BTreeMap, DefaultMemoryImpl};
1155    /// let mut map: BTreeMap<u64, String, _> = BTreeMap::init(DefaultMemoryImpl::default());
1156    ///
1157    /// map.insert(1, "one".to_string());
1158    /// map.insert(2, "two".to_string());
1159    ///
1160    /// for entry in map.iter() {
1161    ///     println!("{}: {}", entry.key(), entry.value());
1162    /// }
1163    /// ```
1164    pub fn iter(&self) -> Iter<'_, K, V, M> {
1165        self.iter_internal().into()
1166    }
1167
1168    /// Returns an iterator over the entries in the map where keys
1169    /// belong to the specified range.
1170    ///
1171    /// # Example
1172    ///
1173    /// ```rust
1174    /// use ic_stable_structures::{BTreeMap, DefaultMemoryImpl};
1175    /// use std::ops::Bound;
1176    ///
1177    /// let mut map: BTreeMap<u64, String, _> = BTreeMap::init(DefaultMemoryImpl::default());
1178    /// map.insert(1, "one".to_string());
1179    /// map.insert(2, "two".to_string());
1180    /// map.insert(3, "three".to_string());
1181    ///
1182    /// // Get entries with keys between 1 and 3 (inclusive)
1183    /// for entry in map.range((Bound::Included(1), Bound::Included(3))) {
1184    ///     println!("{}: {}", entry.key(), entry.value());
1185    /// }
1186    /// ```
1187    pub fn range(&self, key_range: impl RangeBounds<K>) -> Iter<'_, K, V, M> {
1188        self.range_internal(key_range).into()
1189    }
1190
1191    /// Returns an iterator starting just before the given key.
1192    ///
1193    /// Finds the largest key strictly less than `bound` and starts from it.
1194    /// Useful when `range(bound..)` skips the previous element.
1195    ///
1196    /// Returns an empty iterator if no smaller key exists.
1197    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**: use `iter_from_prev_key` instead.
1207    ///
1208    /// The name `iter_upper_bound` was misleading — it suggested an inclusive
1209    /// upper bound. In reality, it starts from the largest key strictly less
1210    /// than the given bound.
1211    ///
1212    /// The new name, `iter_from_prev_key`, better reflects this behavior and
1213    /// improves code clarity.
1214    #[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    /// Returns an iterator over the keys of the map.
1220    pub fn keys(&self) -> KeysIter<'_, K, V, M> {
1221        self.iter_internal().into()
1222    }
1223
1224    /// Returns an iterator over the keys of the map which belong to the specified range.
1225    pub fn keys_range(&self, key_range: impl RangeBounds<K>) -> KeysIter<'_, K, V, M> {
1226        self.range_internal(key_range).into()
1227    }
1228
1229    /// Returns an iterator over the values of the map, sorted by key.
1230    pub fn values(&self) -> ValuesIter<'_, K, V, M> {
1231        self.iter_internal().into()
1232    }
1233
1234    /// Returns an iterator over the values of the map where keys
1235    /// belong to the specified range.
1236    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            // Map is empty.
1247            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    /// Merges one node (`source`) into another (`into`), along with a median entry.
1259    ///
1260    /// Example (values are not included for brevity):
1261    ///
1262    /// Input:
1263    ///   Source: [1, 2, 3]
1264    ///   Into: [5, 6, 7]
1265    ///   Median: 4
1266    ///
1267    /// Output:
1268    ///   [1, 2, 3, 4, 5, 6, 7] (stored in the `into` node)
1269    ///   `source` is deallocated.
1270    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    /// Allocates a new node of the given type.
1276    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    /// Deallocates a node.
1284    #[inline]
1285    fn deallocate_node(&mut self, node: Node<K>) {
1286        node.deallocate(self.allocator_mut());
1287    }
1288
1289    /// Loads a node from memory.
1290    #[inline]
1291    fn load_node(&self, address: Address) -> Node<K> {
1292        Node::load(address, self.version.page_size(), self.memory())
1293    }
1294
1295    /// Saves the node to memory.
1296    #[inline]
1297    fn save_node(&mut self, node: &mut Node<K>) {
1298        node.save(self.allocator_mut());
1299    }
1300
1301    /// Replaces the value at `idx` in the node, saves the node, and returns the old value.
1302    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    /// Saves the map to memory.
1309    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    /// Write the layout header to the memory.
1320    fn write_header(header: &BTreeHeader, memory: &M) {
1321        // Serialize the header
1322        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        // Write the header
1346        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    /// A helper function to clone and collect borrowed key/value pairs into a `Vec`.
1363    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    /// A helper function to collect owned key/value pairs into a `Vec`.
1370    fn collect<K: Clone, V: Clone>(it: impl Iterator<Item = (K, V)>) -> Vec<(K, V)> {
1371        it.collect()
1372    }
1373
1374    /// A helper function to collect lazy entries into a `Vec`.
1375    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    /// Returns a fixed‑size buffer for the given u32.
1385    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    /// A trait to construct a value from a u32.
1416    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    /// Encodes an object into a byte vector.
1465    fn encode<T: Storable>(object: T) -> Vec<u8> {
1466        object.into_bytes_checked()
1467    }
1468
1469    /// A helper method to succinctly create a blob.
1470    pub(crate) fn b(x: &[u8]) -> Blob<10> {
1471        Blob::<10>::try_from(x).unwrap()
1472    }
1473
1474    /// Creates a new shared memory instance.
1475    pub(crate) fn make_memory() -> Rc<RefCell<Vec<u8>>> {
1476        Rc::new(RefCell::new(Vec::new()))
1477    }
1478
1479    /// A test runner that runs the test using V1, migrated V2, and direct V2.
1480    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        // V1 does not support unbounded types.
1487        if K::BOUND != StorableBound::Unbounded && V::BOUND != StorableBound::Unbounded {
1488            // Test with V1.
1489            let mem = make_memory();
1490            let tree_v1 = BTreeMap::new_v1(mem);
1491            f(tree_v1);
1492
1493            // Test with V2 migrated from V1.
1494            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        // Test with direct V2.
1501        let mem = make_memory();
1502        let tree_v2 = BTreeMap::new(mem);
1503        f(tree_v2);
1504    }
1505
1506    /// Checks that objects from boundary u32 values are strictly increasing.
1507    /// This ensures multi-byte conversions preserve order.
1508    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    /// Asserts that the associated `BOUND` for `$ty` is _not_ `Unbounded`.
1519    macro_rules! assert_bounded {
1520        ($ty:ty) => {
1521            assert_ne!(<$ty>::BOUND, StorableBound::Unbounded, "Must be Bounded");
1522        };
1523    }
1524
1525    /// Asserts that the associated `BOUND` for `$ty` _is_ `Unbounded`.
1526    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            // Empty value.
1537            $runner::<$K, ()>();
1538
1539            // Bounded values.
1540            assert_bounded!(u32);
1541            $runner::<$K, u32>();
1542
1543            assert_bounded!(Blob<20>);
1544            $runner::<$K, Blob<20>>();
1545
1546            // Unbounded values.
1547            assert_unbounded!(MonotonicVec32);
1548            $runner::<$K, MonotonicVec32>();
1549
1550            assert_unbounded!(MonotonicString32);
1551            $runner::<$K, MonotonicString32>();
1552        }};
1553    }
1554
1555    /// Macro to apply a test function to a predefined grid of key/value types.
1556    macro_rules! btree_test {
1557        ($name:ident, $runner:ident) => {
1558            #[test]
1559            fn $name() {
1560                // Bounded keys.
1561                assert_bounded!(u32);
1562                run_with_key!($runner, u32);
1563
1564                assert_bounded!(Blob<10>);
1565                run_with_key!($runner, Blob<10>);
1566
1567                // Unbounded keys.
1568                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    // Define a trait for keys that need the full set of bounds.
1578    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    // Define a trait for values that need the full set of bounds.
1582    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            // Reload the BTreeMap and verify the entry.
1595            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            // The result should look like this:
1643            //                [6]
1644            //               /   \
1645            // [1, 2, 3, 4, 5]   [7, 8, 9, 10, 11, (12), 13, 14, 15, 16, 17]
1646
1647            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            // The right child should now be full, with the median key being "12"
1656            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            // Overwrite the value of the median key.
1663            assert_eq!(btree.insert(median_key.clone(), value(123)), Some(value(0)));
1664            assert_eq!(btree.get(&median_key), Some(value(123)));
1665
1666            // The child has not been split and is still full.
1667            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            // We now have a root that is full and looks like this:
1685            //
1686            // [1, 2, 3, 4, 5, (6), 7, 8, 9, 10, 11]
1687            let root = btree.load_node(btree.root_addr);
1688            assert!(root.is_full());
1689
1690            // Overwrite an element in the root. It should NOT cause the node to be split.
1691            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            // Insert entries until the root node is full.
1722            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            // Only need a single allocation to store up to `CAPACITY` elements.
1733            assert_eq!(btree.allocator.num_allocated_chunks(), 1);
1734
1735            assert_eq!(btree.insert(key(255), V::empty()), None);
1736
1737            // The node had to be split into three nodes.
1738            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            // Should now split a node.
1751            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            // The result should look like this:
1756            //                [6]
1757            //               /   \
1758            // [1, 2, 3, 4, 5]   [7, 8, 9, 10, 11, 12]
1759
1760            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            // Should now split a node.
1775            assert_eq!(btree.insert(key(12), V::empty()), None);
1776
1777            // The result should look like this:
1778            //                [6]
1779            //               /   \
1780            // [1, 2, 3, 4, 5]   [7, 8, 9, 10, 11, 12]
1781
1782            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            // Insert more to cause more splitting.
1806            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            // Should cause another split
1812            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            // Should now split a node.
1936            assert_eq!(btree.insert(key(0), V::empty()), None);
1937
1938            // The result should look like this:
1939            //                    [6]
1940            //                   /   \
1941            // [0, 1, 2, 3, 4, 5]     [7, 8, 9, 10, 11]
1942
1943            for i in 0..=11 {
1944                assert_eq!(btree.get(&key(i)), Some(V::empty()));
1945            }
1946
1947            // Remove node 6. Triggers case 2.a
1948            assert_eq!(btree.remove(&key(6)), Some(V::empty()));
1949
1950            // The result should look like this:
1951            //                [5]
1952            //               /   \
1953            // [0, 1, 2, 3, 4]   [7, 8, 9, 10, 11]
1954            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            // There are three allocated nodes.
1974            assert_eq!(btree.allocator.num_allocated_chunks(), 3);
1975
1976            // Remove node 5. Triggers case 2c
1977            assert_eq!(btree.remove(&key(5)), Some(V::empty()));
1978
1979            // Reload the btree to verify that we saved it correctly.
1980            let btree = BTreeMap::<K, V, _>::load(btree.into_memory());
1981
1982            // The result should look like this:
1983            // [0, 1, 2, 3, 4, 7, 8, 9, 10, 11]
1984            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            // There is only one node allocated.
1991            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            // Should now split a node.
2004            assert_eq!(btree.insert(key(12), V::empty()), None);
2005
2006            // The result should look like this:
2007            //                [6]
2008            //               /   \
2009            // [1, 2, 3, 4, 5]   [7, 8, 9, 10, 11, 12]
2010
2011            for i in 1..=12 {
2012                assert_eq!(btree.get(&key(i)), Some(V::empty()));
2013            }
2014
2015            // Remove node 6. Triggers case 2.b
2016            assert_eq!(btree.remove(&key(6)), Some(V::empty()));
2017
2018            // The result should look like this:
2019            //                [7]
2020            //               /   \
2021            // [1, 2, 3, 4, 5]   [8, 9, 10, 11, 12]
2022            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            // Remove node 7. Triggers case 2.c
2042            assert_eq!(btree.remove(&key(7)), Some(V::empty()));
2043            // The result should look like this:
2044            //
2045            // [1, 2, 3, 4, 5, 8, 9, 10, 11, 12]
2046            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            // Should now split a node.
2067            assert_eq!(btree.insert(key(12), V::empty()), None);
2068
2069            // The result should look like this:
2070            //                [6]
2071            //               /   \
2072            // [1, 2, 3, 4, 5]   [7, 8, 9, 10, 11, 12]
2073
2074            // Remove node 3. Triggers case 3.a
2075            assert_eq!(btree.remove(&key(3)), Some(V::empty()));
2076
2077            // The result should look like this:
2078            //                [7]
2079            //               /   \
2080            // [1, 2, 4, 5, 6]   [8, 9, 10, 11, 12]
2081            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            // There are three allocated nodes.
2101            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            // Should now split a node.
2114            assert_eq!(btree.insert(key(0), V::empty()), None);
2115
2116            // The result should look like this:
2117            //                   [6]
2118            //                  /   \
2119            // [0, 1, 2, 3, 4, 5]   [7, 8, 9, 10, 11]
2120
2121            // Remove node 8. Triggers case 3.a left
2122            assert_eq!(btree.remove(&key(8)), Some(V::empty()));
2123
2124            // The result should look like this:
2125            //                [5]
2126            //               /   \
2127            // [0, 1, 2, 3, 4]   [6, 7, 9, 10, 11]
2128            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            // There are three allocated nodes.
2148            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            // Should now split a node.
2161            assert_eq!(btree.insert(key(12), V::empty()), None);
2162
2163            // The result should look like this:
2164            //                [6]
2165            //               /   \
2166            // [1, 2, 3, 4, 5]   [7, 8, 9, 10, 11, 12]
2167
2168            for i in 1..=12 {
2169                assert_eq!(btree.get(&key(i)), Some(V::empty()));
2170            }
2171
2172            // Remove node 6. Triggers case 2.b
2173            assert_eq!(btree.remove(&key(6)), Some(V::empty()));
2174            // The result should look like this:
2175            //                [7]
2176            //               /   \
2177            // [1, 2, 3, 4, 5]   [8, 9, 10, 11, 12]
2178            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            // There are three allocated nodes.
2198            assert_eq!(btree.allocator.num_allocated_chunks(), 3);
2199
2200            // Remove node 3. Triggers case 3.b
2201            assert_eq!(btree.remove(&key(3)), Some(V::empty()));
2202
2203            // Reload the btree to verify that we saved it correctly.
2204            let btree = BTreeMap::<K, V, _>::load(btree.into_memory());
2205
2206            // The result should look like this:
2207            //
2208            // [1, 2, 4, 5, 7, 8, 9, 10, 11, 12]
2209            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            // There is only one allocated node remaining.
2217            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            // Should now split a node.
2234            assert_eq!(btree.insert(key(12), V::empty()), None);
2235
2236            // The result should look like this:
2237            //                [6]
2238            //               /   \
2239            // [1, 2, 3, 4, 5]   [7, 8, 9, 10, 11, 12]
2240
2241            for i in 1..=12 {
2242                assert_eq!(btree.get(&key(i)), Some(V::empty()));
2243            }
2244
2245            // Remove node 6. Triggers case 2.b
2246            assert_eq!(btree.remove(&key(6)), Some(V::empty()));
2247
2248            // The result should look like this:
2249            //                [7]
2250            //               /   \
2251            // [1, 2, 3, 4, 5]   [8, 9, 10, 11, 12]
2252            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            // There are three allocated nodes.
2272            assert_eq!(btree.allocator.num_allocated_chunks(), 3);
2273
2274            // Remove node 10. Triggers case 3.b where we merge the right into the left.
2275            assert_eq!(btree.remove(&key(10)), Some(V::empty()));
2276
2277            // Reload the btree to verify that we saved it correctly.
2278            let btree = BTreeMap::<K, V, _>::load(btree.into_memory());
2279
2280            // The result should look like this:
2281            //
2282            // [1, 2, 3, 4, 5, 7, 8, 9, 11, 12], 10 is gone
2283            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            // There is only one allocated node remaining.
2291            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            // We've deallocated everything.
2315            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            // We've deallocated everything.
2340            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            // We've deallocated everything.
2366            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            // Reload the BTreeMap and verify the entry.
2383            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            // Remove all entries and verify the entry.
2390            for i in 0..n {
2391                assert_eq!(btree.remove(&key(i)), Some(value(i)));
2392            }
2393            assert_eq!(btree.len(), 0);
2394
2395            // Reload the BTreeMap and verify the entry.
2396            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            // Only even keys should be present.
2432            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            // Test prefixes that don't exist in the map.
2443            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    // Tests the case where the prefix is larger than all the entries in a leaf node.
2451    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            // Test a prefix that's larger than the value in the leaf node. Should be empty.
2457            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    // Tests the case where the prefix is larger than all the entries in an internal node.
2466    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            // The result should look like this:
2474            //                [6]
2475            //               /   \
2476            // [1, 2, 3, 4, 5]   [7, 8, 9, 10, 11, 12]
2477
2478            // Test a prefix that's larger than the key in the internal node.
2479            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            // The result should look like this:
2507            //                 [12]
2508            //                /    \
2509            // [1, 2, 3, 4, 11]    [13, 14, 21, 22, 23, 24]
2510
2511            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            // Tests a prefix that's smaller than the key in the internal node.
2520            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            // Tests a prefix that crosses several nodes.
2531            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            // Tests a prefix that's larger than the key in the internal node.
2542            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            // The result should look like this:
2578            //                 [14, 23]
2579            //                /    |   \
2580            // [1, 2, 3, 4, 12]    |   [24, 25, 26, 27, 28, 29]
2581            //                     |
2582            //           [16, 18, 19, 21, 22]
2583            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            // Tests a prefix that doesn't exist, but is in the middle of the root node.
2633            assert_eq!(collect_entry(btree.range(key(15)..key(16))), vec![]);
2634
2635            // Tests a prefix beginning in the middle of the tree and crossing several nodes.
2636            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            // Tests a prefix that crosses several nodes.
2652            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            // Tests a prefix that starts from a leaf node, then iterates through the root and right
2664            // sibling.
2665            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            // Check range [0, MID): should yield exactly the first MID entries.
2694            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            // Check range [MID, TOTAL): should yield the remaining entries.
2703            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            // The result should look like this:
2731            //                 [12]
2732            //                /    \
2733            // [1, 2, 3, 4, 11]    [13, 14, 21, 22, 23, 24]
2734
2735            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            // Tests an offset that has a keys somewhere in the range of keys of an internal node.
2754            assert_eq!(
2755                collect_entry(btree.range(key(13)..key(20))),
2756                vec![(key(13), value(700)), (key(14), value(800))]
2757            );
2758
2759            // Tests an offset that's larger than the key in the internal node.
2760            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            // The result should look like this:
2791            //                 [14, 23]
2792            //                /    |   \
2793            // [1, 2, 3, 4, 12]    |   [24, 25, 26, 27, 28, 29]
2794            //                     |
2795            //           [16, 18, 19, 21, 22]
2796            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            // Tests a offset that crosses several nodes.
2844            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            // Tests a offset that starts from a leaf node, then iterates through the root and right
2855            // sibling.
2856            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        // Reinitialize the BTree with larger keys and value sizes.
2894        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        // Still able to retrieve all the entries inserted.
2901        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        // Smaller key size
2917        let btree: BTreeMap<Blob<3>, Blob<3>, _> = BTreeMap::init(btree.into_memory());
2918        // Equal key size
2919        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        // Smaller key size
2935        let btree: BTreeMap<Blob<4>, Blob<2>, _> = BTreeMap::init(btree.into_memory());
2936        // Equal key size
2937        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    // A buggy implementation of storable where the max_size is smaller than the serialized size.
3020    #[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    // To generate the memory dump file for the current version:
3070    //   cargo test create_btreemap_dump_file -- --include-ignored
3071    #[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        // A type that is bounded.
3163        #[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        // Same as the above type, but unbounded.
3186        #[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        // Create a v1 btreemap with the bounded type.
3206        let mem = make_memory();
3207        let mut btree: BTreeMap<T, T, _> = BTreeMap::new_v1(mem);
3208        btree.insert(T, T);
3209
3210        // Migrate to v2 and the unbounded type.
3211        let btree: BTreeMap<T2, T2, _> = BTreeMap::init(btree.into_memory());
3212        btree.save_header();
3213
3214        // Reload the BTree again and try to read the value.
3215        let btree: BTreeMap<T2, T2, _> = BTreeMap::init(btree.into_memory());
3216        assert_eq!(btree.get(&T2), Some(T2));
3217
3218        // Reload the BTree again with bounded type.
3219        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        // No allocated chunks yet.
3275        assert_eq!(btree.allocator.num_allocated_chunks(), 0);
3276
3277        // Insert and remove an entry that's large and requires overflow pages.
3278        btree.insert(vec![0; 10_000], vec![]);
3279
3280        // At least two chunks should be allocated.
3281        // One for the node itself and at least one overflow page.
3282        assert!(btree.allocator.num_allocated_chunks() >= 2);
3283        btree.remove(&vec![0; 10_000]);
3284
3285        // All chunks have been deallocated.
3286        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        // No allocated chunks yet.
3295        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        // All chunks have been deallocated.
3308        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            // Large keys are stored so that each node overflows.
3318            assert_eq!(btree.insert(vec![i; 10_000], ()), None);
3319        }
3320
3321        // Should now split a node.
3322        assert_eq!(btree.insert(vec![0; 10_000], ()), None);
3323
3324        // The btree should look like this:
3325        //                    [6]
3326        //                   /   \
3327        // [0, 1, 2, 3, 4, 5]     [7, 8, 9, 10, 11]
3328        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        // Remove the element in the root.
3334        btree.remove(&vec![6; 10_000]);
3335
3336        // The btree should look like this:
3337        //                 [5]
3338        //                /   \
3339        // [0, 1, 2, 3, 4]     [7, 8, 9, 10, 11]
3340        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        // Remove the element in the root. This triggers the case where the root
3346        // node is deallocated and the children are merged into a single node.
3347        btree.remove(&vec![5; 10_000]);
3348
3349        // The btree should look like this:
3350        //      [0, 1, 2, 3, 4, 7, 8, 9, 10, 11]
3351        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        // Delete everything else.
3370        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}