ic-storage 0.68.0

// Adapted from ic-stable-structures 0.7.2, DFINITY Foundation.
// SPDX-License-Identifier: Apache-2.0
// See the package LICENSE-APACHE and docs/storage/icss-source.json.
//! This module implements a key/value store based on a B-Tree
//! in stable memory.
//!
//! # V2 layout
//!
//! ```text
//! ---------------------------------------- <- Address 0
//! Magic "BTR"                 ↕ 3 bytes
//! ----------------------------------------
//! Layout version              ↕ 1 byte
//! ----------------------------------------
//! Max key size                ↕ 4 bytes             Page size                   ↕ 4 bytes
//! ----------------------------------------   OR   ----------------------------------------
//! Max value size              ↕ 4 bytes             PAGE_SIZE_VALUE_MARKER      ↕ 4 bytes
//! ----------------------------------------
//! Root node address           ↕ 8 bytes
//! ----------------------------------------
//! Length (number of elements) ↕ 8 bytes
//! ---------------------------------------- <- Address 28 (PACKED_HEADER_SIZE)
//! Reserved space              ↕ 24 bytes
//! ---------------------------------------- <- Address 52 (ALLOCATOR_OFFSET)
//! Allocator
//! ----------------------------------------
//! ... free memory for nodes
//! ----------------------------------------
//! ```
//!
//! # V1 layout
//!
//! ```text
//! ---------------------------------------- <- Address 0
//! Magic "BTR"                 ↕ 3 bytes
//! ----------------------------------------
//! Layout version              ↕ 1 byte
//! ----------------------------------------
//! Max key size                ↕ 4 bytes
//! ----------------------------------------
//! Max value size              ↕ 4 bytes
//! ----------------------------------------
//! Root node address           ↕ 8 bytes
//! ----------------------------------------
//! Length (number of elements) ↕ 8 bytes
//! ---------------------------------------- <- Address 28 (PACKED_HEADER_SIZE)
//! Reserved space              ↕ 24 bytes
//! ---------------------------------------- <- Address 52 (ALLOCATOR_OFFSET)
//! Allocator
//! ----------------------------------------
//! ... free memory for nodes
//! ----------------------------------------
//! ```
mod allocator;
mod iter;
mod node;
use crate::engine::btreemap::iter::{IterInternal, KeysIter, ValuesIter};
use crate::engine::{
    Memory, Storable,
    storable::Bound as StorableBound,
    types::{Address, NULL},
};
use allocator::Allocator;
pub use iter::{Iter, LazyEntry};
use node::{Entry, Node, NodeType, PageSize, Version};
use std::borrow::Cow;
use std::marker::PhantomData;
use std::ops::RangeBounds;

const MAGIC: &[u8; 3] = b"BTR";
const LAYOUT_VERSION_2: u8 = 2;
// The sum of all the header fields, i.e. size of a packed header.
const PACKED_HEADER_SIZE: usize = 28;
// The offset where the allocator begins.
const ALLOCATOR_OFFSET: usize = 52;

// The default page size to use in BTreeMap V2 in bytes.
const DEFAULT_PAGE_SIZE: u32 = 1024;

// A marker to indicate that the `PageSize` stored in the header is a `PageSize::Value`.
const PAGE_SIZE_VALUE_MARKER: u32 = u32::MAX;
pub struct BTreeMap<K, V, M>
where
    K: Storable + Ord + Clone,
    V: Storable,
    M: Memory,
{
    // The address of the root node. If a root node doesn't exist, the address
    // is set to NULL.
    root_addr: Address,

    version: Version,

    // An allocator used for managing memory and allocating nodes.
    allocator: Allocator<M>,

    // The number of elements in the map.
    length: u64,

    // A marker to communicate to the Rust compiler that we own these types.
    _phantom: PhantomData<(K, V)>,
}

#[derive(PartialEq, Debug)]
struct BTreeHeader {
    version: Version,
    root_addr: Address,
    length: u64,
    // Reserved bytes for future extensions
}

