ic_stable_structures/
btreemap.rs

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