From 3957b8f261f543838cd60c327f77b9c0153c4a0a Mon Sep 17 00:00:00 2001 From: Hamish Peebles Date: Fri, 11 Sep 2026 09:25:30 +0100 Subject: [PATCH] feat(btreemap): allow choosing the page size of a new map A map whose keys or values are unbounded stores its nodes in 1024-byte pages. When entries are typically small that wastes most of each page: a full node of ~20-byte entries (key and value combined) needs ~420 bytes, and most nodes aren't full. Inserting 100k such entries uses 13.4 MiB with the default page size, against 5.2 MiB (-61%) with 384-byte pages, for ~0.2% more stable memory reads per `get`. Add `BTreeMap::new_with_page_size` and `BTreeMap::init_with_page_size` so that users can pick a page size to fit their entries. Nodes that outgrow their page already continue into overflow pages, so entries of any size remain supported at any page size. No change to the memory layout is needed: the V2 header already stores an arbitrary page size, which `load` reads back. The page size is fixed once a map is created, so `init_with_page_size` ignores its argument when loading an existing map. Page sizes below the existing 128-byte minimum are rejected up front with a clear message. New proptest variants cover random operations and memory leaks with 128-byte pages, where nodes routinely spill into overflow pages. Co-Authored-By: Claude Opus 5 --- src/btreemap.rs | 135 +++++++++++++++++++++++++++++++------- src/btreemap/node.rs | 1 + src/btreemap/node/v2.rs | 2 +- src/btreemap/proptests.rs | 34 +++++++++- src/btreemap/tests.rs | 51 ++++++++++++++ 5 files changed, 194 insertions(+), 29 deletions(-) diff --git a/src/btreemap.rs b/src/btreemap.rs index afdf6547..16c64d17 100644 --- a/src/btreemap.rs +++ b/src/btreemap.rs @@ -60,7 +60,7 @@ use crate::{ }; use allocator::Allocator; pub use iter::Iter; -use node::{DerivedPageSize, Entry, Node, NodeType, PageSize, Version}; +use node::{DerivedPageSize, Entry, Node, NodeType, PageSize, Version, MINIMUM_PAGE_SIZE}; use node_cache::NodeCache; pub use node_cache::NodeCacheMetrics; use std::borrow::Cow; @@ -236,6 +236,9 @@ const DEFAULT_NODE_CACHE_NUM_SLOTS: usize = 16; /// - Use when your type's serialized size can vary or has no fixed maximum /// - Recommended for most custom types, especially those containing Strings or Vecs /// - Example: `const BOUND: Bound = Bound::Unbounded;` +/// - The map stores its nodes in 1024-byte pages by default. If your entries are +/// typically much smaller, pick a smaller page size with +/// [`init_with_page_size`](BTreeMap::init_with_page_size) to save memory. /// /// - **Bounded (`Bound::Bounded{ max_size, is_fixed_size }`)**: /// - Use when you know the maximum serialized size of your type @@ -295,21 +298,54 @@ where /// If the memory provided already contains a `BTreeMap`, then that /// map is loaded. Otherwise, a new `BTreeMap` instance is created. pub fn init(memory: M) -> Self { + if Self::contains_map(&memory) { + BTreeMap::load(memory) + } else { + BTreeMap::new(memory) + } + } + + /// Initializes a `BTreeMap` that stores its nodes in pages of `page_size` bytes. + /// + /// If the memory provided already contains a `BTreeMap`, then that map is + /// loaded and `page_size` is ignored, as a map's page size is fixed when + /// it's created. Otherwise, a new `BTreeMap` is created with the given + /// page size. See [`new_with_page_size`](Self::new_with_page_size) for + /// guidance on choosing a page size. + /// + /// # Panics + /// + /// Panics if a new map is created and `page_size` is less than 128 bytes. + /// + /// # Examples + /// + /// ```rust + /// use ic_stable_structures::{BTreeMap, DefaultMemoryImpl}; + /// + /// // Entries are unbounded but typically ~20 bytes (key and value + /// // combined), so a page much smaller than the default 1024 bytes fits + /// // a typical node. + /// let map: BTreeMap = + /// BTreeMap::init_with_page_size(DefaultMemoryImpl::default(), 384); + /// ``` + pub fn init_with_page_size(memory: M, page_size: u32) -> Self { + if Self::contains_map(&memory) { + BTreeMap::load(memory) + } else { + BTreeMap::new_with_page_size(memory, page_size) + } + } + + // Returns true if the memory already contains a `BTreeMap`. + fn contains_map(memory: &M) -> bool { if memory.size() == 0 { - // Memory is empty. Create a new map. - return BTreeMap::new(memory); + return false; } // Check if the magic in the memory corresponds to a BTreeMap. - let mut dst = vec![0; 3]; + let mut dst = [0; 3]; memory.read(0, &mut dst); - if dst != MAGIC { - // No BTreeMap found. Create a new instance. - BTreeMap::new(memory) - } else { - // The memory already contains a BTreeMap. Load it. - BTreeMap::load(memory) - } + dst == *MAGIC } /// Configures the number of node-cache slots during construction. @@ -438,21 +474,11 @@ where /// This is exposed only in testing. #[cfg(test)] pub fn init_v1(memory: M) -> Self { - if memory.size() == 0 { - // Memory is empty. Create a new map. - return BTreeMap::new_v1(memory); - } - - // Check if the magic in the memory corresponds to a BTreeMap. - let mut dst = vec![0; 3]; - memory.read(0, &mut dst); - if dst != MAGIC { - // No BTreeMap found. Create a new instance. - BTreeMap::new_v1(memory) - } else { - // The memory already contains a BTreeMap. Load it, making sure - // we don't migrate the BTreeMap to v2. + if Self::contains_map(&memory) { + // Load the map, making sure we don't migrate the BTreeMap to v2. BTreeMap::load_helper(memory, false) + } else { + BTreeMap::new_v1(memory) } } @@ -493,6 +519,65 @@ where _ => PageSize::Value(DEFAULT_PAGE_SIZE), }; + Self::new_helper(memory, page_size) + } + + /// Creates a new `BTreeMap` that stores its nodes in pages of `page_size` bytes. + /// + /// Each node of the tree is allocated a page, and a node that doesn't fit + /// in its page continues into overflow pages, so keys and values of any + /// size can be stored regardless of the page size. The page size trades + /// off memory usage against performance: + /// + /// * A page that's too large wastes memory, as every node occupies a full + /// page no matter how few bytes it uses. + /// * A page that's too small makes nodes spill into overflow pages, each + /// of which adds a few bytes of overhead and extra reads and writes + /// whenever the node is loaded or saved. + /// + /// A good page size is one that fits a typical node. A node holds at most + /// 11 entries, and a 4-byte length is stored alongside each unbounded key + /// and value, so a full leaf node takes roughly + /// `15 + 11 * (4 + key_size) + 11 * (4 + value_size)` bytes, and an + /// internal node another 96 bytes for the addresses of its children. + /// + /// [`new`](Self::new) picks a page size automatically: if both keys and + /// values are bounded, it's derived from their maximum sizes, and + /// otherwise it's 1024 bytes. That can waste most of each page when + /// entries are unbounded but typically small, which is when a smaller + /// page size pays off. + /// + /// The page size is stored in memory, and can't be changed once the map + /// is created. + /// + /// # Panics + /// + /// Panics if `page_size` is less than 128 bytes. + /// + /// # Examples + /// + /// ```rust + /// use ic_stable_structures::{BTreeMap, DefaultMemoryImpl}; + /// + /// // Entries of ~20 bytes (key and value combined) need ~420 bytes per + /// // full internal node, and nodes are rarely full, so 384-byte pages + /// // fit nearly every node. + /// let mut map: BTreeMap, Vec, _> = + /// BTreeMap::new_with_page_size(DefaultMemoryImpl::default(), 384); + /// + /// // Entries larger than a page are still supported. + /// map.insert(vec![1], vec![0; 10_000]); + /// assert_eq!(map.get(&vec![1]), Some(vec![0; 10_000])); + /// ``` + pub fn new_with_page_size(memory: M, page_size: u32) -> Self { + assert!( + page_size >= MINIMUM_PAGE_SIZE, + "page_size must be at least {MINIMUM_PAGE_SIZE} bytes, got {page_size}", + ); + Self::new_helper(memory, PageSize::Value(page_size)) + } + + fn new_helper(memory: M, page_size: PageSize) -> Self { let btree = Self { root_addr: NULL, allocator: Allocator::new( diff --git a/src/btreemap/node.rs b/src/btreemap/node.rs index 2efb5cd6..8b704240 100644 --- a/src/btreemap/node.rs +++ b/src/btreemap/node.rs @@ -15,6 +15,7 @@ mod v1; mod v2; use io::NodeReader; +pub(crate) use v2::MINIMUM_PAGE_SIZE; // The minimum degree to use in the btree. // This constant is taken from Rust's std implementation of BTreeMap. diff --git a/src/btreemap/node/v2.rs b/src/btreemap/node/v2.rs index 2037c0a9..d5e402c9 100644 --- a/src/btreemap/node/v2.rs +++ b/src/btreemap/node/v2.rs @@ -87,7 +87,7 @@ pub(super) const PAGE_OVERFLOW_DATA_OFFSET: Bytes = Bytes::new(11); // The minimum size a page can have. // Rationale: a page size needs to at least store the header (15 bytes) + all the children // addresses (88 bytes). We round that up to 128 to get a nice binary number. -const MINIMUM_PAGE_SIZE: u32 = 128; +pub(crate) const MINIMUM_PAGE_SIZE: u32 = 128; impl Node { /// Creates a new v2 node at the given address. diff --git a/src/btreemap/proptests.rs b/src/btreemap/proptests.rs index 8e193823..c0505d4a 100644 --- a/src/btreemap/proptests.rs +++ b/src/btreemap/proptests.rs @@ -1,7 +1,7 @@ use crate::{ btreemap::{ tests::{b, make_memory, run_btree_test}, - BTreeMap, + BTreeMap, MINIMUM_PAGE_SIZE, }, storable::Blob, Memory, @@ -69,6 +69,21 @@ fn comprehensive(#[strategy(pvec(operation_strategy(), 100..5_000))] ops: Vec, +) { + let mem = make_memory(); + let mut btree = BTreeMap::new_with_page_size(mem, MINIMUM_PAGE_SIZE); + let mut std_btree = StdBTreeMap::new(); + + for op in ops.into_iter() { + execute_operation(&mut std_btree, &mut btree, op); + } +} + // Same as `comprehensive` but with the node cache enabled. #[proptest(cases = 10)] fn comprehensive_cached(#[strategy(pvec(operation_strategy(), 100..5_000))] ops: Vec) { @@ -239,9 +254,22 @@ fn iter_count_test(#[strategy(0..250u8)] start: u8, #[strategy(#start..255u8)] e #[proptest] fn no_memory_leaks(#[strategy(pvec(pvec(0..u8::MAX, 100..10_000), 100))] keys: Vec>) { - let mem = make_memory(); - let mut btree = BTreeMap::new(mem); + run_no_memory_leaks(keys, BTreeMap::new(make_memory())); +} + +// Same as `no_memory_leaks` but with the smallest page size allowed, where +// every node spans many overflow pages. +#[proptest(cases = 32)] +fn no_memory_leaks_min_page_size( + #[strategy(pvec(pvec(0..u8::MAX, 100..10_000), 100))] keys: Vec>, +) { + run_no_memory_leaks( + keys, + BTreeMap::new_with_page_size(make_memory(), MINIMUM_PAGE_SIZE), + ); +} +fn run_no_memory_leaks(keys: Vec>, mut btree: BTreeMap, (), M>) { // Insert entries. for k in keys.iter() { btree.insert(k.clone(), ()); diff --git a/src/btreemap/tests.rs b/src/btreemap/tests.rs index 0d7216fc..ed81290a 100644 --- a/src/btreemap/tests.rs +++ b/src/btreemap/tests.rs @@ -1586,6 +1586,57 @@ fn accepts_small_or_equal_value_sizes() { let _btree: BTreeMap, Blob<3>, _> = BTreeMap::init(btree.into_memory()); } +#[test] +fn new_with_page_size_persists_the_page_size() { + let mut btree: BTreeMap, Vec, _> = BTreeMap::new_with_page_size(make_memory(), 200); + btree.insert(vec![1], vec![2]); + + let btree: BTreeMap, Vec, _> = BTreeMap::load(btree.into_memory()); + assert_eq!(btree.version, Version::V2(PageSize::Value(200))); + assert_eq!(btree.get(&vec![1]), Some(vec![2])); +} + +#[test] +fn new_with_page_size_overrides_the_page_size_of_bounded_types() { + let btree: BTreeMap = BTreeMap::new_with_page_size(make_memory(), 4096); + assert_eq!(btree.version, Version::V2(PageSize::Value(4096))); +} + +#[test] +#[should_panic(expected = "page_size must be at least 128 bytes, got 127")] +fn new_with_page_size_rejects_pages_below_the_minimum() { + let _btree: BTreeMap, Vec, _> = BTreeMap::new_with_page_size(make_memory(), 127); +} + +#[test] +fn init_with_page_size_creates_a_map_with_the_page_size() { + let btree: BTreeMap, Vec, _> = BTreeMap::init_with_page_size(make_memory(), 256); + assert_eq!(btree.version, Version::V2(PageSize::Value(256))); +} + +#[test] +fn init_with_page_size_keeps_the_page_size_of_an_existing_map() { + let mut btree: BTreeMap, Vec, _> = + BTreeMap::init_with_page_size(make_memory(), 256); + btree.insert(vec![1], vec![2]); + + // The page size is fixed when the map is created, so a different page + // size passed when reinitializing is ignored. + let btree: BTreeMap, Vec, _> = + BTreeMap::init_with_page_size(btree.into_memory(), 1024); + assert_eq!(btree.version, Version::V2(PageSize::Value(256))); + assert_eq!(btree.get(&vec![1]), Some(vec![2])); + + // Likewise for a map created with the default page size. + let btree: BTreeMap, Vec, _> = BTreeMap::new(make_memory()); + let btree: BTreeMap, Vec, _> = + BTreeMap::init_with_page_size(btree.into_memory(), 256); + assert_eq!( + btree.version, + Version::V2(PageSize::Value(DEFAULT_PAGE_SIZE)) + ); +} + fn bruteforce_range_search() { let (key, value) = (K::build, V::build); run_btree_test(|mut stable_map| {