impl<K, V, M> BTreeMap<K, V, M>
where
    K: Storable + Ord + Clone,
    V: Storable,
    M: Memory,
{
    pub fn new(memory: M) -> Self {
        let page_size = match (K::BOUND, V::BOUND) {
            // The keys and values are both bounded.
            (
                StorableBound::Bounded {
                    max_size: max_key_size,
                    ..
                },
                StorableBound::Bounded {
                    max_size: max_value_size,
                    ..
                },
            ) => {
                // Get the maximum possible node size.
                let max_node_size = Node::<K>::max_size(max_key_size, max_value_size).get();

                // A node can have at most 11 entries, and an analysis has shown that ~70% of all
                // nodes have <= 8 entries. We can therefore use a page size that's 8/11 the
                // maximum size with little performance difference but with a significant storage
                // saving. We round the 8/11 to be 3/4.
                PageSize::Value(
                    u32::try_from(max_node_size * 3 / 4).expect("bounded node page size"),
                )
            }
            // Use a default page size.
            _ => PageSize::Value(DEFAULT_PAGE_SIZE),
        };

        let btree = Self {
            root_addr: NULL,
            allocator: Allocator::new(
                memory,
                Address::from(ALLOCATOR_OFFSET as u64),
                page_size.get().into(),
            ),
            version: Version::V2(page_size),
            length: 0,
            _phantom: PhantomData,
        };

        btree.save_header();
        btree
    }

    pub fn load(memory: M) -> Self {
        let header = Self::read_header(&memory);
        let allocator_addr = Address::from(ALLOCATOR_OFFSET as u64);
        Self {
            root_addr: header.root_addr,
            allocator: Allocator::load(memory, allocator_addr),
            version: header.version,
            length: header.length,
            _phantom: PhantomData,
        }
    }

    fn read_header(memory: &M) -> BTreeHeader {
        // Read the header
        let mut buf = [0; PACKED_HEADER_SIZE];
        memory.read(0, &mut buf);

        assert_eq!(&buf[0..3], MAGIC, "Bad magic.");

        match buf[3] {
            LAYOUT_VERSION_2 => {
                // Load the page size.
                let page_size = {
                    // Page sizes can be stored either as a direct value or as max/value sizes.
                    let a = u32::from_le_bytes(buf[4..8].try_into().unwrap());
                    let b = u32::from_le_bytes(buf[8..12].try_into().unwrap());

                    assert_eq!(b, PAGE_SIZE_VALUE_MARKER, "Unsupported derived-page layout");
                    PageSize::Value(a)
                };

                // Deserialize the fields
                BTreeHeader {
                    version: Version::V2(page_size),
                    root_addr: Address::from(u64::from_le_bytes(buf[12..20].try_into().unwrap())),
                    length: u64::from_le_bytes(buf[20..28].try_into().unwrap()),
                }
            }
            version => {
                panic!("Unsupported version: {version}.");
            }
        }
    }
    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
        let value = value.into_bytes_checked();

        let root = if self.root_addr == NULL {
            // No root present. Allocate one.
            let node = self.allocate_node(NodeType::Leaf);
            self.root_addr = node.address();
            self.save_header();
            node
        } else {
            // Load the root from memory.
            let mut root = self.load_node(self.root_addr);

            // Check if the key already exists in the root.
            if let Ok(idx) = root.search(&key, self.memory()) {
                // Key found, replace its value and return the old one.
                return Some(V::from_bytes(Cow::Owned(
                    self.update_value(&mut root, idx, value),
                )));
            }

            // If the root is full, we need to introduce a new node as the root.
            //
            // NOTE: In the case where we are overwriting an existing key, then introducing
            // a new root node isn't strictly necessary. However, that's a micro-optimization
            // that adds more complexity than it's worth.
            if root.is_full() {
                // The root is full. Allocate a new node that will be used as the new root.
                let mut new_root = self.allocate_node(NodeType::Internal);

                // The new root has the old root as its only child.
                new_root.push_child(self.root_addr);

                // Update the root address.
                self.root_addr = new_root.address();
                self.save_header();

                // Split the old (full) root.
                self.split_child(&mut new_root, 0);

                new_root
            } else {
                root
            }
        };

        self.insert_nonfull(root, key, value)
            .map(Cow::Owned)
            .map(V::from_bytes)
    }
    fn insert_nonfull(&mut self, mut node: Node<K>, key: K, value: Vec<u8>) -> Option<Vec<u8>> {
        // We're guaranteed by the caller that the provided node is not full.
        assert!(!node.is_full());

        // Look for the key in the node.
        match node.search(&key, self.memory()) {
            Ok(idx) => {
                // Key found, replace its value and return the old one.
                Some(self.update_value(&mut node, idx, value))
            }
            Err(idx) => {
                // The key isn't in the node. `idx` is where that key should be inserted.

                match node.node_type() {
                    NodeType::Leaf => {
                        // The node is a non-full leaf.
                        // Insert the entry at the proper location.
                        node.insert_entry(idx, (key, value));
                        self.save_node(&mut node);

                        // Update the length.
                        self.length += 1;
                        self.save_header();

                        // No previous value to return.
                        None
                    }
                    NodeType::Internal => {
                        // The node is an internal node.
                        // Load the child that we should add the entry to.
                        let mut child = self.load_node(node.child(idx));

                        if child.is_full() {
                            // Check if the key already exists in the child.
                            if let Ok(idx) = child.search(&key, self.memory()) {
                                // Key found, replace its value and return the old one.
                                return Some(self.update_value(&mut child, idx, value));
                            }

                            // The child is full. Split the child.
                            self.split_child(&mut node, idx);

                            // The children have now changed. Search again for
                            // the child where we need to store the entry in.
                            let idx = node.search(&key, self.memory()).unwrap_or_else(|idx| idx);
                            child = self.load_node(node.child(idx));
                        }

                        // The child should now be not full.
                        assert!(!child.is_full());

                        self.insert_nonfull(child, key, value)
                    }
                }
            }
        }
    }
    fn split_child(&mut self, node: &mut Node<K>, full_child_idx: usize) {
        // The node must not be full.
        assert!(!node.is_full());

        // The node's child must be full.
        let mut full_child = self.load_node(node.child(full_child_idx));
        assert!(full_child.is_full());

        // Create a sibling to this full child (which has to be the same type).
        let mut sibling = self.allocate_node(full_child.node_type());
        assert_eq!(sibling.node_type(), full_child.node_type());

        // Add sibling as a new child in the node.
        node.insert_child(full_child_idx + 1, sibling.address());

        let (median_key, median_value) = full_child.split(&mut sibling, self.memory());

        node.insert_entry(full_child_idx, (median_key, median_value));

        self.save_node(&mut sibling);
        self.save_node(&mut full_child);
        self.save_node(node);
    }
    pub fn get(&self, key: &K) -> Option<V> {
        if self.root_addr == NULL {
            return None;
        }
        self.traverse(self.root_addr, key, |node, idx| {
            node.extract_entry_at(idx, self.memory()).1 // Extract value.
        })
        .map(Cow::Owned)
        .map(V::from_bytes)
    }
    pub fn contains_key(&self, key: &K) -> bool {
        // An empty closure returns Some(()) if the key is found.
        self.root_addr != NULL && self.traverse(self.root_addr, key, |_, _| ()).is_some()
    }
    fn traverse<F, R>(&self, node_addr: Address, key: &K, f: F) -> Option<R>
    where
        F: Fn(&mut Node<K>, usize) -> R,
    {
        let mut node = self.load_node(node_addr);
        // Look for the key in the current node.
        match node.search(key, self.memory()) {
            Ok(idx) => Some(f(&mut node, idx)), // Key found: apply `f`.
            Err(idx) => match node.node_type() {
                NodeType::Leaf => None, // At a leaf: key not present.
                NodeType::Internal => self.traverse(node.child(idx), key, f), // Continue search in child.
            },
        }
    }
    pub const fn is_empty(&self) -> bool {
        self.length == 0
    }
    pub const fn len(&self) -> u64 {
        self.length
    }

    // TODO: In next major release (v1.0), remove this deprecated method and rename
    // `clear_new` to `clear` for consistency with standard Rust collections.

    // TODO: In next major release (v1.0), rename this method to `clear` to follow
    // standard Rust collection naming conventions.
    pub fn clear_new(&mut self) {
        self.root_addr = NULL;
        self.length = 0;
        self.allocator.clear();
        self.save_header();
    }
    pub fn first_key_value(&self) -> Option<(K, V)> {
        if self.root_addr == NULL {
            return None;
        }
        let root = self.load_node(self.root_addr);
        let (k, encoded_v) = root.get_min(self.memory());
        Some((k, V::from_bytes(Cow::Owned(encoded_v))))
    }
    pub fn last_key_value(&self) -> Option<(K, V)> {
        if self.root_addr == NULL {
            return None;
        }
        let root = self.load_node(self.root_addr);
        let (k, encoded_v) = root.get_max(self.memory());
        Some((k, V::from_bytes(Cow::Owned(encoded_v))))
    }

    pub(crate) const fn memory(&self) -> &M {
        self.allocator.memory()
    }

    const fn allocator_mut(&mut self) -> &mut Allocator<M> {
        &mut self.allocator
    }
    pub fn remove(&mut self, key: &K) -> Option<V> {
        if self.root_addr == NULL {
            return None;
        }

        let root_node = self.load_node(self.root_addr);
        self.remove_helper(root_node, key)
            .map(Cow::Owned)
            .map(V::from_bytes)
    }

    #[expect(
        clippy::too_many_lines,
        reason = "Keep the upstream deletion algorithm intact for the baseline"
    )]
    fn remove_helper(&mut self, mut node: Node<K>, key: &K) -> Option<Vec<u8>> {
        if node.address() != self.root_addr {
            // We're guaranteed that whenever this method is called an entry can be
            // removed from the node without it needing to be merged into a sibling.
            // This strengthened condition allows us to delete an entry in a single
            // pass most of the time without having to back up.
            assert!(node.can_remove_entry_without_merging());
        }

        match node.node_type() {
            NodeType::Leaf => {
                match node.search(key, self.memory()) {
                    Ok(idx) => {
                        // Case 1: The node is a leaf node and the key exists in it.
                        // This is the simplest case. The key is removed from the leaf.
                        let value = node.remove_entry(idx, self.memory()).1;
                        self.length -= 1;

                        if node.entries_len() == 0 {
                            assert_eq!(
                                node.address(),
                                self.root_addr,
                                "Removal can only result in an empty leaf node if that node is the root"
                            );

                            // Deallocate the empty node.
                            self.deallocate_node(node);
                            self.root_addr = NULL;
                        } else {
                            self.save_node(&mut node);
                        }

                        self.save_header();
                        Some(value)
                    }
                    _ => None, // Key not found.
                }
            }
            NodeType::Internal => {
                match node.search(key, self.memory()) {
                    Ok(idx) => {
                        // Case 2: The node is an internal node and the key exists in it.

                        let left_child = self.load_node(node.child(idx));
                        if left_child.can_remove_entry_without_merging() {
                            // Case 2.a: A key can be removed from the left child without merging.
                            //
                            //                       parent
                            //                  [..., key, ...]
                            //                       /   \
                            //            [left child]   [...]
                            //           /            \
                            //        [...]         [..., key predecessor]
                            //
                            // In this case, we replace `key` with the key's predecessor from the
                            // left child's subtree, then we recursively delete the key's
                            // predecessor for the following end result:
                            //
                            //                       parent
                            //            [..., key predecessor, ...]
                            //                       /   \
                            //            [left child]   [...]
                            //           /            \
                            //        [...]          [...]

                            // Recursively delete the predecessor.
                            // TODO(EXC-1034): Do this in a single pass.
                            let predecessor = left_child.get_max(self.memory());
                            self.remove_helper(left_child, &predecessor.0)?;

                            // Replace the `key` with its predecessor.
                            let (_, old_value) = node.swap_entry(idx, predecessor, self.memory());

                            // Save the parent node.
                            self.save_node(&mut node);
                            return Some(old_value);
                        }

                        let right_child = self.load_node(node.child(idx + 1));
                        if right_child.can_remove_entry_without_merging() {
                            // Case 2.b: A key can be removed from the right child without merging.
                            //
                            //                       parent
                            //                  [..., key, ...]
                            //                       /   \
                            //                   [...]   [right child]
                            //                          /             \
                            //              [key successor, ...]     [...]
                            //
                            // In this case, we replace `key` with the key's successor from the
                            // right child's subtree, then we recursively delete the key's
                            // successor for the following end result:
                            //
                            //                       parent
                            //            [..., key successor, ...]
                            //                       /   \
                            //                  [...]   [right child]
                            //                           /            \
                            //                        [...]          [...]

                            // Recursively delete the successor.
                            // TODO(EXC-1034): Do this in a single pass.
                            let successor = right_child.get_min(self.memory());
                            self.remove_helper(right_child, &successor.0)?;

                            // Replace the `key` with its successor.
                            let (_, old_value) = node.swap_entry(idx, successor, self.memory());

                            // Save the parent node.
                            self.save_node(&mut node);
                            return Some(old_value);
                        }

                        // Case 2.c: Both the left and right child are at their minimum sizes.
                        //
                        //                       parent
                        //                  [..., key, ...]
                        //                       /   \
                        //            [left child]   [right child]
                        //
                        // In this case, we merge (left child, key, right child) into a single
                        // node. The result will look like this:
                        //
                        //                       parent
                        //                     [...  ...]
                        //                         |
                        //          [left child, `key`, right child] <= new child
                        //
                        // We then recurse on this new child to delete `key`.
                        //
                        // If `parent` becomes empty (which can only happen if it's the root),
                        // then `parent` is deleted and `new_child` becomes the new root.
                        assert!(left_child.at_minimum());
                        assert!(right_child.at_minimum());

                        // Merge the right child into the left child.
                        let mut new_child = self.merge(
                            right_child,
                            left_child,
                            node.remove_entry(idx, self.memory()),
                        );

                        // Remove the right child from the parent node.
                        node.remove_child(idx + 1);

                        if node.entries_len() == 0 {
                            // Can only happen if this node is root.
                            assert_eq!(node.address(), self.root_addr);
                            assert_eq!(node.child(0), new_child.address());
                            assert_eq!(node.children_len(), 1);

                            self.root_addr = new_child.address();

                            // Deallocate the root node.
                            self.deallocate_node(node);
                            self.save_header();
                        } else {
                            self.save_node(&mut node);
                        }

                        self.save_node(&mut new_child);

                        // Recursively delete the key.
                        self.remove_helper(new_child, key)
                    }
                    Err(idx) => {
                        // Case 3: The node is an internal node and the key does NOT exist in it.

                        // If the key does exist in the tree, it will exist in the subtree at index
                        // `idx`.
                        let mut child = self.load_node(node.child(idx));

                        if child.can_remove_entry_without_merging() {
                            // The child has enough nodes. Recurse to delete the `key` from the
                            // `child`.
                            return self.remove_helper(child, key);
                        }

                        // An entry can't be removed from the child without merging.
                        // See if it has a sibling where an entry can be removed without merging.
                        let mut left_sibling = if idx > 0 {
                            Some(self.load_node(node.child(idx - 1)))
                        } else {
                            None
                        };

                        let mut right_sibling = if idx + 1 < node.children_len() {
                            Some(self.load_node(node.child(idx + 1)))
                        } else {
                            None
                        };

                        if let Some(ref mut left_sibling) = left_sibling
                            && left_sibling.can_remove_entry_without_merging()
                        {
                            // Case 3.a (left):
                            // A key can be removed from the left child without merging.
                            //
                            //                            [d] (parent)
                            //                           /   \
                            //  (left sibling) [a, b, c]     [e, f] (child)
                            //                         \
                            //                         [c']
                            //
                            // In this case, we move a key down from the parent into the child
                            // and move a key from the left sibling up into the parent
                            // resulting in the following tree:
                            //
                            //                            [c] (parent)
                            //                           /   \
                            //       (left sibling) [a, b]   [d, e, f] (child)
                            //                              /
                            //                            [c']
                            //
                            // We then recurse to delete the key from the child.

                            // Remove the last entry from the left sibling.
                            let (left_sibling_key, left_sibling_value) =
                                left_sibling.pop_entry(self.memory()).unwrap();

                            // Replace the parent's entry with the one from the left sibling.
                            let (parent_key, parent_value) = node.swap_entry(
                                idx - 1,
                                (left_sibling_key, left_sibling_value),
                                self.memory(),
                            );

                            // Move the entry from the parent into the child.
                            child.insert_entry(0, (parent_key, parent_value));

                            // Move the last child from left sibling into child.
                            if let Some(last_child) = left_sibling.pop_child() {
                                assert_eq!(left_sibling.node_type(), NodeType::Internal);
                                assert_eq!(child.node_type(), NodeType::Internal);

                                child.insert_child(0, last_child);
                            } else {
                                assert_eq!(left_sibling.node_type(), NodeType::Leaf);
                                assert_eq!(child.node_type(), NodeType::Leaf);
                            }

                            self.save_node(left_sibling);
                            self.save_node(&mut child);
                            self.save_node(&mut node);
                            return self.remove_helper(child, key);
                        }

                        if let Some(right_sibling) = &mut right_sibling
                            && right_sibling.can_remove_entry_without_merging()
                        {
                            // Case 3.a (right):
                            // A key can be removed from the right child without merging.
                            //
                            //                            [c] (parent)
                            //                           /   \
                            //             (child) [a, b]     [d, e, f] (right sibling)
                            //                               /
                            //                            [d']
                            //
                            // In this case, we move a key down from the parent into the child
                            // and move a key from the right sibling up into the parent
                            // resulting in the following tree:
                            //
                            //                            [d] (parent)
                            //                           /   \
                            //          (child) [a, b, c]     [e, f] (right sibling)
                            //                          \
                            //                           [d']
                            //
                            // We then recurse to delete the key from the child.

                            // Remove the first entry from the right sibling.
                            let (right_sibling_key, right_sibling_value) =
                                right_sibling.remove_entry(0, self.memory());

                            // Replace the parent's entry with the one from the right sibling.
                            let parent_entry = node.swap_entry(
                                idx,
                                (right_sibling_key, right_sibling_value),
                                self.memory(),
                            );

                            // Move the entry from the parent into the child.
                            child.push_entry(parent_entry);

                            // Move the first child of right_sibling into `child`.
                            match right_sibling.node_type() {
                                NodeType::Internal => {
                                    assert_eq!(child.node_type(), NodeType::Internal);
                                    child.push_child(right_sibling.remove_child(0));
                                }
                                NodeType::Leaf => {
                                    assert_eq!(child.node_type(), NodeType::Leaf);
                                }
                            }

                            self.save_node(right_sibling);
                            self.save_node(&mut child);
                            self.save_node(&mut node);
                            return self.remove_helper(child, key);
                        }

                        // Case 3.b: Both the left and right siblings are at their minimum sizes.

                        if let Some(left_sibling) = left_sibling {
                            // Merge child into left sibling if it exists.

                            assert!(left_sibling.at_minimum());
                            let left_sibling = self.merge(
                                child,
                                left_sibling,
                                node.remove_entry(idx - 1, self.memory()),
                            );
                            // Removing child from parent.
                            node.remove_child(idx);

                            if node.entries_len() == 0 {
                                let node_address = node.address();
                                self.deallocate_node(node);

                                if node_address == self.root_addr {
                                    // Update the root.
                                    self.root_addr = left_sibling.address();
                                    self.save_header();
                                }
                            } else {
                                self.save_node(&mut node);
                            }

                            return self.remove_helper(left_sibling, key);
                        }

                        if let Some(right_sibling) = right_sibling {
                            // Merge child into right sibling.

                            assert!(right_sibling.at_minimum());
                            let right_sibling = self.merge(
                                child,
                                right_sibling,
                                node.remove_entry(idx, self.memory()),
                            );

                            // Removing child from parent.
                            node.remove_child(idx);

                            if node.entries_len() == 0 {
                                let node_address = node.address();
                                self.deallocate_node(node);

                                if node_address == self.root_addr {
                                    // Update the root.
                                    self.root_addr = right_sibling.address();
                                    self.save_header();
                                }
                            } else {
                                self.save_node(&mut node);
                            }

                            return self.remove_helper(right_sibling, key);
                        }

                        unreachable!("At least one of the siblings must exist.");
                    }
                }
            }
        }
    }
    pub fn iter(&self) -> Iter<'_, K, V, M> {
        self.iter_internal().into()
    }
    pub fn range(&self, key_range: impl RangeBounds<K>) -> Iter<'_, K, V, M> {
        self.range_internal(key_range).into()
    }

    // TODO: In next major release (v1.0), remove this deprecated method to clean up the API.

    pub fn keys(&self) -> KeysIter<'_, K, V, M> {
        self.iter_internal().into()
    }
    pub fn keys_range(&self, key_range: impl RangeBounds<K>) -> KeysIter<'_, K, V, M> {
        self.range_internal(key_range).into()
    }
    pub fn values(&self) -> ValuesIter<'_, K, V, M> {
        self.iter_internal().into()
    }

    const fn iter_internal(&self) -> IterInternal<'_, K, V, M> {
        IterInternal::new(self)
    }

    fn range_internal(&self, key_range: impl RangeBounds<K>) -> IterInternal<'_, K, V, M> {
        if self.root_addr == NULL {
            // Map is empty.
            return IterInternal::null(self);
        }

        let range = (
            key_range.start_bound().cloned(),
            key_range.end_bound().cloned(),
        );

        IterInternal::new_in_range(self, range)
    }
    fn merge(&mut self, source: Node<K>, mut into: Node<K>, median: Entry<K>) -> Node<K> {
        into.merge(source, median, &mut self.allocator);
        into
    }
    fn allocate_node(&mut self, node_type: NodeType) -> Node<K> {
        match self.version {
            Version::V2(page_size) => Node::new_v2(self.allocator.allocate(), node_type, page_size),
        }
    }
    #[inline]
    fn deallocate_node(&mut self, node: Node<K>) {
        node.deallocate(self.allocator_mut());
    }
    #[inline]
    fn load_node(&self, address: Address) -> Node<K> {
        Node::load(address, self.version.page_size(), self.memory())
    }
    #[inline]
    fn save_node(&mut self, node: &mut Node<K>) {
        node.save(self.allocator_mut());
    }
    fn update_value(&mut self, node: &mut Node<K>, idx: usize, new_value: Vec<u8>) -> Vec<u8> {
        let old_value = node.swap_value(idx, new_value, self.memory());
        self.save_node(node);
        old_value
    }
    fn save_header(&self) {
        let header = BTreeHeader {
            version: self.version,
            root_addr: self.root_addr,
            length: self.length,
        };

        Self::write_header(&header, self.memory());
    }
    fn write_header(header: &BTreeHeader, memory: &M) {
        // Serialize the header
        let mut buf = [0; PACKED_HEADER_SIZE];
        buf[0..3].copy_from_slice(MAGIC.as_slice());
        match header.version {
            Version::V2(PageSize::Value(page_size)) => {
                buf[3] = LAYOUT_VERSION_2;
                buf[4..8].copy_from_slice(&page_size.to_le_bytes());
                buf[8..12].copy_from_slice(&PAGE_SIZE_VALUE_MARKER.to_le_bytes());
            }
        }
        buf[12..20].copy_from_slice(&header.root_addr.get().to_le_bytes());
        buf[20..28].copy_from_slice(&header.length.to_le_bytes());
        // Write the header
        crate::engine::write(memory, 0, &buf);
    }
}