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 from_bytes(bytes: Cow<[u8]>) -> Self {
182///         // TODO: Convert bytes back to your struct
183///         let (id, name) = (0, "".to_string());
184///         User { id, name }
185///     }
186///
187///     // Types can be bounded or unbounded:
188///     // - Use Bound::Unbounded if the size can vary or isn't known in advance (recommended for most cases)
189///     // - Use Bound::Bounded if you know the maximum size and want memory optimization
190///     const BOUND: Bound = Bound::Unbounded;
191/// }
192///
193/// // Now you can use your custom type in a BTreeMap
194/// let mut map: BTreeMap<u64, User, _> = BTreeMap::init(DefaultMemoryImpl::default());
195///
196/// let user = User { id: 1, name: "Alice".to_string() };
197/// map.insert(1, user);
198///
199/// // Retrieving the custom type
200/// if let Some(user) = map.get(&1) {
201///     println!("Found user: {}", user.name);
202/// }
203/// ```
204///
205/// ### Bounded vs Unbounded Types
206///
207/// When implementing `Storable`, you must specify whether your type is bounded or unbounded:
208///
209/// - **Unbounded (`Bound::Unbounded`)**:
210///   - Use when your type's serialized size can vary or has no fixed maximum
211///   - Recommended for most custom types, especially those containing Strings or Vecs
212///   - Example: `const BOUND: Bound = Bound::Unbounded;`
213///
214/// - **Bounded (`Bound::Bounded{ max_size, is_fixed_size }`)**:
215///   - Use when you know the maximum serialized size of your type
216///   - Enables memory optimizations in the BTreeMap
217///   - Example: `const BOUND: Bound = Bound::Bounded { max_size: 100, is_fixed_size: false };`
218///   - For types with truly fixed size (like primitive types), set `is_fixed_size: true`
219///
220/// If unsure, use `Bound::Unbounded` as it's the safer choice.
221///
222/// # Warning
223///
224/// Once you've deployed with a bounded type, you cannot increase its `max_size` in
225/// future versions without risking data corruption. You can, however, migrate from a bounded type
226/// to an unbounded type if needed. For evolving data structures, prefer `Bound::Unbounded`.
227pub struct BTreeMap<K, V, M>
228where
229    K: Storable + Ord + Clone,
230    V: Storable,
231    M: Memory,
232{
233    // The address of the root node. If a root node doesn't exist, the address
234    // is set to NULL.
235    root_addr: Address,
236
237    version: Version,
238
239    // An allocator used for managing memory and allocating nodes.
240    allocator: Allocator<M>,
241
242    // The number of elements in the map.
243    length: u64,
244
245    // A marker to communicate to the Rust compiler that we own these types.
246    _phantom: PhantomData<(K, V)>,
247}
248
249#[derive(PartialEq, Debug)]
250/// The packed header size must be <= ALLOCATOR_OFFSET.
251struct BTreeHeader {
252    version: Version,
253    root_addr: Address,
254    length: u64,
255    // Reserved bytes for future extensions
256}
257
258impl<K, V, M> BTreeMap<K, V, M>
259where
260    K: Storable + Ord + Clone,
261    V: Storable,
262    M: Memory,
263{
264    /// Initializes a `BTreeMap`.
265    ///
266    /// If the memory provided already contains a `BTreeMap`, then that
267    /// map is loaded. Otherwise, a new `BTreeMap` instance is created.
268    pub fn init(memory: M) -> Self {
269        if memory.size() == 0 {
270            // Memory is empty. Create a new map.
271            return BTreeMap::new(memory);
272        }
273
274        // Check if the magic in the memory corresponds to a BTreeMap.
275        let mut dst = vec![0; 3];
276        memory.read(0, &mut dst);
277        if dst != MAGIC {
278            // No BTreeMap found. Create a new instance.
279            BTreeMap::new(memory)
280        } else {
281            // The memory already contains a BTreeMap. Load it.
282            BTreeMap::load(memory)
283        }
284    }
285
286    /// Initializes a v1 `BTreeMap`.
287    ///
288    /// This is exposed only in testing.
289    #[cfg(test)]
290    pub fn init_v1(memory: M) -> Self {
291        if memory.size() == 0 {
292            // Memory is empty. Create a new map.
293            return BTreeMap::new_v1(memory);
294        }
295
296        // Check if the magic in the memory corresponds to a BTreeMap.
297        let mut dst = vec![0; 3];
298        memory.read(0, &mut dst);
299        if dst != MAGIC {
300            // No BTreeMap found. Create a new instance.
301            BTreeMap::new_v1(memory)
302        } else {
303            // The memory already contains a BTreeMap. Load it, making sure
304            // we don't migrate the BTreeMap to v2.
305            BTreeMap::load_helper(memory, false)
306        }
307    }
308
309    /// Creates a new instance a `BTreeMap`.
310    ///
311    /// The given `memory` is assumed to be exclusively reserved for this data
312    /// structure and that it starts at address zero. Typically `memory` will
313    /// be an instance of `RestrictedMemory`.
314    ///
315    /// When initialized, the data structure has the following memory layout:
316    ///
317    ///    |  BTreeHeader  |  Allocator | ... free memory for nodes |
318    ///
319    /// See `Allocator` for more details on its own memory layout.
320    pub fn new(memory: M) -> Self {
321        let page_size = match (K::BOUND, V::BOUND) {
322            // The keys and values are both bounded.
323            (
324                StorableBound::Bounded {
325                    max_size: max_key_size,
326                    ..
327                },
328                StorableBound::Bounded {
329                    max_size: max_value_size,
330                    ..
331                },
332            ) => {
333                // Get the maximum possible node size.
334                let max_node_size = Node::<K>::max_size(max_key_size, max_value_size).get();
335
336                // A node can have at most 11 entries, and an analysis has shown that ~70% of all
337                // nodes have <= 8 entries. We can therefore use a page size that's 8/11 the
338                // maximum size with little performance difference but with a significant storage
339                // saving. We round the 8/11 to be 3/4.
340                PageSize::Value((max_node_size * 3 / 4) as u32)
341            }
342            // Use a default page size.
343            _ => PageSize::Value(DEFAULT_PAGE_SIZE),
344        };
345
346        let btree = Self {
347            root_addr: NULL,
348            allocator: Allocator::new(
349                memory,
350                Address::from(ALLOCATOR_OFFSET as u64),
351                page_size.get().into(),
352            ),
353            version: Version::V2(page_size),
354            length: 0,
355            _phantom: PhantomData,
356        };
357
358        btree.save_header();
359        btree
360    }
361
362    /// Create a v1 instance of the BTree.
363    ///
364    /// This is only exposed for testing.
365    #[cfg(test)]
366    pub fn new_v1(memory: M) -> Self {
367        let max_key_size = K::BOUND.max_size();
368        let max_value_size = V::BOUND.max_size();
369
370        let btree = Self {
371            root_addr: NULL,
372            allocator: Allocator::new(
373                memory,
374                Address::from(ALLOCATOR_OFFSET as u64),
375                Node::<K>::max_size(max_key_size, max_value_size),
376            ),
377            version: Version::V1(DerivedPageSize {
378                max_key_size,
379                max_value_size,
380            }),
381            length: 0,
382            _phantom: PhantomData,
383        };
384
385        btree.save_header();
386        btree
387    }
388
389    /// Loads the map from memory.
390    pub fn load(memory: M) -> Self {
391        Self::load_helper(memory, true)
392    }
393
394    /// Loads the map from memory, potentially migrating the map from V1 to V2.
395    fn load_helper(memory: M, migrate_to_v2: bool) -> Self {
396        // Read the header from memory.
397        let header = Self::read_header(&memory);
398
399        let version = match header.version {
400            Version::V1(derived_page_size) => {
401                // Migrate to V2 if flag is enabled.
402                if migrate_to_v2 {
403                    Version::V2(PageSize::Derived(derived_page_size))
404                } else {
405                    // Assert that the bounds are correct.
406                    let max_key_size = derived_page_size.max_key_size;
407                    let max_value_size = derived_page_size.max_value_size;
408
409                    assert!(
410                        K::BOUND.max_size() <= max_key_size,
411                        "max_key_size must be <= {max_key_size}",
412                    );
413
414                    assert!(
415                        V::BOUND.max_size() <= max_value_size,
416                        "max_value_size must be <= {max_value_size}"
417                    );
418
419                    Version::V1(derived_page_size)
420                }
421            }
422            other => other,
423        };
424
425        let allocator_addr = Address::from(ALLOCATOR_OFFSET as u64);
426        Self {
427            root_addr: header.root_addr,
428            allocator: Allocator::load(memory, allocator_addr),
429            version,
430            length: header.length,
431            _phantom: PhantomData,
432        }
433    }
434
435    /// Reads the header from the specified memory.
436    fn read_header(memory: &M) -> BTreeHeader {
437        // Read the header
438        let mut buf = [0; PACKED_HEADER_SIZE];
439        memory.read(0, &mut buf);
440
441        assert_eq!(&buf[0..3], MAGIC, "Bad magic.");
442
443        match buf[3] {
444            LAYOUT_VERSION => {
445                // Deserialize the fields
446                BTreeHeader {
447                    version: Version::V1(DerivedPageSize {
448                        max_key_size: u32::from_le_bytes(buf[4..8].try_into().unwrap()),
449                        max_value_size: u32::from_le_bytes(buf[8..12].try_into().unwrap()),
450                    }),
451                    root_addr: Address::from(u64::from_le_bytes(buf[12..20].try_into().unwrap())),
452                    length: u64::from_le_bytes(buf[20..28].try_into().unwrap()),
453                }
454            }
455            LAYOUT_VERSION_2 => {
456                // Load the page size.
457                let page_size = {
458                    // Page sizes can be stored either as a direct value or as max/value sizes.
459                    let a = u32::from_le_bytes(buf[4..8].try_into().unwrap());
460                    let b = u32::from_le_bytes(buf[8..12].try_into().unwrap());
461
462                    if b == PAGE_SIZE_VALUE_MARKER {
463                        // Page size is stored as a direct value
464                        PageSize::Value(a)
465                    } else {
466                        // Page size is stored as a derived value.
467                        PageSize::Derived(DerivedPageSize {
468                            max_key_size: a,
469                            max_value_size: b,
470                        })
471                    }
472                };
473
474                // Deserialize the fields
475                BTreeHeader {
476                    version: Version::V2(page_size),
477                    root_addr: Address::from(u64::from_le_bytes(buf[12..20].try_into().unwrap())),
478                    length: u64::from_le_bytes(buf[20..28].try_into().unwrap()),
479                }
480            }
481            version => {
482                panic!("Unsupported version: {version}.");
483            }
484        }
485    }
486
487    /// Inserts a key-value pair into the map.
488    ///
489    /// The previous value of the key, if present, is returned.
490    ///
491    /// PRECONDITION:
492    ///   key.to_bytes().len() <= max_size(Key)
493    ///   value.to_bytes().len() <= max_size(Value)
494    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
495        let value = value.to_bytes_checked().into_owned();
496
497        let root = if self.root_addr == NULL {
498            // No root present. Allocate one.
499            let node = self.allocate_node(NodeType::Leaf);
500            self.root_addr = node.address();
501            self.save_header();
502            node
503        } else {
504            // Load the root from memory.
505            let mut root = self.load_node(self.root_addr);
506
507            // Check if the key already exists in the root.
508            if let Ok(idx) = root.search(&key, self.memory()) {
509                // Key found, replace its value and return the old one.
510                return Some(V::from_bytes(Cow::Owned(
511                    self.update_value(&mut root, idx, value),
512                )));
513            }
514
515            // If the root is full, we need to introduce a new node as the root.
516            //
517            // NOTE: In the case where we are overwriting an existing key, then introducing
518            // a new root node isn't strictly necessary. However, that's a micro-optimization
519            // that adds more complexity than it's worth.
520            if root.is_full() {
521                // The root is full. Allocate a new node that will be used as the new root.
522                let mut new_root = self.allocate_node(NodeType::Internal);
523
524                // The new root has the old root as its only child.
525                new_root.push_child(self.root_addr);
526
527                // Update the root address.
528                self.root_addr = new_root.address();
529                self.save_header();
530
531                // Split the old (full) root.
532                self.split_child(&mut new_root, 0);
533
534                new_root
535            } else {
536                root
537            }
538        };
539
540        self.insert_nonfull(root, key, value)
541            .map(Cow::Owned)
542            .map(V::from_bytes)
543    }
544
545    /// Inserts an entry into a node that is *not full*.
546    fn insert_nonfull(&mut self, mut node: Node<K>, key: K, value: Vec<u8>) -> Option<Vec<u8>> {
547        // We're guaranteed by the caller that the provided node is not full.
548        assert!(!node.is_full());
549
550        // Look for the key in the node.
551        match node.search(&key, self.memory()) {
552            Ok(idx) => {
553                // Key found, replace its value and return the old one.
554                Some(self.update_value(&mut node, idx, value))
555            }
556            Err(idx) => {
557                // The key isn't in the node. `idx` is where that key should be inserted.
558
559                match node.node_type() {
560                    NodeType::Leaf => {
561                        // The node is a non-full leaf.
562                        // Insert the entry at the proper location.
563                        node.insert_entry(idx, (key, value));
564                        self.save_node(&mut node);
565
566                        // Update the length.
567                        self.length += 1;
568                        self.save_header();
569
570                        // No previous value to return.
571                        None
572                    }
573                    NodeType::Internal => {
574                        // The node is an internal node.
575                        // Load the child that we should add the entry to.
576                        let mut child = self.load_node(node.child(idx));
577
578                        if child.is_full() {
579                            // Check if the key already exists in the child.
580                            if let Ok(idx) = child.search(&key, self.memory()) {
581                                // Key found, replace its value and return the old one.
582                                return Some(self.update_value(&mut child, idx, value));
583                            }
584
585                            // The child is full. Split the child.
586                            self.split_child(&mut node, idx);
587
588                            // The children have now changed. Search again for
589                            // the child where we need to store the entry in.
590                            let idx = node.search(&key, self.memory()).unwrap_or_else(|idx| idx);
591                            child = self.load_node(node.child(idx));
592                        }
593
594                        // The child should now be not full.
595                        assert!(!child.is_full());
596
597                        self.insert_nonfull(child, key, value)
598                    }
599                }
600            }
601        }
602    }
603
604    /// Takes as input a nonfull internal `node` and index to its full child, then
605    /// splits this child into two, adding an additional child to `node`.
606    ///
607    /// Example:
608    /// ```ignore
609    ///                          [ ... M   Y ... ]
610    ///                                  |
611    ///                 [ N  O  P  Q  R  S  T  U  V  W  X ]
612    /// ```
613    ///
614    /// After splitting becomes:
615    /// ```ignore
616    ///                         [ ... M  S  Y ... ]
617    ///                                 / \
618    ///                [ N  O  P  Q  R ]   [ T  U  V  W  X ]
619    /// ```
620    ///
621    fn split_child(&mut self, node: &mut Node<K>, full_child_idx: usize) {
622        // The node must not be full.
623        assert!(!node.is_full());
624
625        // The node's child must be full.
626        let mut full_child = self.load_node(node.child(full_child_idx));
627        assert!(full_child.is_full());
628
629        // Create a sibling to this full child (which has to be the same type).
630        let mut sibling = self.allocate_node(full_child.node_type());
631        assert_eq!(sibling.node_type(), full_child.node_type());
632
633        // Add sibling as a new child in the node.
634        node.insert_child(full_child_idx + 1, sibling.address());
635
636        let (median_key, median_value) = full_child.split(&mut sibling, self.memory());
637
638        node.insert_entry(full_child_idx, (median_key, median_value));
639
640        self.save_node(&mut sibling);
641        self.save_node(&mut full_child);
642        self.save_node(node);
643    }
644
645    /// Returns the value for the given key, if it exists.
646    pub fn get(&self, key: &K) -> Option<V> {
647        if self.root_addr == NULL {
648            return None;
649        }
650        self.traverse(self.root_addr, key, |node, idx| {
651            node.into_entry(idx, self.memory()).1 // Extract value.
652        })
653        .map(Cow::Owned)
654        .map(V::from_bytes)
655    }
656
657    /// Returns true if the key exists.
658    pub fn contains_key(&self, key: &K) -> bool {
659        // An empty closure returns Some(()) if the key is found.
660        self.root_addr != NULL && self.traverse(self.root_addr, key, |_, _| ()).is_some()
661    }
662
663    /// Recursively traverses from `node_addr`, invoking `f` if `key` is found. Stops at a leaf if not.
664    fn traverse<F, R>(&self, node_addr: Address, key: &K, f: F) -> Option<R>
665    where
666        F: Fn(Node<K>, usize) -> R + Clone,
667    {
668        let node = self.load_node(node_addr);
669        // Look for the key in the current node.
670        match node.search(key, self.memory()) {
671            Ok(idx) => Some(f(node, idx)), // Key found: apply `f`.
672            Err(idx) => match node.node_type() {
673                NodeType::Leaf => None, // At a leaf: key not present.
674                NodeType::Internal => self.traverse(node.child(idx), key, f), // Continue search in child.
675            },
676        }
677    }
678
679    /// Returns `true` if the map contains no elements.
680    pub fn is_empty(&self) -> bool {
681        self.length == 0
682    }
683
684    /// Returns the number of elements in the map.
685    pub fn len(&self) -> u64 {
686        self.length
687    }
688
689    /// Returns the underlying memory.
690    pub fn into_memory(self) -> M {
691        self.allocator.into_memory()
692    }
693
694    /// Removes all elements from the map.
695    #[deprecated(since = "0.6.3", note = "please use `clear_new` instead")]
696    pub fn clear(self) -> Self {
697        let mem = self.allocator.into_memory();
698        Self::new(mem)
699    }
700
701    /// Removes all elements from the map.
702    pub fn clear_new(&mut self) {
703        self.root_addr = NULL;
704        self.length = 0;
705        self.allocator.clear();
706        self.save_header();
707    }
708
709    /// Returns the first key-value pair in the map. The key in this
710    /// pair is the minimum key in the map.
711    pub fn first_key_value(&self) -> Option<(K, V)> {
712        if self.root_addr == NULL {
713            return None;
714        }
715        let root = self.load_node(self.root_addr);
716        let (k, encoded_v) = root.get_min(self.memory());
717        Some((k, V::from_bytes(Cow::Owned(encoded_v))))
718    }
719
720    /// Returns the last key-value pair in the map. The key in this
721    /// pair is the maximum key in the map.
722    pub fn last_key_value(&self) -> Option<(K, V)> {
723        if self.root_addr == NULL {
724            return None;
725        }
726        let root = self.load_node(self.root_addr);
727        let (k, encoded_v) = root.get_max(self.memory());
728        Some((k, V::from_bytes(Cow::Owned(encoded_v))))
729    }
730
731    fn memory(&self) -> &M {
732        self.allocator.memory()
733    }
734
735    fn allocator_mut(&mut self) -> &mut Allocator<M> {
736        &mut self.allocator
737    }
738
739    /// Removes a key from the map, returning the previous value at the key if it exists.
740    pub fn remove(&mut self, key: &K) -> Option<V> {
741        if self.root_addr == NULL {
742            return None;
743        }
744
745        let root_node = self.load_node(self.root_addr);
746        self.remove_helper(root_node, key)
747            .map(Cow::Owned)
748            .map(V::from_bytes)
749    }
750
751    /// Removes and returns the last element in the map. The key of this element is the maximum key that was in the map
752    pub fn pop_last(&mut self) -> Option<(K, V)> {
753        if self.root_addr == NULL {
754            return None;
755        }
756
757        let root = self.load_node(self.root_addr);
758        let (max_key, _) = root.get_max(self.memory());
759        self.remove_helper(root, &max_key)
760            .map(|v| (max_key, V::from_bytes(Cow::Owned(v))))
761    }
762
763    /// Removes and returns the first element in the map. The key of this element is the minimum key that was in the map
764    pub fn pop_first(&mut self) -> Option<(K, V)> {
765        if self.root_addr == NULL {
766            return None;
767        }
768
769        let root = self.load_node(self.root_addr);
770        let (min_key, _) = root.get_min(self.memory());
771        self.remove_helper(root, &min_key)
772            .map(|v| (min_key, V::from_bytes(Cow::Owned(v))))
773    }
774
775    /// A helper method for recursively removing a key from the B-tree.
776    fn remove_helper(&mut self, mut node: Node<K>, key: &K) -> Option<Vec<u8>> {
777        if node.address() != self.root_addr {
778            // We're guaranteed that whenever this method is called an entry can be
779            // removed from the node without it needing to be merged into a sibling.
780            // This strengthened condition allows us to delete an entry in a single
781            // pass most of the time without having to back up.
782            assert!(node.can_remove_entry_without_merging());
783        }
784
785        match node.node_type() {
786            NodeType::Leaf => {
787                match node.search(key, self.memory()) {
788                    Ok(idx) => {
789                        // Case 1: The node is a leaf node and the key exists in it.
790                        // This is the simplest case. The key is removed from the leaf.
791                        let value = node.remove_entry(idx, self.memory()).1;
792                        self.length -= 1;
793
794                        if node.entries_len() == 0 {
795                            assert_eq!(
796                                node.address(), self.root_addr,
797                                "Removal can only result in an empty leaf node if that node is the root"
798                            );
799
800                            // Deallocate the empty node.
801                            self.deallocate_node(node);
802                            self.root_addr = NULL;
803                        } else {
804                            self.save_node(&mut node);
805                        }
806
807                        self.save_header();
808                        Some(value)
809                    }
810                    _ => None, // Key not found.
811                }
812            }
813            NodeType::Internal => {
814                match node.search(key, self.memory()) {
815                    Ok(idx) => {
816                        // Case 2: The node is an internal node and the key exists in it.
817
818                        let left_child = self.load_node(node.child(idx));
819                        if left_child.can_remove_entry_without_merging() {
820                            // Case 2.a: A key can be removed from the left child without merging.
821                            //
822                            //                       parent
823                            //                  [..., key, ...]
824                            //                       /   \
825                            //            [left child]   [...]
826                            //           /            \
827                            //        [...]         [..., key predecessor]
828                            //
829                            // In this case, we replace `key` with the key's predecessor from the
830                            // left child's subtree, then we recursively delete the key's
831                            // predecessor for the following end result:
832                            //
833                            //                       parent
834                            //            [..., key predecessor, ...]
835                            //                       /   \
836                            //            [left child]   [...]
837                            //           /            \
838                            //        [...]          [...]
839
840                            // Recursively delete the predecessor.
841                            // TODO(EXC-1034): Do this in a single pass.
842                            let predecessor = left_child.get_max(self.memory());
843                            self.remove_helper(left_child, &predecessor.0)?;
844
845                            // Replace the `key` with its predecessor.
846                            let (_, old_value) = node.swap_entry(idx, predecessor, self.memory());
847
848                            // Save the parent node.
849                            self.save_node(&mut node);
850                            return Some(old_value);
851                        }
852
853                        let right_child = self.load_node(node.child(idx + 1));
854                        if right_child.can_remove_entry_without_merging() {
855                            // Case 2.b: A key can be removed from the right child without merging.
856                            //
857                            //                       parent
858                            //                  [..., key, ...]
859                            //                       /   \
860                            //                   [...]   [right child]
861                            //                          /             \
862                            //              [key successor, ...]     [...]
863                            //
864                            // In this case, we replace `key` with the key's successor from the
865                            // right child's subtree, then we recursively delete the key's
866                            // successor for the following end result:
867                            //
868                            //                       parent
869                            //            [..., key successor, ...]
870                            //                       /   \
871                            //                  [...]   [right child]
872                            //                           /            \
873                            //                        [...]          [...]
874
875                            // Recursively delete the successor.
876                            // TODO(EXC-1034): Do this in a single pass.
877                            let successor = right_child.get_min(self.memory());
878                            self.remove_helper(right_child, &successor.0)?;
879
880                            // Replace the `key` with its successor.
881                            let (_, old_value) = node.swap_entry(idx, successor, self.memory());
882
883                            // Save the parent node.
884                            self.save_node(&mut node);
885                            return Some(old_value);
886                        }
887
888                        // Case 2.c: Both the left and right child are at their minimum sizes.
889                        //
890                        //                       parent
891                        //                  [..., key, ...]
892                        //                       /   \
893                        //            [left child]   [right child]
894                        //
895                        // In this case, we merge (left child, key, right child) into a single
896                        // node. The result will look like this:
897                        //
898                        //                       parent
899                        //                     [...  ...]
900                        //                         |
901                        //          [left child, `key`, right child] <= new child
902                        //
903                        // We then recurse on this new child to delete `key`.
904                        //
905                        // If `parent` becomes empty (which can only happen if it's the root),
906                        // then `parent` is deleted and `new_child` becomes the new root.
907                        assert!(left_child.at_minimum());
908                        assert!(right_child.at_minimum());
909
910                        // Merge the right child into the left child.
911                        let mut new_child = self.merge(
912                            right_child,
913                            left_child,
914                            node.remove_entry(idx, self.memory()),
915                        );
916
917                        // Remove the right child from the parent node.
918                        node.remove_child(idx + 1);
919
920                        if node.entries_len() == 0 {
921                            // Can only happen if this node is root.
922                            assert_eq!(node.address(), self.root_addr);
923                            assert_eq!(node.child(0), new_child.address());
924                            assert_eq!(node.children_len(), 1);
925
926                            self.root_addr = new_child.address();
927
928                            // Deallocate the root node.
929                            self.deallocate_node(node);
930                            self.save_header();
931                        } else {
932                            self.save_node(&mut node);
933                        }
934
935                        self.save_node(&mut new_child);
936
937                        // Recursively delete the key.
938                        self.remove_helper(new_child, key)
939                    }
940                    Err(idx) => {
941                        // Case 3: The node is an internal node and the key does NOT exist in it.
942
943                        // If the key does exist in the tree, it will exist in the subtree at index
944                        // `idx`.
945                        let mut child = self.load_node(node.child(idx));
946
947                        if child.can_remove_entry_without_merging() {
948                            // The child has enough nodes. Recurse to delete the `key` from the
949                            // `child`.
950                            return self.remove_helper(child, key);
951                        }
952
953                        // An entry can't be removed from the child without merging.
954                        // See if it has a sibling where an entry can be removed without merging.
955                        let mut left_sibling = if idx > 0 {
956                            Some(self.load_node(node.child(idx - 1)))
957                        } else {
958                            None
959                        };
960
961                        let mut right_sibling = if idx + 1 < node.children_len() {
962                            Some(self.load_node(node.child(idx + 1)))
963                        } else {
964                            None
965                        };
966
967                        if let Some(ref mut left_sibling) = left_sibling {
968                            if left_sibling.can_remove_entry_without_merging() {
969                                // Case 3.a (left):
970                                // A key can be removed from the left child without merging.
971                                //
972                                //                            [d] (parent)
973                                //                           /   \
974                                //  (left sibling) [a, b, c]     [e, f] (child)
975                                //                         \
976                                //                         [c']
977                                //
978                                // In this case, we move a key down from the parent into the child
979                                // and move a key from the left sibling up into the parent
980                                // resulting in the following tree:
981                                //
982                                //                            [c] (parent)
983                                //                           /   \
984                                //       (left sibling) [a, b]   [d, e, f] (child)
985                                //                              /
986                                //                            [c']
987                                //
988                                // We then recurse to delete the key from the child.
989
990                                // Remove the last entry from the left sibling.
991                                let (left_sibling_key, left_sibling_value) =
992                                    left_sibling.pop_entry(self.memory()).unwrap();
993
994                                // Replace the parent's entry with the one from the left sibling.
995                                let (parent_key, parent_value) = node.swap_entry(
996                                    idx - 1,
997                                    (left_sibling_key, left_sibling_value),
998                                    self.memory(),
999                                );
1000
1001                                // Move the entry from the parent into the child.
1002                                child.insert_entry(0, (parent_key, parent_value));
1003
1004                                // Move the last child from left sibling into child.
1005                                if let Some(last_child) = left_sibling.pop_child() {
1006                                    assert_eq!(left_sibling.node_type(), NodeType::Internal);
1007                                    assert_eq!(child.node_type(), NodeType::Internal);
1008
1009                                    child.insert_child(0, last_child);
1010                                } else {
1011                                    assert_eq!(left_sibling.node_type(), NodeType::Leaf);
1012                                    assert_eq!(child.node_type(), NodeType::Leaf);
1013                                }
1014
1015                                self.save_node(left_sibling);
1016                                self.save_node(&mut child);
1017                                self.save_node(&mut node);
1018                                return self.remove_helper(child, key);
1019                            }
1020                        }
1021
1022                        if let Some(right_sibling) = &mut right_sibling {
1023                            if right_sibling.can_remove_entry_without_merging() {
1024                                // Case 3.a (right):
1025                                // A key can be removed from the right child without merging.
1026                                //
1027                                //                            [c] (parent)
1028                                //                           /   \
1029                                //             (child) [a, b]     [d, e, f] (right sibling)
1030                                //                               /
1031                                //                            [d']
1032                                //
1033                                // In this case, we move a key down from the parent into the child
1034                                // and move a key from the right sibling up into the parent
1035                                // resulting in the following tree:
1036                                //
1037                                //                            [d] (parent)
1038                                //                           /   \
1039                                //          (child) [a, b, c]     [e, f] (right sibling)
1040                                //                          \
1041                                //                           [d']
1042                                //
1043                                // We then recurse to delete the key from the child.
1044
1045                                // Remove the first entry from the right sibling.
1046                                let (right_sibling_key, right_sibling_value) =
1047                                    right_sibling.remove_entry(0, self.memory());
1048
1049                                // Replace the parent's entry with the one from the right sibling.
1050                                let parent_entry = node.swap_entry(
1051                                    idx,
1052                                    (right_sibling_key, right_sibling_value),
1053                                    self.memory(),
1054                                );
1055
1056                                // Move the entry from the parent into the child.
1057                                child.push_entry(parent_entry);
1058
1059                                // Move the first child of right_sibling into `child`.
1060                                match right_sibling.node_type() {
1061                                    NodeType::Internal => {
1062                                        assert_eq!(child.node_type(), NodeType::Internal);
1063                                        child.push_child(right_sibling.remove_child(0));
1064                                    }
1065                                    NodeType::Leaf => {
1066                                        assert_eq!(child.node_type(), NodeType::Leaf);
1067                                    }
1068                                }
1069
1070                                self.save_node(right_sibling);
1071                                self.save_node(&mut child);
1072                                self.save_node(&mut node);
1073                                return self.remove_helper(child, key);
1074                            }
1075                        }
1076
1077                        // Case 3.b: Both the left and right siblings are at their minimum sizes.
1078
1079                        if let Some(left_sibling) = left_sibling {
1080                            // Merge child into left sibling if it exists.
1081
1082                            assert!(left_sibling.at_minimum());
1083                            let left_sibling = self.merge(
1084                                child,
1085                                left_sibling,
1086                                node.remove_entry(idx - 1, self.memory()),
1087                            );
1088                            // Removing child from parent.
1089                            node.remove_child(idx);
1090
1091                            if node.entries_len() == 0 {
1092                                let node_address = node.address();
1093                                self.deallocate_node(node);
1094
1095                                if node_address == self.root_addr {
1096                                    // Update the root.
1097                                    self.root_addr = left_sibling.address();
1098                                    self.save_header();
1099                                }
1100                            } else {
1101                                self.save_node(&mut node);
1102                            }
1103
1104                            return self.remove_helper(left_sibling, key);
1105                        }
1106
1107                        if let Some(right_sibling) = right_sibling {
1108                            // Merge child into right sibling.
1109
1110                            assert!(right_sibling.at_minimum());
1111                            let right_sibling = self.merge(
1112                                child,
1113                                right_sibling,
1114                                node.remove_entry(idx, self.memory()),
1115                            );
1116
1117                            // Removing child from parent.
1118                            node.remove_child(idx);
1119
1120                            if node.entries_len() == 0 {
1121                                let node_address = node.address();
1122                                self.deallocate_node(node);
1123
1124                                if node_address == self.root_addr {
1125                                    // Update the root.
1126                                    self.root_addr = right_sibling.address();
1127                                    self.save_header();
1128                                }
1129                            } else {
1130                                self.save_node(&mut node);
1131                            }
1132
1133                            return self.remove_helper(right_sibling, key);
1134                        }
1135
1136                        unreachable!("At least one of the siblings must exist.");
1137                    }
1138                }
1139            }
1140        }
1141    }
1142
1143    /// Returns an iterator over the entries of the map, sorted by key.
1144    ///
1145    /// # Example
1146    ///
1147    /// ```rust
1148    /// use ic_stable_structures::{BTreeMap, DefaultMemoryImpl};
1149    /// let mut map: BTreeMap<u64, String, _> = BTreeMap::init(DefaultMemoryImpl::default());
1150    ///
1151    /// map.insert(1, "one".to_string());
1152    /// map.insert(2, "two".to_string());
1153    ///
1154    /// for (key, value) in map.iter() {
1155    ///     println!("{}: {}", key, value);
1156    /// }
1157    /// ```
1158    pub fn iter(&self) -> Iter<K, V, M> {
1159        self.iter_internal().into()
1160    }
1161
1162    /// Returns an iterator over the entries in the map where keys
1163    /// belong to the specified range.
1164    ///
1165    /// # Example
1166    ///
1167    /// ```rust
1168    /// use ic_stable_structures::{BTreeMap, DefaultMemoryImpl};
1169    /// use std::ops::Bound;
1170    ///
1171    /// let mut map: BTreeMap<u64, String, _> = BTreeMap::init(DefaultMemoryImpl::default());
1172    /// map.insert(1, "one".to_string());
1173    /// map.insert(2, "two".to_string());
1174    /// map.insert(3, "three".to_string());
1175    ///
1176    /// // Get entries with keys between 1 and 3 (inclusive)
1177    /// for (key, value) in map.range((Bound::Included(1), Bound::Included(3))) {
1178    ///     println!("{}: {}", key, value);
1179    /// }
1180    /// ```
1181    pub fn range(&self, key_range: impl RangeBounds<K>) -> Iter<K, V, M> {
1182        self.range_internal(key_range).into()
1183    }
1184
1185    /// Returns an iterator pointing to the first element below the given bound.
1186    /// Returns an empty iterator if there are no keys below the given bound.
1187    pub fn iter_upper_bound(&self, bound: &K) -> Iter<K, V, M> {
1188        if let Some((start_key, _)) = self.range(..bound).next_back() {
1189            IterInternal::new_in_range(self, (Bound::Included(start_key), Bound::Unbounded)).into()
1190        } else {
1191            IterInternal::null(self).into()
1192        }
1193    }
1194
1195    /// Returns an iterator over the keys of the map.
1196    pub fn keys(&self) -> KeysIter<K, V, M> {
1197        self.iter_internal().into()
1198    }
1199
1200    /// Returns an iterator over the keys of the map which belong to the specified range.
1201    pub fn keys_range(&self, key_range: impl RangeBounds<K>) -> KeysIter<K, V, M> {
1202        self.range_internal(key_range).into()
1203    }
1204
1205    /// Returns an iterator over the values of the map, sorted by key.
1206    pub fn values(&self) -> ValuesIter<K, V, M> {
1207        self.iter_internal().into()
1208    }
1209
1210    /// Returns an iterator over the values of the map where keys
1211    /// belong to the specified range.
1212    pub fn values_range(&self, key_range: impl RangeBounds<K>) -> ValuesIter<K, V, M> {
1213        self.range_internal(key_range).into()
1214    }
1215
1216    fn iter_internal(&self) -> IterInternal<K, V, M> {
1217        IterInternal::new(self)
1218    }
1219
1220    fn range_internal(&self, key_range: impl RangeBounds<K>) -> IterInternal<K, V, M> {
1221        if self.root_addr == NULL {
1222            // Map is empty.
1223            return IterInternal::null(self);
1224        }
1225
1226        let range = (
1227            key_range.start_bound().cloned(),
1228            key_range.end_bound().cloned(),
1229        );
1230
1231        IterInternal::new_in_range(self, range)
1232    }
1233
1234    /// Merges one node (`source`) into another (`into`), along with a median entry.
1235    ///
1236    /// Example (values are not included for brevity):
1237    ///
1238    /// Input:
1239    ///   Source: [1, 2, 3]
1240    ///   Into: [5, 6, 7]
1241    ///   Median: 4
1242    ///
1243    /// Output:
1244    ///   [1, 2, 3, 4, 5, 6, 7] (stored in the `into` node)
1245    ///   `source` is deallocated.
1246    fn merge(&mut self, source: Node<K>, mut into: Node<K>, median: Entry<K>) -> Node<K> {
1247        into.merge(source, median, &mut self.allocator);
1248        into
1249    }
1250
1251    /// Allocates a new node of the given type.
1252    fn allocate_node(&mut self, node_type: NodeType) -> Node<K> {
1253        match self.version {
1254            Version::V1(page_size) => Node::new_v1(self.allocator.allocate(), node_type, page_size),
1255            Version::V2(page_size) => Node::new_v2(self.allocator.allocate(), node_type, page_size),
1256        }
1257    }
1258
1259    /// Deallocates a node.
1260    #[inline]
1261    fn deallocate_node(&mut self, node: Node<K>) {
1262        node.deallocate(self.allocator_mut());
1263    }
1264
1265    /// Loads a node from memory.
1266    #[inline]
1267    fn load_node(&self, address: Address) -> Node<K> {
1268        Node::load(address, self.version.page_size(), self.memory())
1269    }
1270
1271    /// Saves the node to memory.
1272    #[inline]
1273    fn save_node(&mut self, node: &mut Node<K>) {
1274        node.save(self.allocator_mut());
1275    }
1276
1277    /// Replaces the value at `idx` in the node, saves the node, and returns the old value.
1278    fn update_value(&mut self, node: &mut Node<K>, idx: usize, new_value: Vec<u8>) -> Vec<u8> {
1279        let old_value = node.swap_value(idx, new_value, self.memory());
1280        self.save_node(node);
1281        old_value
1282    }
1283
1284    /// Saves the map to memory.
1285    fn save_header(&self) {
1286        let header = BTreeHeader {
1287            version: self.version,
1288            root_addr: self.root_addr,
1289            length: self.length,
1290        };
1291
1292        Self::write_header(&header, self.memory());
1293    }
1294
1295    /// Write the layout header to the memory.
1296    fn write_header(header: &BTreeHeader, memory: &M) {
1297        // Serialize the header
1298        let mut buf = [0; PACKED_HEADER_SIZE];
1299        buf[0..3].copy_from_slice(MAGIC.as_slice());
1300        match header.version {
1301            Version::V1(DerivedPageSize {
1302                max_key_size,
1303                max_value_size,
1304            })
1305            | Version::V2(PageSize::Derived(DerivedPageSize {
1306                max_key_size,
1307                max_value_size,
1308            })) => {
1309                buf[3] = LAYOUT_VERSION;
1310                buf[4..8].copy_from_slice(&max_key_size.to_le_bytes());
1311                buf[8..12].copy_from_slice(&max_value_size.to_le_bytes());
1312            }
1313            Version::V2(PageSize::Value(page_size)) => {
1314                buf[3] = LAYOUT_VERSION_2;
1315                buf[4..8].copy_from_slice(&page_size.to_le_bytes());
1316                buf[8..12].copy_from_slice(&PAGE_SIZE_VALUE_MARKER.to_le_bytes());
1317            }
1318        };
1319        buf[12..20].copy_from_slice(&header.root_addr.get().to_le_bytes());
1320        buf[20..28].copy_from_slice(&header.length.to_le_bytes());
1321        // Write the header
1322        crate::write(memory, 0, &buf);
1323    }
1324}
1325
1326#[cfg(test)]
1327mod test {
1328    use super::*;
1329    use crate::{
1330        storable::{Blob, Bound as StorableBound},
1331        VectorMemory,
1332    };
1333    use std::cell::RefCell;
1334    use std::convert::TryFrom;
1335    use std::rc::Rc;
1336
1337    /// A helper function to clone and collect borrowed key/value pairs into a `Vec`.
1338    fn collect_kv<'a, K: Clone + 'a, V: Clone + 'a>(
1339        iter: impl Iterator<Item = (&'a K, &'a V)>,
1340    ) -> Vec<(K, V)> {
1341        iter.map(|(k, v)| (k.clone(), v.clone())).collect()
1342    }
1343
1344    /// A helper function to collect owned key/value pairs into a `Vec`.
1345    fn collect<K: Clone, V: Clone>(it: impl Iterator<Item = (K, V)>) -> Vec<(K, V)> {
1346        it.collect()
1347    }
1348
1349    /// Returns a fixed‑size buffer for the given u32.
1350    fn monotonic_buffer<const N: usize>(i: u32) -> [u8; N] {
1351        let mut buf = [0u8; N];
1352        let bytes = i.to_be_bytes();
1353        buf[N - bytes.len()..].copy_from_slice(&bytes);
1354        buf
1355    }
1356
1357    #[test]
1358    fn test_monotonic_buffer() {
1359        let cases: &[(u32, [u8; 4])] = &[
1360            (1, [0, 0, 0, 1]),
1361            (2, [0, 0, 0, 2]),
1362            ((1 << 8) - 1, [0, 0, 0, 255]),
1363            ((1 << 8), [0, 0, 1, 0]),
1364            ((1 << 16) - 1, [0, 0, 255, 255]),
1365            (1 << 16, [0, 1, 0, 0]),
1366            ((1 << 24) - 1, [0, 255, 255, 255]),
1367            (1 << 24, [1, 0, 0, 0]),
1368        ];
1369
1370        for &(input, expected) in cases {
1371            let output = monotonic_buffer::<4>(input);
1372            assert_eq!(
1373                output, expected,
1374                "monotonic_buffer::<4>({}) returned {:?}, expected {:?}",
1375                input, output, expected
1376            );
1377        }
1378    }
1379
1380    /// A trait to construct a value from a u32.
1381    trait Builder {
1382        fn build(i: u32) -> Self;
1383        fn empty() -> Self;
1384    }
1385
1386    impl Builder for () {
1387        fn build(_i: u32) -> Self {}
1388        fn empty() -> Self {}
1389    }
1390
1391    impl Builder for u32 {
1392        fn build(i: u32) -> Self {
1393            i
1394        }
1395        fn empty() -> Self {
1396            0
1397        }
1398    }
1399
1400    impl<const N: usize> Builder for Blob<N> {
1401        fn build(i: u32) -> Self {
1402            Blob::try_from(&monotonic_buffer::<N>(i)[..]).unwrap()
1403        }
1404        fn empty() -> Self {
1405            Blob::try_from(&[][..]).unwrap()
1406        }
1407    }
1408
1409    type MonotonicVec32 = Vec<u8>;
1410    impl Builder for MonotonicVec32 {
1411        fn build(i: u32) -> Self {
1412            monotonic_buffer::<32>(i).to_vec()
1413        }
1414        fn empty() -> Self {
1415            Vec::new()
1416        }
1417    }
1418
1419    type MonotonicString32 = String;
1420    impl Builder for MonotonicString32 {
1421        fn build(i: u32) -> Self {
1422            format!("{i:0>32}")
1423        }
1424        fn empty() -> Self {
1425            String::new()
1426        }
1427    }
1428
1429    /// Encodes an object into a byte vector.
1430    fn encode<T: Storable>(object: T) -> Vec<u8> {
1431        object.to_bytes_checked().into_owned()
1432    }
1433
1434    /// A helper method to succinctly create a blob.
1435    pub(crate) fn b(x: &[u8]) -> Blob<10> {
1436        Blob::<10>::try_from(x).unwrap()
1437    }
1438
1439    /// Creates a new shared memory instance.
1440    pub(crate) fn make_memory() -> Rc<RefCell<Vec<u8>>> {
1441        Rc::new(RefCell::new(Vec::new()))
1442    }
1443
1444    /// A test runner that runs the test using V1, migrated V2, and direct V2.
1445    pub fn run_btree_test<K, V, R, F>(f: F)
1446    where
1447        K: Storable + Ord + Clone,
1448        V: Storable,
1449        F: Fn(BTreeMap<K, V, VectorMemory>) -> R,
1450    {
1451        // V1 does not support unbounded types.
1452        if K::BOUND != StorableBound::Unbounded && V::BOUND != StorableBound::Unbounded {
1453            // Test with V1.
1454            let mem = make_memory();
1455            let tree_v1 = BTreeMap::new_v1(mem);
1456            f(tree_v1);
1457
1458            // Test with V2 migrated from V1.
1459            let mem = make_memory();
1460            let tree_v1: BTreeMap<K, V, _> = BTreeMap::new_v1(mem);
1461            let tree_v2_migrated = BTreeMap::load_helper(tree_v1.into_memory(), true);
1462            f(tree_v2_migrated);
1463        }
1464
1465        // Test with direct V2.
1466        let mem = make_memory();
1467        let tree_v2 = BTreeMap::new(mem);
1468        f(tree_v2);
1469    }
1470
1471    /// Checks that objects from boundary u32 values are strictly increasing.
1472    /// This ensures multi-byte conversions preserve order.
1473    fn verify_monotonic<T: Builder + PartialOrd>() {
1474        for shift_bits in [8, 16, 24] {
1475            let i = (1 << shift_bits) - 1;
1476            assert!(
1477                T::build(i) < T::build(i + 1),
1478                "Monotonicity failed at i: {i}",
1479            );
1480        }
1481    }
1482
1483    /// Asserts that the associated `BOUND` for `$ty` is _not_ `Unbounded`.
1484    macro_rules! assert_bounded {
1485        ($ty:ty) => {
1486            assert_ne!(<$ty>::BOUND, StorableBound::Unbounded, "Must be Bounded");
1487        };
1488    }
1489
1490    /// Asserts that the associated `BOUND` for `$ty` _is_ `Unbounded`.
1491    macro_rules! assert_unbounded {
1492        ($ty:ty) => {
1493            assert_eq!(<$ty>::BOUND, StorableBound::Unbounded, "Must be Unbounded");
1494        };
1495    }
1496
1497    macro_rules! run_with_key {
1498        ($runner:ident, $K:ty) => {{
1499            verify_monotonic::<$K>();
1500
1501            // Empty value.
1502            $runner::<$K, ()>();
1503
1504            // Bounded values.
1505            assert_bounded!(u32);
1506            $runner::<$K, u32>();
1507
1508            assert_bounded!(Blob<20>);
1509            $runner::<$K, Blob<20>>();
1510
1511            // Unbounded values.
1512            assert_unbounded!(MonotonicVec32);
1513            $runner::<$K, MonotonicVec32>();
1514
1515            assert_unbounded!(MonotonicString32);
1516            $runner::<$K, MonotonicString32>();
1517        }};
1518    }
1519
1520    /// Macro to apply a test function to a predefined grid of key/value types.
1521    macro_rules! btree_test {
1522        ($name:ident, $runner:ident) => {
1523            #[test]
1524            fn $name() {
1525                // Bounded keys.
1526                assert_bounded!(u32);
1527                run_with_key!($runner, u32);
1528
1529                assert_bounded!(Blob<10>);
1530                run_with_key!($runner, Blob<10>);
1531
1532                // Unbounded keys.
1533                assert_unbounded!(MonotonicVec32);
1534                run_with_key!($runner, MonotonicVec32);
1535
1536                assert_unbounded!(MonotonicString32);
1537                run_with_key!($runner, MonotonicString32);
1538            }
1539        };
1540    }
1541
1542    // Define a trait for keys that need the full set of bounds.
1543    trait TestKey: Storable + Ord + Clone + Builder + std::fmt::Debug {}
1544    impl<T> TestKey for T where T: Storable + Ord + Clone + Builder + std::fmt::Debug {}
1545
1546    // Define a trait for values that need the full set of bounds.
1547    trait TestValue: Storable + Clone + Builder + std::fmt::Debug + PartialEq {}
1548    impl<T> TestValue for T where T: Storable + Clone + Builder + std::fmt::Debug + PartialEq {}
1549
1550    fn insert_get_init_preserves_data<K: TestKey, V: TestValue>() {
1551        let (key, value) = (K::build, V::build);
1552        run_btree_test(|mut btree| {
1553            let n = 1_000;
1554            for i in 0..n {
1555                assert_eq!(btree.insert(key(i), value(i)), None);
1556                assert_eq!(btree.get(&key(i)), Some(value(i)));
1557            }
1558
1559            // Reload the BTreeMap and verify the entry.
1560            let btree = BTreeMap::<K, V, VectorMemory>::init(btree.into_memory());
1561            for i in 0..n {
1562                assert_eq!(btree.get(&key(i)), Some(value(i)));
1563            }
1564        });
1565    }
1566    btree_test!(
1567        test_insert_get_init_preserves_data,
1568        insert_get_init_preserves_data
1569    );
1570
1571    fn insert_overwrites_previous_value<K: TestKey, V: TestValue>() {
1572        let (key, value) = (K::build, V::build);
1573        run_btree_test(|mut btree| {
1574            let n = 1_000;
1575            for i in 0..n {
1576                assert_eq!(btree.insert(key(i), value(i)), None);
1577                assert_eq!(btree.insert(key(i), value(i + 1)), Some(value(i)));
1578                assert_eq!(btree.get(&key(i)), Some(value(i + 1)));
1579            }
1580        });
1581    }
1582    btree_test!(
1583        test_insert_overwrites_previous_value,
1584        insert_overwrites_previous_value
1585    );
1586
1587    fn insert_same_key_many<K: TestKey, V: TestValue>() {
1588        let (key, value) = (K::build, V::build);
1589        run_btree_test(|mut btree| {
1590            let n = 1_000;
1591            assert_eq!(btree.insert(key(1), value(2)), None);
1592            for i in 2..n {
1593                assert_eq!(btree.insert(key(1), value(i + 1)), Some(value(i)));
1594            }
1595            assert_eq!(btree.get(&key(1)), Some(value(n)));
1596        });
1597    }
1598    btree_test!(test_insert_same_key_many, insert_same_key_many);
1599
1600    fn insert_overwrite_median_key_in_full_child_node<K: TestKey, V: TestValue>() {
1601        let (key, value) = (K::build, V::build);
1602        run_btree_test(|mut btree| {
1603            for i in 1..=17 {
1604                assert_eq!(btree.insert(key(i), value(0)), None);
1605            }
1606
1607            // The result should look like this:
1608            //                [6]
1609            //               /   \
1610            // [1, 2, 3, 4, 5]   [7, 8, 9, 10, 11, (12), 13, 14, 15, 16, 17]
1611
1612            let root = btree.load_node(btree.root_addr);
1613            assert_eq!(root.node_type(), NodeType::Internal);
1614            assert_eq!(
1615                root.entries(btree.memory()),
1616                vec![(key(6), encode(value(0)))]
1617            );
1618            assert_eq!(root.children_len(), 2);
1619
1620            // The right child should now be full, with the median key being "12"
1621            let right_child = btree.load_node(root.child(1));
1622            assert!(right_child.is_full());
1623            let median_index = right_child.entries_len() / 2;
1624            let median_key = key(12);
1625            assert_eq!(right_child.key(median_index, btree.memory()), &median_key);
1626
1627            // Overwrite the value of the median key.
1628            assert_eq!(btree.insert(median_key.clone(), value(123)), Some(value(0)));
1629            assert_eq!(btree.get(&median_key), Some(value(123)));
1630
1631            // The child has not been split and is still full.
1632            let right_child = btree.load_node(root.child(1));
1633            assert_eq!(right_child.node_type(), NodeType::Leaf);
1634            assert!(right_child.is_full());
1635        });
1636    }
1637    btree_test!(
1638        test_insert_overwrite_median_key_in_full_child_node,
1639        insert_overwrite_median_key_in_full_child_node
1640    );
1641
1642    fn insert_overwrite_key_in_full_root_node<K: TestKey, V: TestValue>() {
1643        let (key, value) = (K::build, V::build);
1644        run_btree_test(|mut btree| {
1645            for i in 1..=11 {
1646                assert_eq!(btree.insert(key(i), value(0)), None);
1647            }
1648
1649            // We now have a root that is full and looks like this:
1650            //
1651            // [1, 2, 3, 4, 5, (6), 7, 8, 9, 10, 11]
1652            let root = btree.load_node(btree.root_addr);
1653            assert!(root.is_full());
1654
1655            // Overwrite an element in the root. It should NOT cause the node to be split.
1656            assert_eq!(btree.insert(key(6), value(456)), Some(value(0)));
1657
1658            let root = btree.load_node(btree.root_addr);
1659            assert_eq!(root.node_type(), NodeType::Leaf);
1660            assert_eq!(btree.get(&key(6)), Some(value(456)));
1661            assert_eq!(root.entries_len(), 11);
1662        });
1663    }
1664    btree_test!(
1665        test_insert_overwrite_key_in_full_root_node,
1666        insert_overwrite_key_in_full_root_node
1667    );
1668
1669    fn allocations_without_split<K: TestKey, V: TestValue>() {
1670        let key = K::build;
1671        run_btree_test(|mut btree| {
1672            assert_eq!(btree.allocator.num_allocated_chunks(), 0);
1673
1674            assert_eq!(btree.insert(key(1), V::empty()), None);
1675            assert_eq!(btree.allocator.num_allocated_chunks(), 1);
1676
1677            assert_eq!(btree.remove(&key(1)), Some(V::empty()));
1678            assert_eq!(btree.allocator.num_allocated_chunks(), 0);
1679        });
1680    }
1681    btree_test!(test_allocations_without_split, allocations_without_split);
1682
1683    fn allocations_with_split<K: TestKey, V: TestValue>() {
1684        let key = K::build;
1685        run_btree_test(|mut btree| {
1686            // Insert entries until the root node is full.
1687            let mut i = 0;
1688            loop {
1689                assert_eq!(btree.insert(key(i), V::empty()), None);
1690                let root = btree.load_node(btree.root_addr);
1691                if root.is_full() {
1692                    break;
1693                }
1694                i += 1;
1695            }
1696
1697            // Only need a single allocation to store up to `CAPACITY` elements.
1698            assert_eq!(btree.allocator.num_allocated_chunks(), 1);
1699
1700            assert_eq!(btree.insert(key(255), V::empty()), None);
1701
1702            // The node had to be split into three nodes.
1703            assert_eq!(btree.allocator.num_allocated_chunks(), 3);
1704        });
1705    }
1706    btree_test!(test_allocations_with_split, allocations_with_split);
1707
1708    fn insert_split_node<K: TestKey, V: TestValue>() {
1709        let (key, value) = (K::build, V::build);
1710        run_btree_test(|mut btree| {
1711            for i in 1..=11 {
1712                assert_eq!(btree.insert(key(i), value(i)), None);
1713            }
1714
1715            // Should now split a node.
1716            let root = btree.load_node(btree.root_addr);
1717            assert!(root.is_full());
1718            assert_eq!(btree.insert(key(12), value(12)), None);
1719
1720            // The result should look like this:
1721            //                [6]
1722            //               /   \
1723            // [1, 2, 3, 4, 5]   [7, 8, 9, 10, 11, 12]
1724
1725            for i in 1..=12 {
1726                assert_eq!(btree.get(&key(i)), Some(value(i)));
1727            }
1728        });
1729    }
1730    btree_test!(test_insert_split_node, insert_split_node);
1731
1732    fn insert_split_multiple_nodes<K: TestKey, V: TestValue>() {
1733        let key = K::build;
1734        let e = |i: u32| (key(i), encode(V::empty()));
1735        run_btree_test(|mut btree| {
1736            for i in 1..=11 {
1737                assert_eq!(btree.insert(key(i), V::empty()), None);
1738            }
1739            // Should now split a node.
1740            assert_eq!(btree.insert(key(12), V::empty()), None);
1741
1742            // The result should look like this:
1743            //                [6]
1744            //               /   \
1745            // [1, 2, 3, 4, 5]   [7, 8, 9, 10, 11, 12]
1746
1747            let root = btree.load_node(btree.root_addr);
1748            assert_eq!(root.node_type(), NodeType::Internal);
1749            assert_eq!(root.entries(btree.memory()), vec![e(6)]);
1750            assert_eq!(root.children_len(), 2);
1751
1752            let child_0 = btree.load_node(root.child(0));
1753            assert_eq!(child_0.node_type(), NodeType::Leaf);
1754            assert_eq!(
1755                child_0.entries(btree.memory()),
1756                vec![e(1), e(2), e(3), e(4), e(5)]
1757            );
1758
1759            let child_1 = btree.load_node(root.child(1));
1760            assert_eq!(child_1.node_type(), NodeType::Leaf);
1761            assert_eq!(
1762                child_1.entries(btree.memory()),
1763                vec![e(7), e(8), e(9), e(10), e(11), e(12)]
1764            );
1765
1766            for i in 1..=12 {
1767                assert_eq!(btree.get(&key(i)), Some(V::empty()));
1768            }
1769
1770            // Insert more to cause more splitting.
1771            assert_eq!(btree.insert(key(13), V::empty()), None);
1772            assert_eq!(btree.insert(key(14), V::empty()), None);
1773            assert_eq!(btree.insert(key(15), V::empty()), None);
1774            assert_eq!(btree.insert(key(16), V::empty()), None);
1775            assert_eq!(btree.insert(key(17), V::empty()), None);
1776            // Should cause another split
1777            assert_eq!(btree.insert(key(18), V::empty()), None);
1778
1779            for i in 1..=18 {
1780                assert_eq!(btree.get(&key(i)), Some(V::empty()));
1781            }
1782
1783            let root = btree.load_node(btree.root_addr);
1784            assert_eq!(root.node_type(), NodeType::Internal);
1785            assert_eq!(root.entries(btree.memory()), vec![e(6), e(12)],);
1786            assert_eq!(root.children_len(), 3);
1787
1788            let child_0 = btree.load_node(root.child(0));
1789            assert_eq!(child_0.node_type(), NodeType::Leaf);
1790            assert_eq!(
1791                child_0.entries(btree.memory()),
1792                vec![e(1), e(2), e(3), e(4), e(5)]
1793            );
1794
1795            let child_1 = btree.load_node(root.child(1));
1796            assert_eq!(child_1.node_type(), NodeType::Leaf);
1797            assert_eq!(
1798                child_1.entries(btree.memory()),
1799                vec![e(7), e(8), e(9), e(10), e(11)]
1800            );
1801
1802            let child_2 = btree.load_node(root.child(2));
1803            assert_eq!(child_2.node_type(), NodeType::Leaf);
1804            assert_eq!(
1805                child_2.entries(btree.memory()),
1806                vec![e(13), e(14), e(15), e(16), e(17), e(18)]
1807            );
1808
1809            assert_eq!(btree.allocator.num_allocated_chunks(), 4);
1810        });
1811    }
1812    btree_test!(
1813        test_insert_split_multiple_nodes,
1814        insert_split_multiple_nodes
1815    );
1816
1817    fn first_key_value<K: TestKey, V: TestValue>() {
1818        let (key, value) = (K::build, V::build);
1819        run_btree_test(|mut btree: BTreeMap<K, V, _>| {
1820            assert_eq!(btree.first_key_value(), None);
1821
1822            let n = 1_000;
1823            for i in (0..n).rev() {
1824                assert_eq!(btree.insert(key(i), value(i)), None);
1825                assert_eq!(btree.first_key_value(), Some((key(i), value(i))));
1826            }
1827
1828            for i in 0..n {
1829                assert_eq!(btree.remove(&key(i)), Some(value(i)));
1830                if !btree.is_empty() {
1831                    assert_eq!(btree.first_key_value(), Some((key(i + 1), value(i + 1))));
1832                }
1833            }
1834            assert_eq!(btree.first_key_value(), None);
1835        });
1836    }
1837    btree_test!(test_first_key_value, first_key_value);
1838
1839    fn last_key_value<K: TestKey, V: TestValue>() {
1840        let (key, value) = (K::build, V::build);
1841        run_btree_test(|mut btree: BTreeMap<K, V, _>| {
1842            assert_eq!(btree.last_key_value(), None);
1843
1844            let n = 1_000;
1845            for i in 0..n {
1846                assert_eq!(btree.insert(key(i), value(i)), None);
1847                assert_eq!(btree.last_key_value(), Some((key(i), value(i))));
1848            }
1849
1850            for i in (0..n).rev() {
1851                assert_eq!(btree.remove(&key(i)), Some(value(i)));
1852                if !btree.is_empty() {
1853                    assert_eq!(btree.last_key_value(), Some((key(i - 1), value(i - 1))));
1854                }
1855            }
1856            assert_eq!(btree.last_key_value(), None);
1857        });
1858    }
1859    btree_test!(test_last_key_value, last_key_value);
1860
1861    fn pop_first_single_entry<K: TestKey, V: TestValue>() {
1862        let key = K::build;
1863        run_btree_test(|mut btree| {
1864            assert_eq!(btree.allocator.num_allocated_chunks(), 0);
1865
1866            assert_eq!(btree.insert(key(1), V::empty()), None);
1867            assert!(!btree.is_empty());
1868            assert_eq!(btree.allocator.num_allocated_chunks(), 1);
1869
1870            assert_eq!(btree.pop_first(), Some((key(1), V::empty())));
1871            assert!(btree.is_empty());
1872            assert_eq!(btree.allocator.num_allocated_chunks(), 0);
1873        });
1874    }
1875    btree_test!(test_pop_first_single_entry, pop_first_single_entry);
1876
1877    fn pop_last_single_entry<K: TestKey, V: TestValue>() {
1878        let (key, value) = (K::build, V::build);
1879        run_btree_test(|mut btree| {
1880            assert_eq!(btree.allocator.num_allocated_chunks(), 0);
1881
1882            assert_eq!(btree.insert(key(1), value(1)), None);
1883            assert!(!btree.is_empty());
1884            assert_eq!(btree.allocator.num_allocated_chunks(), 1);
1885
1886            assert_eq!(btree.pop_last(), Some((key(1), value(1))));
1887            assert!(btree.is_empty());
1888            assert_eq!(btree.allocator.num_allocated_chunks(), 0);
1889        });
1890    }
1891    btree_test!(test_pop_last_single_entry, pop_last_single_entry);
1892
1893    fn remove_case_2a_and_2c<K: TestKey, V: TestValue>() {
1894        let key = K::build;
1895        let e = |i: u32| (key(i), encode(V::empty()));
1896        run_btree_test(|mut btree| {
1897            for i in 1..=11 {
1898                assert_eq!(btree.insert(key(i), V::empty()), None);
1899            }
1900            // Should now split a node.
1901            assert_eq!(btree.insert(key(0), V::empty()), None);
1902
1903            // The result should look like this:
1904            //                    [6]
1905            //                   /   \
1906            // [0, 1, 2, 3, 4, 5]     [7, 8, 9, 10, 11]
1907
1908            for i in 0..=11 {
1909                assert_eq!(btree.get(&key(i)), Some(V::empty()));
1910            }
1911
1912            // Remove node 6. Triggers case 2.a
1913            assert_eq!(btree.remove(&key(6)), Some(V::empty()));
1914
1915            // The result should look like this:
1916            //                [5]
1917            //               /   \
1918            // [0, 1, 2, 3, 4]   [7, 8, 9, 10, 11]
1919            let root = btree.load_node(btree.root_addr);
1920            assert_eq!(root.node_type(), NodeType::Internal);
1921            assert_eq!(root.entries(btree.memory()), vec![e(5)]);
1922            assert_eq!(root.children_len(), 2);
1923
1924            let child_0 = btree.load_node(root.child(0));
1925            assert_eq!(child_0.node_type(), NodeType::Leaf);
1926            assert_eq!(
1927                child_0.entries(btree.memory()),
1928                vec![e(0), e(1), e(2), e(3), e(4)]
1929            );
1930
1931            let child_1 = btree.load_node(root.child(1));
1932            assert_eq!(child_1.node_type(), NodeType::Leaf);
1933            assert_eq!(
1934                child_1.entries(btree.memory()),
1935                vec![e(7), e(8), e(9), e(10), e(11)]
1936            );
1937
1938            // There are three allocated nodes.
1939            assert_eq!(btree.allocator.num_allocated_chunks(), 3);
1940
1941            // Remove node 5. Triggers case 2c
1942            assert_eq!(btree.remove(&key(5)), Some(V::empty()));
1943
1944            // Reload the btree to verify that we saved it correctly.
1945            let btree = BTreeMap::<K, V, _>::load(btree.into_memory());
1946
1947            // The result should look like this:
1948            // [0, 1, 2, 3, 4, 7, 8, 9, 10, 11]
1949            let root = btree.load_node(btree.root_addr);
1950            assert_eq!(
1951                root.entries(btree.memory()),
1952                vec![e(0), e(1), e(2), e(3), e(4), e(7), e(8), e(9), e(10), e(11)]
1953            );
1954
1955            // There is only one node allocated.
1956            assert_eq!(btree.allocator.num_allocated_chunks(), 1);
1957        });
1958    }
1959    btree_test!(test_remove_case_2a_and_2c, remove_case_2a_and_2c);
1960
1961    fn remove_case_2b<K: TestKey, V: TestValue>() {
1962        let key = K::build;
1963        let e = |i: u32| (key(i), encode(V::empty()));
1964        run_btree_test(|mut btree| {
1965            for i in 1..=11 {
1966                assert_eq!(btree.insert(key(i), V::empty()), None);
1967            }
1968            // Should now split a node.
1969            assert_eq!(btree.insert(key(12), V::empty()), None);
1970
1971            // The result should look like this:
1972            //                [6]
1973            //               /   \
1974            // [1, 2, 3, 4, 5]   [7, 8, 9, 10, 11, 12]
1975
1976            for i in 1..=12 {
1977                assert_eq!(btree.get(&key(i)), Some(V::empty()));
1978            }
1979
1980            // Remove node 6. Triggers case 2.b
1981            assert_eq!(btree.remove(&key(6)), Some(V::empty()));
1982
1983            // The result should look like this:
1984            //                [7]
1985            //               /   \
1986            // [1, 2, 3, 4, 5]   [8, 9, 10, 11, 12]
1987            let root = btree.load_node(btree.root_addr);
1988            assert_eq!(root.node_type(), NodeType::Internal);
1989            assert_eq!(root.entries(btree.memory()), vec![e(7)]);
1990            assert_eq!(root.children_len(), 2);
1991
1992            let child_0 = btree.load_node(root.child(0));
1993            assert_eq!(child_0.node_type(), NodeType::Leaf);
1994            assert_eq!(
1995                child_0.entries(btree.memory()),
1996                vec![e(1), e(2), e(3), e(4), e(5)]
1997            );
1998
1999            let child_1 = btree.load_node(root.child(1));
2000            assert_eq!(child_1.node_type(), NodeType::Leaf);
2001            assert_eq!(
2002                child_1.entries(btree.memory()),
2003                vec![e(8), e(9), e(10), e(11), e(12)]
2004            );
2005
2006            // Remove node 7. Triggers case 2.c
2007            assert_eq!(btree.remove(&key(7)), Some(V::empty()));
2008            // The result should look like this:
2009            //
2010            // [1, 2, 3, 4, 5, 8, 9, 10, 11, 12]
2011            let root = btree.load_node(btree.root_addr);
2012            assert_eq!(root.node_type(), NodeType::Leaf);
2013            assert_eq!(
2014                root.entries(btree.memory()),
2015                collect([1, 2, 3, 4, 5, 8, 9, 10, 11, 12].into_iter().map(e))
2016            );
2017
2018            assert_eq!(btree.allocator.num_allocated_chunks(), 1);
2019        });
2020    }
2021    btree_test!(test_remove_case_2b, remove_case_2b);
2022
2023    fn remove_case_3a_right<K: TestKey, V: TestValue>() {
2024        let key = K::build;
2025        let e = |i: u32| (key(i), encode(V::empty()));
2026        run_btree_test(|mut btree| {
2027            for i in 1..=11 {
2028                assert_eq!(btree.insert(key(i), V::empty()), None);
2029            }
2030
2031            // Should now split a node.
2032            assert_eq!(btree.insert(key(12), V::empty()), None);
2033
2034            // The result should look like this:
2035            //                [6]
2036            //               /   \
2037            // [1, 2, 3, 4, 5]   [7, 8, 9, 10, 11, 12]
2038
2039            // Remove node 3. Triggers case 3.a
2040            assert_eq!(btree.remove(&key(3)), Some(V::empty()));
2041
2042            // The result should look like this:
2043            //                [7]
2044            //               /   \
2045            // [1, 2, 4, 5, 6]   [8, 9, 10, 11, 12]
2046            let root = btree.load_node(btree.root_addr);
2047            assert_eq!(root.node_type(), NodeType::Internal);
2048            assert_eq!(root.entries(btree.memory()), vec![e(7)]);
2049            assert_eq!(root.children_len(), 2);
2050
2051            let child_0 = btree.load_node(root.child(0));
2052            assert_eq!(child_0.node_type(), NodeType::Leaf);
2053            assert_eq!(
2054                child_0.entries(btree.memory()),
2055                vec![e(1), e(2), e(4), e(5), e(6)]
2056            );
2057
2058            let child_1 = btree.load_node(root.child(1));
2059            assert_eq!(child_1.node_type(), NodeType::Leaf);
2060            assert_eq!(
2061                child_1.entries(btree.memory()),
2062                vec![e(8), e(9), e(10), e(11), e(12)]
2063            );
2064
2065            // There are three allocated nodes.
2066            assert_eq!(btree.allocator.num_allocated_chunks(), 3);
2067        });
2068    }
2069    btree_test!(test_remove_case_3a_right, remove_case_3a_right);
2070
2071    fn remove_case_3a_left<K: TestKey, V: TestValue>() {
2072        let key = K::build;
2073        let e = |i: u32| (key(i), encode(V::empty()));
2074        run_btree_test(|mut btree| {
2075            for i in 1..=11 {
2076                assert_eq!(btree.insert(key(i), V::empty()), None);
2077            }
2078            // Should now split a node.
2079            assert_eq!(btree.insert(key(0), V::empty()), None);
2080
2081            // The result should look like this:
2082            //                   [6]
2083            //                  /   \
2084            // [0, 1, 2, 3, 4, 5]   [7, 8, 9, 10, 11]
2085
2086            // Remove node 8. Triggers case 3.a left
2087            assert_eq!(btree.remove(&key(8)), Some(V::empty()));
2088
2089            // The result should look like this:
2090            //                [5]
2091            //               /   \
2092            // [0, 1, 2, 3, 4]   [6, 7, 9, 10, 11]
2093            let root = btree.load_node(btree.root_addr);
2094            assert_eq!(root.node_type(), NodeType::Internal);
2095            assert_eq!(root.entries(btree.memory()), vec![e(5)]);
2096            assert_eq!(root.children_len(), 2);
2097
2098            let child_0 = btree.load_node(root.child(0));
2099            assert_eq!(child_0.node_type(), NodeType::Leaf);
2100            assert_eq!(
2101                child_0.entries(btree.memory()),
2102                vec![e(0), e(1), e(2), e(3), e(4)]
2103            );
2104
2105            let child_1 = btree.load_node(root.child(1));
2106            assert_eq!(child_1.node_type(), NodeType::Leaf);
2107            assert_eq!(
2108                child_1.entries(btree.memory()),
2109                vec![e(6), e(7), e(9), e(10), e(11)]
2110            );
2111
2112            // There are three allocated nodes.
2113            assert_eq!(btree.allocator.num_allocated_chunks(), 3);
2114        });
2115    }
2116    btree_test!(test_remove_case_3a_left, remove_case_3a_left);
2117
2118    fn remove_case_3b_merge_into_right<K: TestKey, V: TestValue>() {
2119        let key = K::build;
2120        let e = |i: u32| (key(i), encode(V::empty()));
2121        run_btree_test(|mut btree| {
2122            for i in 1..=11 {
2123                assert_eq!(btree.insert(key(i), V::empty()), None);
2124            }
2125            // Should now split a node.
2126            assert_eq!(btree.insert(key(12), V::empty()), None);
2127
2128            // The result should look like this:
2129            //                [6]
2130            //               /   \
2131            // [1, 2, 3, 4, 5]   [7, 8, 9, 10, 11, 12]
2132
2133            for i in 1..=12 {
2134                assert_eq!(btree.get(&key(i)), Some(V::empty()));
2135            }
2136
2137            // Remove node 6. Triggers case 2.b
2138            assert_eq!(btree.remove(&key(6)), Some(V::empty()));
2139            // The result should look like this:
2140            //                [7]
2141            //               /   \
2142            // [1, 2, 3, 4, 5]   [8, 9, 10, 11, 12]
2143            let root = btree.load_node(btree.root_addr);
2144            assert_eq!(root.node_type(), NodeType::Internal);
2145            assert_eq!(root.entries(btree.memory()), vec![e(7)]);
2146            assert_eq!(root.children_len(), 2);
2147
2148            let child_0 = btree.load_node(root.child(0));
2149            assert_eq!(child_0.node_type(), NodeType::Leaf);
2150            assert_eq!(
2151                child_0.entries(btree.memory()),
2152                vec![e(1), e(2), e(3), e(4), e(5)]
2153            );
2154
2155            let child_1 = btree.load_node(root.child(1));
2156            assert_eq!(child_1.node_type(), NodeType::Leaf);
2157            assert_eq!(
2158                child_1.entries(btree.memory()),
2159                vec![e(8), e(9), e(10), e(11), e(12)]
2160            );
2161
2162            // There are three allocated nodes.
2163            assert_eq!(btree.allocator.num_allocated_chunks(), 3);
2164
2165            // Remove node 3. Triggers case 3.b
2166            assert_eq!(btree.remove(&key(3)), Some(V::empty()));
2167
2168            // Reload the btree to verify that we saved it correctly.
2169            let btree = BTreeMap::<K, V, _>::load(btree.into_memory());
2170
2171            // The result should look like this:
2172            //
2173            // [1, 2, 4, 5, 7, 8, 9, 10, 11, 12]
2174            let root = btree.load_node(btree.root_addr);
2175            assert_eq!(root.node_type(), NodeType::Leaf);
2176            assert_eq!(
2177                root.entries(btree.memory()),
2178                collect([1, 2, 4, 5, 7, 8, 9, 10, 11, 12].into_iter().map(e))
2179            );
2180
2181            // There is only one allocated node remaining.
2182            assert_eq!(btree.allocator.num_allocated_chunks(), 1);
2183        });
2184    }
2185    btree_test!(
2186        test_remove_case_3b_merge_into_right,
2187        remove_case_3b_merge_into_right
2188    );
2189
2190    fn remove_case_3b_merge_into_left<K: TestKey, V: TestValue>() {
2191        let key = K::build;
2192        let e = |i: u32| (key(i), encode(V::empty()));
2193        run_btree_test(|mut btree| {
2194            for i in 1..=11 {
2195                assert_eq!(btree.insert(key(i), V::empty()), None);
2196            }
2197
2198            // Should now split a node.
2199            assert_eq!(btree.insert(key(12), V::empty()), None);
2200
2201            // The result should look like this:
2202            //                [6]
2203            //               /   \
2204            // [1, 2, 3, 4, 5]   [7, 8, 9, 10, 11, 12]
2205
2206            for i in 1..=12 {
2207                assert_eq!(btree.get(&key(i)), Some(V::empty()));
2208            }
2209
2210            // Remove node 6. Triggers case 2.b
2211            assert_eq!(btree.remove(&key(6)), Some(V::empty()));
2212
2213            // The result should look like this:
2214            //                [7]
2215            //               /   \
2216            // [1, 2, 3, 4, 5]   [8, 9, 10, 11, 12]
2217            let root = btree.load_node(btree.root_addr);
2218            assert_eq!(root.node_type(), NodeType::Internal);
2219            assert_eq!(root.entries(btree.memory()), vec![e(7)]);
2220            assert_eq!(root.children_len(), 2);
2221
2222            let child_0 = btree.load_node(root.child(0));
2223            assert_eq!(child_0.node_type(), NodeType::Leaf);
2224            assert_eq!(
2225                child_0.entries(btree.memory()),
2226                vec![e(1), e(2), e(3), e(4), e(5)]
2227            );
2228
2229            let child_1 = btree.load_node(root.child(1));
2230            assert_eq!(child_1.node_type(), NodeType::Leaf);
2231            assert_eq!(
2232                child_1.entries(btree.memory()),
2233                vec![e(8), e(9), e(10), e(11), e(12)]
2234            );
2235
2236            // There are three allocated nodes.
2237            assert_eq!(btree.allocator.num_allocated_chunks(), 3);
2238
2239            // Remove node 10. Triggers case 3.b where we merge the right into the left.
2240            assert_eq!(btree.remove(&key(10)), Some(V::empty()));
2241
2242            // Reload the btree to verify that we saved it correctly.
2243            let btree = BTreeMap::<K, V, _>::load(btree.into_memory());
2244
2245            // The result should look like this:
2246            //
2247            // [1, 2, 3, 4, 5, 7, 8, 9, 11, 12], 10 is gone
2248            let root = btree.load_node(btree.root_addr);
2249            assert_eq!(root.node_type(), NodeType::Leaf);
2250            assert_eq!(
2251                root.entries(btree.memory()),
2252                vec![e(1), e(2), e(3), e(4), e(5), e(7), e(8), e(9), e(11), e(12)]
2253            );
2254
2255            // There is only one allocated node remaining.
2256            assert_eq!(btree.allocator.num_allocated_chunks(), 1);
2257        });
2258    }
2259    btree_test!(
2260        test_remove_case_3b_merge_into_left,
2261        remove_case_3b_merge_into_left
2262    );
2263
2264    fn insert_remove_many<K: TestKey, V: TestValue>() {
2265        let (key, value) = (K::build, V::build);
2266        run_btree_test(|mut btree| {
2267            let n = 10_000;
2268            for i in 0..n {
2269                assert_eq!(btree.insert(key(i), value(i)), None);
2270                assert_eq!(btree.get(&key(i)), Some(value(i)));
2271            }
2272
2273            let mut btree = BTreeMap::<K, V, _>::load(btree.into_memory());
2274            for i in 0..n {
2275                assert_eq!(btree.remove(&key(i)), Some(value(i)));
2276                assert_eq!(btree.get(&key(i)), None);
2277            }
2278
2279            // We've deallocated everything.
2280            assert_eq!(btree.allocator.num_allocated_chunks(), 0);
2281        });
2282    }
2283    btree_test!(test_insert_remove_many, insert_remove_many);
2284
2285    fn pop_first_many<K: TestKey, V: TestValue>() {
2286        let (key, value) = (K::build, V::build);
2287        run_btree_test(|mut btree| {
2288            let n = 10_000;
2289
2290            assert!(btree.is_empty());
2291            assert_eq!(btree.pop_first(), None);
2292
2293            for i in 0..n {
2294                assert_eq!(btree.insert(key(i), value(i)), None);
2295            }
2296            assert_eq!(btree.len(), n as u64);
2297
2298            let mut btree = BTreeMap::<K, V, _>::load(btree.into_memory());
2299            for i in 0..n {
2300                assert_eq!(btree.pop_first(), Some((key(i), value(i))));
2301                assert_eq!(btree.get(&key(i)), None);
2302            }
2303
2304            // We've deallocated everything.
2305            assert!(btree.is_empty());
2306            assert_eq!(btree.allocator.num_allocated_chunks(), 0);
2307        });
2308    }
2309    btree_test!(test_pop_first_many, pop_first_many);
2310
2311    fn pop_last_many<K: TestKey, V: TestValue>() {
2312        let (key, value) = (K::build, V::build);
2313        run_btree_test(|mut btree| {
2314            let n = 10_000;
2315
2316            assert!(btree.is_empty());
2317            assert_eq!(btree.pop_last(), None);
2318
2319            for i in 0..n {
2320                assert_eq!(btree.insert(key(i), value(i)), None);
2321            }
2322            assert_eq!(btree.len(), n as u64);
2323
2324            let mut btree = BTreeMap::<K, V, _>::load(btree.into_memory());
2325            for i in (0..n).rev() {
2326                assert_eq!(btree.pop_last(), Some((key(i), value(i))));
2327                assert_eq!(btree.get(&key(i)), None);
2328            }
2329
2330            // We've deallocated everything.
2331            assert!(btree.is_empty());
2332            assert_eq!(btree.allocator.num_allocated_chunks(), 0);
2333        });
2334    }
2335    btree_test!(test_pop_last_many, pop_last_many);
2336
2337    fn reloading<K: TestKey, V: TestValue>() {
2338        let (key, value) = (K::build, V::build);
2339        run_btree_test(|mut btree| {
2340            let n = 1_000;
2341            for i in 0..n {
2342                assert_eq!(btree.insert(key(i), value(i)), None);
2343                assert_eq!(btree.get(&key(i)), Some(value(i)));
2344            }
2345            assert_eq!(btree.len(), n as u64);
2346
2347            // Reload the BTreeMap and verify the entry.
2348            let mut btree = BTreeMap::<K, V, VectorMemory>::load(btree.into_memory());
2349            assert_eq!(btree.len(), n as u64);
2350            for i in 0..n {
2351                assert_eq!(btree.get(&key(i)), Some(value(i)));
2352            }
2353
2354            // Remove all entries and verify the entry.
2355            for i in 0..n {
2356                assert_eq!(btree.remove(&key(i)), Some(value(i)));
2357            }
2358            assert_eq!(btree.len(), 0);
2359
2360            // Reload the BTreeMap and verify the entry.
2361            let btree = BTreeMap::<K, V, VectorMemory>::load(btree.into_memory());
2362            assert_eq!(btree.len(), 0);
2363        });
2364    }
2365    btree_test!(test_reloading, reloading);
2366
2367    fn len<K: TestKey, V: TestValue>() {
2368        let (key, value) = (K::build, V::build);
2369        run_btree_test(|mut btree| {
2370            let n = 1_000;
2371            for i in 0..n {
2372                assert_eq!(btree.insert(key(i), value(i)), None);
2373            }
2374
2375            assert_eq!(btree.len(), n as u64);
2376            assert!(!btree.is_empty());
2377
2378            for i in 0..n {
2379                assert_eq!(btree.remove(&key(i)), Some(value(i)));
2380            }
2381
2382            assert_eq!(btree.len(), 0);
2383            assert!(btree.is_empty());
2384        });
2385    }
2386    btree_test!(test_len, len);
2387
2388    fn contains_key<K: TestKey, V: TestValue>() {
2389        let (key, value) = (K::build, V::build);
2390        run_btree_test(|mut btree| {
2391            let n = 1_000;
2392            for i in (0..n).step_by(2) {
2393                assert_eq!(btree.insert(key(i), value(i)), None);
2394            }
2395
2396            // Only even keys should be present.
2397            for i in 0..n {
2398                assert_eq!(btree.contains_key(&key(i)), i % 2 == 0);
2399            }
2400        });
2401    }
2402    btree_test!(test_contains_key, contains_key);
2403
2404    fn range_empty<K: TestKey, V: TestValue>() {
2405        let key = K::build;
2406        run_btree_test(|btree: BTreeMap<K, V, _>| {
2407            // Test prefixes that don't exist in the map.
2408            assert_eq!(collect(btree.range(key(0)..)), vec![]);
2409            assert_eq!(collect(btree.range(key(10)..)), vec![]);
2410            assert_eq!(collect(btree.range(key(100)..)), vec![]);
2411        });
2412    }
2413    btree_test!(test_range_empty, range_empty);
2414
2415    // Tests the case where the prefix is larger than all the entries in a leaf node.
2416    fn range_leaf_prefix_greater_than_all_entries<K: TestKey, V: TestValue>() {
2417        let (key, value) = (K::build, V::build);
2418        run_btree_test(|mut btree| {
2419            btree.insert(key(0), value(0));
2420
2421            // Test a prefix that's larger than the value in the leaf node. Should be empty.
2422            assert_eq!(collect(btree.range(key(1)..)), vec![]);
2423        });
2424    }
2425    btree_test!(
2426        test_range_leaf_prefix_greater_than_all_entries,
2427        range_leaf_prefix_greater_than_all_entries
2428    );
2429
2430    // Tests the case where the prefix is larger than all the entries in an internal node.
2431    fn range_internal_prefix_greater_than_all_entries<K: TestKey, V: TestValue>() {
2432        let (key, value) = (K::build, V::build);
2433        run_btree_test(|mut btree| {
2434            for i in 1..=12 {
2435                assert_eq!(btree.insert(key(i), value(i)), None);
2436            }
2437
2438            // The result should look like this:
2439            //                [6]
2440            //               /   \
2441            // [1, 2, 3, 4, 5]   [7, 8, 9, 10, 11, 12]
2442
2443            // Test a prefix that's larger than the key in the internal node.
2444            assert_eq!(
2445                collect(btree.range(key(7)..key(8))),
2446                vec![(key(7), value(7))]
2447            );
2448        });
2449    }
2450    btree_test!(
2451        test_range_internal_prefix_greater_than_all_entries,
2452        range_internal_prefix_greater_than_all_entries
2453    );
2454
2455    fn range_various_prefixes<K: TestKey, V: TestValue>() {
2456        let (key, value) = (K::build, V::build);
2457        run_btree_test(|mut btree| {
2458            btree.insert(key(1), value(100));
2459            btree.insert(key(2), value(200));
2460            btree.insert(key(3), value(300));
2461            btree.insert(key(4), value(400));
2462            btree.insert(key(11), value(500));
2463            btree.insert(key(12), value(600));
2464            btree.insert(key(13), value(700));
2465            btree.insert(key(14), value(800));
2466            btree.insert(key(21), value(900));
2467            btree.insert(key(22), value(1_000));
2468            btree.insert(key(23), value(1_100));
2469            btree.insert(key(24), value(1_200));
2470
2471            // The result should look like this:
2472            //                 [12]
2473            //                /    \
2474            // [1, 2, 3, 4, 11]    [13, 14, 21, 22, 23, 24]
2475
2476            let root = btree.load_node(btree.root_addr);
2477            assert_eq!(root.node_type(), NodeType::Internal);
2478            assert_eq!(
2479                root.entries(btree.memory()),
2480                vec![(key(12), encode(value(600)))]
2481            );
2482            assert_eq!(root.children_len(), 2);
2483
2484            // Tests a prefix that's smaller than the key in the internal node.
2485            assert_eq!(
2486                collect(btree.range(key(0)..key(11))),
2487                vec![
2488                    (key(1), value(100)),
2489                    (key(2), value(200)),
2490                    (key(3), value(300)),
2491                    (key(4), value(400)),
2492                ]
2493            );
2494
2495            // Tests a prefix that crosses several nodes.
2496            assert_eq!(
2497                collect(btree.range(key(10)..key(20))),
2498                vec![
2499                    (key(11), value(500)),
2500                    (key(12), value(600)),
2501                    (key(13), value(700)),
2502                    (key(14), value(800)),
2503                ]
2504            );
2505
2506            // Tests a prefix that's larger than the key in the internal node.
2507            assert_eq!(
2508                collect(btree.range(key(20)..key(30))),
2509                vec![
2510                    (key(21), value(900)),
2511                    (key(22), value(1_000)),
2512                    (key(23), value(1_100)),
2513                    (key(24), value(1_200)),
2514                ]
2515            );
2516        });
2517    }
2518    btree_test!(test_range_various_prefixes, range_various_prefixes);
2519
2520    fn range_various_prefixes_2<K: TestKey, V: TestValue>() {
2521        let (key, value) = (K::build, V::build);
2522        run_btree_test(|mut btree| {
2523            btree.insert(key(1), value(100));
2524            btree.insert(key(2), value(200));
2525            btree.insert(key(3), value(300));
2526            btree.insert(key(4), value(400));
2527            btree.insert(key(12), value(500));
2528            btree.insert(key(14), value(600));
2529            btree.insert(key(16), value(700));
2530            btree.insert(key(18), value(800));
2531            btree.insert(key(19), value(900));
2532            btree.insert(key(21), value(1000));
2533            btree.insert(key(22), value(1100));
2534            btree.insert(key(23), value(1200));
2535            btree.insert(key(24), value(1300));
2536            btree.insert(key(25), value(1400));
2537            btree.insert(key(26), value(1500));
2538            btree.insert(key(27), value(1600));
2539            btree.insert(key(28), value(1700));
2540            btree.insert(key(29), value(1800));
2541
2542            // The result should look like this:
2543            //                 [14, 23]
2544            //                /    |   \
2545            // [1, 2, 3, 4, 12]    |   [24, 25, 26, 27, 28, 29]
2546            //                     |
2547            //           [16, 18, 19, 21, 22]
2548            let root = btree.load_node(btree.root_addr);
2549            assert_eq!(root.node_type(), NodeType::Internal);
2550            assert_eq!(
2551                root.entries(btree.memory()),
2552                vec![
2553                    (key(14), encode(value(600))),
2554                    (key(23), encode(value(1200)))
2555                ]
2556            );
2557            assert_eq!(root.children_len(), 3);
2558            let child_0 = btree.load_node(root.child(0));
2559            assert_eq!(child_0.node_type(), NodeType::Leaf);
2560            assert_eq!(
2561                child_0.entries(btree.memory()),
2562                vec![
2563                    (key(1), encode(value(100))),
2564                    (key(2), encode(value(200))),
2565                    (key(3), encode(value(300))),
2566                    (key(4), encode(value(400))),
2567                    (key(12), encode(value(500))),
2568                ]
2569            );
2570
2571            let child_1 = btree.load_node(root.child(1));
2572            assert_eq!(child_1.node_type(), NodeType::Leaf);
2573            assert_eq!(
2574                child_1.entries(btree.memory()),
2575                vec![
2576                    (key(16), encode(value(700))),
2577                    (key(18), encode(value(800))),
2578                    (key(19), encode(value(900))),
2579                    (key(21), encode(value(1000))),
2580                    (key(22), encode(value(1100))),
2581                ]
2582            );
2583
2584            let child_2 = btree.load_node(root.child(2));
2585            assert_eq!(
2586                child_2.entries(btree.memory()),
2587                vec![
2588                    (key(24), encode(value(1300))),
2589                    (key(25), encode(value(1400))),
2590                    (key(26), encode(value(1500))),
2591                    (key(27), encode(value(1600))),
2592                    (key(28), encode(value(1700))),
2593                    (key(29), encode(value(1800))),
2594                ]
2595            );
2596
2597            // Tests a prefix that doesn't exist, but is in the middle of the root node.
2598            assert_eq!(collect(btree.range(key(15)..key(16))), vec![]);
2599
2600            // Tests a prefix beginning in the middle of the tree and crossing several nodes.
2601            assert_eq!(
2602                collect(btree.range(key(15)..=key(26))),
2603                vec![
2604                    (key(16), value(700)),
2605                    (key(18), value(800)),
2606                    (key(19), value(900)),
2607                    (key(21), value(1000)),
2608                    (key(22), value(1100)),
2609                    (key(23), value(1200)),
2610                    (key(24), value(1300)),
2611                    (key(25), value(1400)),
2612                    (key(26), value(1500)),
2613                ]
2614            );
2615
2616            // Tests a prefix that crosses several nodes.
2617            assert_eq!(
2618                collect(btree.range(key(10)..key(20))),
2619                vec![
2620                    (key(12), value(500)),
2621                    (key(14), value(600)),
2622                    (key(16), value(700)),
2623                    (key(18), value(800)),
2624                    (key(19), value(900)),
2625                ]
2626            );
2627
2628            // Tests a prefix that starts from a leaf node, then iterates through the root and right
2629            // sibling.
2630            assert_eq!(
2631                collect(btree.range(key(20)..key(30))),
2632                vec![
2633                    (key(21), value(1000)),
2634                    (key(22), value(1100)),
2635                    (key(23), value(1200)),
2636                    (key(24), value(1300)),
2637                    (key(25), value(1400)),
2638                    (key(26), value(1500)),
2639                    (key(27), value(1600)),
2640                    (key(28), value(1700)),
2641                    (key(29), value(1800)),
2642                ]
2643            );
2644        });
2645    }
2646    btree_test!(test_range_various_prefixes_2, range_various_prefixes_2);
2647
2648    fn range_large<K: TestKey, V: TestValue>() {
2649        let (key, value) = (K::build, V::build);
2650        run_btree_test(|mut btree| {
2651            const TOTAL: u32 = 2_000;
2652            const MID: u32 = TOTAL / 2;
2653
2654            for i in 0..TOTAL {
2655                assert_eq!(btree.insert(key(i), value(i)), None);
2656            }
2657
2658            // Check range [0, MID): should yield exactly the first MID entries.
2659            let keys: Vec<_> = btree.range(key(0)..key(MID)).collect();
2660            assert_eq!(keys.len(), MID as usize);
2661            for (i, (k, v)) in keys.into_iter().enumerate() {
2662                let j = i as u32;
2663                assert_eq!(k, key(j));
2664                assert_eq!(v, value(j));
2665            }
2666
2667            // Check range [MID, TOTAL): should yield the remaining entries.
2668            let keys: Vec<_> = btree.range(key(MID)..).collect();
2669            assert_eq!(keys.len(), (TOTAL - MID) as usize);
2670            for (i, (k, v)) in keys.into_iter().enumerate() {
2671                let j = MID + i as u32;
2672                assert_eq!(k, key(j));
2673                assert_eq!(v, value(j));
2674            }
2675        });
2676    }
2677    btree_test!(test_range_large, range_large);
2678
2679    fn range_various_prefixes_with_offset<K: TestKey, V: TestValue>() {
2680        let (key, value) = (K::build, V::build);
2681        run_btree_test(|mut btree| {
2682            btree.insert(key(1), value(100));
2683            btree.insert(key(2), value(200));
2684            btree.insert(key(3), value(300));
2685            btree.insert(key(4), value(400));
2686            btree.insert(key(11), value(500));
2687            btree.insert(key(12), value(600));
2688            btree.insert(key(13), value(700));
2689            btree.insert(key(14), value(800));
2690            btree.insert(key(21), value(900));
2691            btree.insert(key(22), value(1000));
2692            btree.insert(key(23), value(1100));
2693            btree.insert(key(24), value(1200));
2694
2695            // The result should look like this:
2696            //                 [12]
2697            //                /    \
2698            // [1, 2, 3, 4, 11]    [13, 14, 21, 22, 23, 24]
2699
2700            let root = btree.load_node(btree.root_addr);
2701            assert_eq!(root.node_type(), NodeType::Internal);
2702            assert_eq!(
2703                root.entries(btree.memory()),
2704                vec![(key(12), encode(value(600)))]
2705            );
2706            assert_eq!(root.children_len(), 2);
2707
2708            assert_eq!(
2709                collect(btree.range(key(0)..key(10))),
2710                vec![
2711                    (key(1), value(100)),
2712                    (key(2), value(200)),
2713                    (key(3), value(300)),
2714                    (key(4), value(400)),
2715                ]
2716            );
2717
2718            // Tests an offset that has a keys somewhere in the range of keys of an internal node.
2719            assert_eq!(
2720                collect(btree.range(key(13)..key(20))),
2721                vec![(key(13), value(700)), (key(14), value(800))]
2722            );
2723
2724            // Tests an offset that's larger than the key in the internal node.
2725            assert_eq!(collect(btree.range(key(25)..)), vec![]);
2726        });
2727    }
2728    btree_test!(
2729        test_range_various_prefixes_with_offset,
2730        range_various_prefixes_with_offset
2731    );
2732
2733    fn range_various_prefixes_with_offset_2<K: TestKey, V: TestValue>() {
2734        let (key, value) = (K::build, V::build);
2735        run_btree_test(|mut btree| {
2736            btree.insert(key(1), value(0));
2737            btree.insert(key(2), value(0));
2738            btree.insert(key(3), value(0));
2739            btree.insert(key(4), value(0));
2740            btree.insert(key(12), value(0));
2741            btree.insert(key(14), value(0));
2742            btree.insert(key(16), value(0));
2743            btree.insert(key(18), value(0));
2744            btree.insert(key(19), value(0));
2745            btree.insert(key(21), value(0));
2746            btree.insert(key(22), value(0));
2747            btree.insert(key(23), value(0));
2748            btree.insert(key(24), value(0));
2749            btree.insert(key(25), value(0));
2750            btree.insert(key(26), value(0));
2751            btree.insert(key(27), value(0));
2752            btree.insert(key(28), value(0));
2753            btree.insert(key(29), value(0));
2754
2755            // The result should look like this:
2756            //                 [14, 23]
2757            //                /    |   \
2758            // [1, 2, 3, 4, 12]    |   [24, 25, 26, 27, 28, 29]
2759            //                     |
2760            //           [16, 18, 19, 21, 22]
2761            let root = btree.load_node(btree.root_addr);
2762            assert_eq!(root.node_type(), NodeType::Internal);
2763            assert_eq!(
2764                root.entries(btree.memory()),
2765                vec![(key(14), encode(value(0))), (key(23), encode(value(0)))]
2766            );
2767            assert_eq!(root.children_len(), 3);
2768
2769            let child_0 = btree.load_node(root.child(0));
2770            assert_eq!(child_0.node_type(), NodeType::Leaf);
2771            assert_eq!(
2772                child_0.entries(btree.memory()),
2773                vec![
2774                    (key(1), encode(value(0))),
2775                    (key(2), encode(value(0))),
2776                    (key(3), encode(value(0))),
2777                    (key(4), encode(value(0))),
2778                    (key(12), encode(value(0))),
2779                ]
2780            );
2781
2782            let child_1 = btree.load_node(root.child(1));
2783            assert_eq!(child_1.node_type(), NodeType::Leaf);
2784            assert_eq!(
2785                child_1.entries(btree.memory()),
2786                vec![
2787                    (key(16), encode(value(0))),
2788                    (key(18), encode(value(0))),
2789                    (key(19), encode(value(0))),
2790                    (key(21), encode(value(0))),
2791                    (key(22), encode(value(0))),
2792                ]
2793            );
2794
2795            let child_2 = btree.load_node(root.child(2));
2796            assert_eq!(
2797                child_2.entries(btree.memory()),
2798                vec![
2799                    (key(24), encode(value(0))),
2800                    (key(25), encode(value(0))),
2801                    (key(26), encode(value(0))),
2802                    (key(27), encode(value(0))),
2803                    (key(28), encode(value(0))),
2804                    (key(29), encode(value(0))),
2805                ]
2806            );
2807
2808            // Tests a offset that crosses several nodes.
2809            assert_eq!(
2810                collect(btree.range(key(14)..key(20))),
2811                vec![
2812                    (key(14), value(0)),
2813                    (key(16), value(0)),
2814                    (key(18), value(0)),
2815                    (key(19), value(0)),
2816                ]
2817            );
2818
2819            // Tests a offset that starts from a leaf node, then iterates through the root and right
2820            // sibling.
2821            assert_eq!(
2822                collect(btree.range(key(22)..key(30))),
2823                vec![
2824                    (key(22), value(0)),
2825                    (key(23), value(0)),
2826                    (key(24), value(0)),
2827                    (key(25), value(0)),
2828                    (key(26), value(0)),
2829                    (key(27), value(0)),
2830                    (key(28), value(0)),
2831                    (key(29), value(0)),
2832                ]
2833            );
2834        });
2835    }
2836    btree_test!(
2837        test_range_various_prefixes_with_offset_2,
2838        range_various_prefixes_with_offset_2
2839    );
2840
2841    #[test]
2842    #[should_panic(expected = "max_key_size must be <= 4")]
2843    fn v1_rejects_increases_in_max_key_size() {
2844        let mem = make_memory();
2845        let btree: BTreeMap<Blob<4>, Blob<3>, _> = BTreeMap::init_v1(mem);
2846        let _btree: BTreeMap<Blob<5>, Blob<3>, _> = BTreeMap::init_v1(btree.into_memory());
2847    }
2848
2849    #[test]
2850    fn v2_handles_increases_in_max_key_size_and_max_value_size() {
2851        let mem = make_memory();
2852        let mut btree: BTreeMap<Blob<4>, Blob<4>, _> = BTreeMap::init(mem);
2853        btree.insert(
2854            [1u8; 4].as_slice().try_into().unwrap(),
2855            [1u8; 4].as_slice().try_into().unwrap(),
2856        );
2857
2858        // Reinitialize the BTree with larger keys and value sizes.
2859        let mut btree: BTreeMap<Blob<5>, Blob<5>, _> = BTreeMap::init(btree.into_memory());
2860        btree.insert(
2861            [2u8; 5].as_slice().try_into().unwrap(),
2862            [2u8; 5].as_slice().try_into().unwrap(),
2863        );
2864
2865        // Still able to retrieve all the entries inserted.
2866        assert_eq!(
2867            btree.get(&([1u8; 4].as_slice().try_into().unwrap())),
2868            Some([1u8; 4].as_slice().try_into().unwrap())
2869        );
2870
2871        assert_eq!(
2872            btree.get(&([2u8; 5].as_slice().try_into().unwrap())),
2873            Some([2u8; 5].as_slice().try_into().unwrap())
2874        );
2875    }
2876
2877    #[test]
2878    fn accepts_small_or_equal_key_sizes() {
2879        let mem = make_memory();
2880        let btree: BTreeMap<Blob<4>, Blob<3>, _> = BTreeMap::init(mem);
2881        // Smaller key size
2882        let btree: BTreeMap<Blob<3>, Blob<3>, _> = BTreeMap::init(btree.into_memory());
2883        // Equal key size
2884        let _btree: BTreeMap<Blob<4>, Blob<3>, _> = BTreeMap::init(btree.into_memory());
2885    }
2886
2887    #[test]
2888    #[should_panic(expected = "max_value_size must be <= 3")]
2889    fn v1_rejects_larger_value_sizes() {
2890        let mem = make_memory();
2891        let btree: BTreeMap<Blob<4>, Blob<3>, _> = BTreeMap::init_v1(mem);
2892        let _btree: BTreeMap<Blob<4>, Blob<4>, _> = BTreeMap::init_v1(btree.into_memory());
2893    }
2894
2895    #[test]
2896    fn accepts_small_or_equal_value_sizes() {
2897        let mem = make_memory();
2898        let btree: BTreeMap<Blob<4>, Blob<3>, _> = BTreeMap::init(mem);
2899        // Smaller key size
2900        let btree: BTreeMap<Blob<4>, Blob<2>, _> = BTreeMap::init(btree.into_memory());
2901        // Equal key size
2902        let _btree: BTreeMap<Blob<4>, Blob<3>, _> = BTreeMap::init(btree.into_memory());
2903    }
2904
2905    fn bruteforce_range_search<K: TestKey, V: TestValue>() {
2906        let (key, value) = (K::build, V::build);
2907        run_btree_test(|mut stable_map| {
2908            use std::collections::BTreeMap;
2909            const NKEYS: u32 = 60;
2910            let mut std_map = BTreeMap::new();
2911
2912            for i in 0..NKEYS {
2913                std_map.insert(key(i), value(i));
2914                stable_map.insert(key(i), value(i));
2915            }
2916
2917            assert_eq!(collect_kv(std_map.range(..)), collect(stable_map.range(..)));
2918
2919            for l in 0..=NKEYS {
2920                assert_eq!(
2921                    collect_kv(std_map.range(key(l)..)),
2922                    collect(stable_map.range(key(l)..))
2923                );
2924
2925                assert_eq!(
2926                    collect_kv(std_map.range(..key(l))),
2927                    collect(stable_map.range(..key(l)))
2928                );
2929
2930                assert_eq!(
2931                    collect_kv(std_map.range(..=key(l))),
2932                    collect(stable_map.range(..=key(l)))
2933                );
2934
2935                for r in l + 1..=NKEYS {
2936                    for lbound in [Bound::Included(key(l)), Bound::Excluded(key(l))] {
2937                        for rbound in [Bound::Included(key(r)), Bound::Excluded(key(r))] {
2938                            let range = (lbound.clone(), rbound);
2939                            assert_eq!(
2940                                collect_kv(std_map.range(range.clone())),
2941                                collect(stable_map.range(range.clone())),
2942                                "range: {range:?}"
2943                            );
2944                        }
2945                    }
2946                }
2947            }
2948        });
2949    }
2950    btree_test!(test_bruteforce_range_search, bruteforce_range_search);
2951
2952    fn test_iter_upper_bound<K: TestKey, V: TestValue>() {
2953        let (key, value) = (K::build, V::build);
2954        run_btree_test(|mut btree| {
2955            for j in 0..100 {
2956                btree.insert(key(j), value(j));
2957                for i in 0..=j {
2958                    assert_eq!(
2959                        btree.iter_upper_bound(&key(i + 1)).next(),
2960                        Some((key(i), value(i))),
2961                        "failed to get an upper bound for key({})",
2962                        i + 1
2963                    );
2964                }
2965                assert_eq!(
2966                    btree.iter_upper_bound(&key(0)).next(),
2967                    None,
2968                    "key(0) must not have an upper bound"
2969                );
2970            }
2971        });
2972    }
2973    btree_test!(test_test_iter_upper_bound, test_iter_upper_bound);
2974
2975    // A buggy implementation of storable where the max_size is smaller than the serialized size.
2976    #[derive(Clone, Ord, PartialOrd, Eq, PartialEq)]
2977    struct BuggyStruct;
2978    impl crate::Storable for BuggyStruct {
2979        fn to_bytes(&self) -> Cow<[u8]> {
2980            Cow::Borrowed(&[1, 2, 3, 4])
2981        }
2982
2983        fn from_bytes(_: Cow<[u8]>) -> Self {
2984            unimplemented!();
2985        }
2986
2987        const BOUND: StorableBound = StorableBound::Bounded {
2988            max_size: 1,
2989            is_fixed_size: false,
2990        };
2991    }
2992
2993    #[test]
2994    #[should_panic(expected = "expected an element with length <= 1 bytes, but found 4")]
2995    fn v1_panics_if_key_is_bigger_than_max_size() {
2996        let mut btree = BTreeMap::init_v1(make_memory());
2997        btree.insert(BuggyStruct, ());
2998    }
2999
3000    #[test]
3001    #[should_panic(expected = "expected an element with length <= 1 bytes, but found 4")]
3002    fn v2_panics_if_key_is_bigger_than_max_size() {
3003        let mut btree = BTreeMap::init(make_memory());
3004        btree.insert(BuggyStruct, ());
3005    }
3006
3007    #[test]
3008    #[should_panic(expected = "expected an element with length <= 1 bytes, but found 4")]
3009    fn v1_panics_if_value_is_bigger_than_max_size() {
3010        let mut btree = BTreeMap::init(make_memory());
3011        btree.insert((), BuggyStruct);
3012    }
3013
3014    #[test]
3015    #[should_panic(expected = "expected an element with length <= 1 bytes, but found 4")]
3016    fn v2_panics_if_value_is_bigger_than_max_size() {
3017        let mut btree = BTreeMap::init(make_memory());
3018        btree.insert((), BuggyStruct);
3019    }
3020
3021    // To generate the memory dump file for the current version:
3022    //   cargo test create_btreemap_dump_file -- --include-ignored
3023    #[test]
3024    #[ignore]
3025    fn create_btreemap_dump_file() {
3026        let mem = make_memory();
3027        let mut btree = BTreeMap::init_v1(mem.clone());
3028
3029        let key = b(&[1, 2, 3]);
3030        let value = b(&[4, 5, 6]);
3031        assert_eq!(btree.insert(key, value), None);
3032        assert_eq!(btree.get(&key), Some(value));
3033
3034        use std::io::prelude::*;
3035        let mut file =
3036            std::fs::File::create(format!("dumps/btreemap_v{LAYOUT_VERSION}.dump")).unwrap();
3037        file.write_all(&mem.borrow()).unwrap();
3038    }
3039
3040    #[test]
3041    fn produces_layout_identical_to_layout_version_1_with_packed_headers() {
3042        let mem = make_memory();
3043        let mut btree = BTreeMap::init_v1(mem.clone());
3044
3045        let key = b(&[1, 2, 3]);
3046        let value = b(&[4, 5, 6]);
3047        assert_eq!(btree.insert(key, value), None);
3048        assert_eq!(btree.get(&key), Some(value));
3049
3050        let btreemap_v1 = include_bytes!("../dumps/btreemap_v1_packed_headers.dump");
3051        assert_eq!(*mem.borrow(), btreemap_v1);
3052    }
3053
3054    #[test]
3055    fn read_write_header_is_identical_to_read_write_struct() {
3056        #[repr(C, packed)]
3057        struct BTreePackedHeader {
3058            magic: [u8; 3],
3059            version: u8,
3060            max_key_size: u32,
3061            max_value_size: u32,
3062            root_addr: Address,
3063            length: u64,
3064            _buffer: [u8; 24],
3065        }
3066        let packed_header = BTreePackedHeader {
3067            magic: *MAGIC,
3068            version: LAYOUT_VERSION,
3069            root_addr: Address::from(0xDEADBEEF),
3070            max_key_size: 0x12345678,
3071            max_value_size: 0x87654321,
3072            length: 0xA1B2D3C4,
3073            _buffer: [0; 24],
3074        };
3075
3076        let packed_mem = make_memory();
3077        crate::write_struct(&packed_header, Address::from(0), &packed_mem);
3078
3079        let v1_header = BTreeHeader {
3080            version: Version::V1(DerivedPageSize {
3081                max_key_size: 0x12345678,
3082                max_value_size: 0x87654321,
3083            }),
3084            root_addr: Address::from(0xDEADBEEF),
3085            length: 0xA1B2D3C4,
3086        };
3087
3088        let v1_mem = make_memory();
3089        BTreeMap::<Vec<_>, Vec<_>, RefCell<Vec<_>>>::write_header(&v1_header, &v1_mem);
3090
3091        assert_eq!(packed_mem, v1_mem);
3092
3093        let packed_header: BTreePackedHeader = crate::read_struct(Address::from(0), &v1_mem);
3094        let v1_header = BTreeMap::<Vec<_>, Vec<_>, RefCell<Vec<_>>>::read_header(&v1_mem);
3095        assert!(packed_header.magic == *MAGIC);
3096        assert!(packed_header.version == LAYOUT_VERSION);
3097        match v1_header.version {
3098            Version::V1(DerivedPageSize {
3099                max_key_size,
3100                max_value_size,
3101            }) => {
3102                assert!(packed_header.max_key_size == max_key_size);
3103                assert!(packed_header.max_value_size == max_value_size);
3104            }
3105            _ => unreachable!("version must be v1"),
3106        };
3107
3108        assert!(packed_header.root_addr == v1_header.root_addr);
3109        assert!(packed_header.length == v1_header.length);
3110    }
3111
3112    #[test]
3113    fn migrate_from_bounded_to_unbounded_and_back() {
3114        // A type that is bounded.
3115        #[derive(PartialOrd, Ord, Clone, Eq, PartialEq, Debug)]
3116        struct T;
3117        impl Storable for T {
3118            fn to_bytes(&self) -> Cow<[u8]> {
3119                Cow::Owned(vec![1, 2, 3])
3120            }
3121
3122            fn from_bytes(bytes: Cow<[u8]>) -> Self {
3123                assert_eq!(bytes.to_vec(), vec![1, 2, 3]);
3124                T
3125            }
3126
3127            const BOUND: StorableBound = StorableBound::Bounded {
3128                max_size: 3,
3129                is_fixed_size: true,
3130            };
3131        }
3132
3133        // Same as the above type, but unbounded.
3134        #[derive(PartialOrd, Ord, Clone, Eq, PartialEq, Debug)]
3135        struct T2;
3136        impl Storable for T2 {
3137            fn to_bytes(&self) -> Cow<[u8]> {
3138                Cow::Owned(vec![1, 2, 3])
3139            }
3140
3141            fn from_bytes(bytes: Cow<[u8]>) -> Self {
3142                assert_eq!(bytes.to_vec(), vec![1, 2, 3]);
3143                T2
3144            }
3145
3146            const BOUND: StorableBound = StorableBound::Unbounded;
3147        }
3148
3149        // Create a v1 btreemap with the bounded type.
3150        let mem = make_memory();
3151        let mut btree: BTreeMap<T, T, _> = BTreeMap::new_v1(mem);
3152        btree.insert(T, T);
3153
3154        // Migrate to v2 and the unbounded type.
3155        let btree: BTreeMap<T2, T2, _> = BTreeMap::init(btree.into_memory());
3156        btree.save_header();
3157
3158        // Reload the BTree again and try to read the value.
3159        let btree: BTreeMap<T2, T2, _> = BTreeMap::init(btree.into_memory());
3160        assert_eq!(btree.get(&T2), Some(T2));
3161
3162        // Reload the BTree again with bounded type.
3163        let btree: BTreeMap<T, T, _> = BTreeMap::init(btree.into_memory());
3164        assert_eq!(btree.get(&T), Some(T));
3165    }
3166
3167    #[test]
3168    fn test_clear_new_bounded_type() {
3169        let mem = make_memory();
3170        let mut btree: BTreeMap<Blob<4>, Blob<4>, _> = BTreeMap::new(mem.clone());
3171
3172        btree.insert(
3173            [1u8; 4].as_slice().try_into().unwrap(),
3174            [1u8; 4].as_slice().try_into().unwrap(),
3175        );
3176
3177        assert_ne!(btree.len(), 0);
3178        assert_ne!(btree.allocator.num_allocated_chunks(), 0);
3179        assert_ne!(btree.root_addr, NULL);
3180
3181        btree.clear_new();
3182
3183        let header_actual = BTreeMap::<Blob<4>, Blob<4>, _>::read_header(&mem);
3184
3185        BTreeMap::<Blob<4>, Blob<4>, _>::new(mem.clone());
3186
3187        let header_expected = BTreeMap::<Blob<4>, Blob<4>, _>::read_header(&mem);
3188
3189        assert_eq!(header_actual, header_expected);
3190    }
3191
3192    #[test]
3193    fn test_clear_new_unbounded_type() {
3194        let mem = make_memory();
3195        let mut btree: BTreeMap<String, String, _> = BTreeMap::new(mem.clone());
3196        btree.insert("asd".into(), "bce".into());
3197
3198        assert_ne!(btree.len(), 0);
3199        assert_ne!(btree.allocator.num_allocated_chunks(), 0);
3200        assert_ne!(btree.root_addr, NULL);
3201
3202        btree.clear_new();
3203
3204        let header_actual = BTreeMap::<String, String, _>::read_header(&mem);
3205
3206        BTreeMap::<String, String, _>::new(mem.clone());
3207
3208        let header_expected = BTreeMap::<String, String, _>::read_header(&mem);
3209
3210        assert_eq!(header_actual, header_expected);
3211    }
3212
3213    #[test]
3214    fn deallocating_node_with_overflows() {
3215        let mem = make_memory();
3216        let mut btree: BTreeMap<Vec<u8>, Vec<u8>, _> = BTreeMap::new(mem.clone());
3217
3218        // No allocated chunks yet.
3219        assert_eq!(btree.allocator.num_allocated_chunks(), 0);
3220
3221        // Insert and remove an entry that's large and requires overflow pages.
3222        btree.insert(vec![0; 10_000], vec![]);
3223
3224        // At least two chunks should be allocated.
3225        // One for the node itself and at least one overflow page.
3226        assert!(btree.allocator.num_allocated_chunks() >= 2);
3227        btree.remove(&vec![0; 10_000]);
3228
3229        // All chunks have been deallocated.
3230        assert_eq!(btree.allocator.num_allocated_chunks(), 0);
3231    }
3232
3233    #[test]
3234    fn repeatedly_deallocating_nodes_with_overflows() {
3235        let mem = make_memory();
3236        let mut btree: BTreeMap<Vec<u8>, Vec<u8>, _> = BTreeMap::new(mem.clone());
3237
3238        // No allocated chunks yet.
3239        assert_eq!(btree.allocator.num_allocated_chunks(), 0);
3240
3241        for _ in 0..100 {
3242            for i in 0..100 {
3243                btree.insert(vec![i; 10_000], vec![]);
3244            }
3245
3246            for i in 0..100 {
3247                btree.remove(&vec![i; 10_000]);
3248            }
3249        }
3250
3251        // All chunks have been deallocated.
3252        assert_eq!(btree.allocator.num_allocated_chunks(), 0);
3253    }
3254
3255    #[test]
3256    fn deallocating_root_does_not_leak_memory() {
3257        let mem = make_memory();
3258        let mut btree: BTreeMap<Vec<u8>, _, _> = BTreeMap::new(mem.clone());
3259
3260        for i in 1..=11 {
3261            // Large keys are stored so that each node overflows.
3262            assert_eq!(btree.insert(vec![i; 10_000], ()), None);
3263        }
3264
3265        // Should now split a node.
3266        assert_eq!(btree.insert(vec![0; 10_000], ()), None);
3267
3268        // The btree should look like this:
3269        //                    [6]
3270        //                   /   \
3271        // [0, 1, 2, 3, 4, 5]     [7, 8, 9, 10, 11]
3272        let root = btree.load_node(btree.root_addr);
3273        assert_eq!(root.node_type(), NodeType::Internal);
3274        assert_eq!(root.keys(btree.memory()), vec![&[6; 10_000]]);
3275        assert_eq!(root.children_len(), 2);
3276
3277        // Remove the element in the root.
3278        btree.remove(&vec![6; 10_000]);
3279
3280        // The btree should look like this:
3281        //                 [5]
3282        //                /   \
3283        // [0, 1, 2, 3, 4]     [7, 8, 9, 10, 11]
3284        let root = btree.load_node(btree.root_addr);
3285        assert_eq!(root.node_type(), NodeType::Internal);
3286        assert_eq!(root.keys(btree.memory()), vec![&[5; 10_000]]);
3287        assert_eq!(root.children_len(), 2);
3288
3289        // Remove the element in the root. This triggers the case where the root
3290        // node is deallocated and the children are merged into a single node.
3291        btree.remove(&vec![5; 10_000]);
3292
3293        // The btree should look like this:
3294        //      [0, 1, 2, 3, 4, 7, 8, 9, 10, 11]
3295        let root = btree.load_node(btree.root_addr);
3296        assert_eq!(root.node_type(), NodeType::Leaf);
3297        assert_eq!(
3298            root.keys(btree.memory()),
3299            vec![
3300                &[0; 10_000],
3301                &[1; 10_000],
3302                &[2; 10_000],
3303                &[3; 10_000],
3304                &[4; 10_000],
3305                &[7; 10_000],
3306                &[8; 10_000],
3307                &[9; 10_000],
3308                &[10; 10_000],
3309                &[11; 10_000],
3310            ]
3311        );
3312
3313        // Delete everything else.
3314        for i in 0..=11 {
3315            btree.remove(&vec![i; 10_000]);
3316        }
3317
3318        assert_eq!(btree.allocator.num_allocated_chunks(), 0);
3319    }
3320}