diff --git a/benchmarks/btreemap/canbench_results.yml b/benchmarks/btreemap/canbench_results.yml index 108c18e9..522b8b00 100644 --- a/benchmarks/btreemap/canbench_results.yml +++ b/benchmarks/btreemap/canbench_results.yml @@ -1,4 +1,18 @@ benches: + btreemap_get_and_incr: + total: + calls: 1 + instructions: 376179070 + heap_increase: 0 + stable_memory_increase: 0 + scopes: {} + btreemap_get_and_incr_via_entry: + total: + calls: 1 + instructions: 317352818 + heap_increase: 0 + stable_memory_increase: 0 + scopes: {} btreemap_v2_contains_10mib_values: total: calls: 1 diff --git a/benchmarks/btreemap/src/main.rs b/benchmarks/btreemap/src/main.rs index ff766c85..52c0572a 100644 --- a/benchmarks/btreemap/src/main.rs +++ b/benchmarks/btreemap/src/main.rs @@ -1,6 +1,7 @@ use benchmarks::{random::Random, vec::UnboundedVecN}; use canbench_rs::{bench, bench_fn, BenchResult}; use candid::Principal; +use ic_stable_structures::btreemap::entry::Entry::Occupied; use ic_stable_structures::memory_manager::{MemoryId, MemoryManager}; use ic_stable_structures::{storable::Blob, BTreeMap, DefaultMemoryImpl, Memory, Storable}; use std::ops::Bound; @@ -441,6 +442,45 @@ pub fn btreemap_v2_get_10mib_values() -> BenchResult { }) } +#[bench(raw)] +pub fn btreemap_get_and_incr() -> BenchResult { + let count = 10000; + let mut btree = BTreeMap::new(DefaultMemoryImpl::default()); + let mut rng = Rng::from_seed(0); + let values = generate_random_kv::(count, &mut rng); + for (key, value) in values.iter().copied() { + btree.insert(key, value); + } + + bench_fn(|| { + for (key, _) in values { + let existing = btree.get(&key).unwrap(); + btree.insert(key, existing + 1); + } + }) +} + +#[bench(raw)] +pub fn btreemap_get_and_incr_via_entry() -> BenchResult { + let count = 10000; + let mut btree = BTreeMap::new(DefaultMemoryImpl::default()); + let mut rng = Rng::from_seed(0); + let values = generate_random_kv::(count, &mut rng); + for (key, value) in values.iter().copied() { + btree.insert(key, value); + } + + bench_fn(|| { + for (key, _) in values { + let Occupied(e) = btree.entry(key) else { + panic!() + }; + let existing = e.get(); + e.insert(existing + 1); + } + }) +} + // Benchmarks for `BTreeMap::contains_key`. // Reduced grid: contains_key traversal is identical to get, only skips value deserialization. bench_tests! { diff --git a/src/btreemap.rs b/src/btreemap.rs index afdf6547..1d2977e7 100644 --- a/src/btreemap.rs +++ b/src/btreemap.rs @@ -49,6 +49,7 @@ //! ---------------------------------------- //! ``` mod allocator; +pub mod entry; mod iter; mod node; mod node_cache; @@ -60,7 +61,7 @@ use crate::{ }; use allocator::Allocator; pub use iter::Iter; -use node::{DerivedPageSize, Entry, Node, NodeType, PageSize, Version}; +use node::{DerivedPageSize, Node, NodeType, PageSize, Version}; use node_cache::NodeCache; pub use node_cache::NodeCacheMetrics; use std::borrow::Cow; @@ -676,37 +677,254 @@ where ))); } - // 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); + self.split_root_if_full(root) + }; - // The new root has the old root as its only child. - new_root.push_child(self.root_addr); + self.insert_nonfull(root, key, value, 0) + .map(Cow::Owned) + .map(V::from_bytes) + } - // Update the root address. - self.root_addr = new_root.address(); - self.save_header(); + /// Gets an [`Entry`](entry::Entry) for the given `key`, which gives efficient in-place access + /// to a map's entries, allowing inspection or modification without redundant key lookups. The + /// API mirrors [`std::collections::btree_map::Entry`] as closely as the stable-memory model + /// allows. + /// + /// Returns [`Entry::Occupied`](entry::Entry::Occupied) when `key` is already present, or + /// [`Entry::Vacant`](entry::Entry::Vacant) when it is not. Both variants expose methods to + /// read, overwrite, or remove the value. + /// + /// # Differences from `std::collections::BTreeMap::entry` + /// + /// The standard library's `or_insert` returns `&mut V`, giving a direct reference into the + /// map. Because values live in stable memory, long-lived references are not possible here. + /// Instead, `or_insert` (and its variants) return an [`OccupiedEntry`](entry::OccupiedEntry), + /// which lets you continue operating on the entry without a second key lookup. + /// + /// # This method writes to stable memory + /// + /// Unlike [`get`](Self::get) and [`contains_key`](Self::contains_key), `entry` **writes**. + /// So that a later insert can complete in a single pass, it splits every full node on the + /// path down from the root, which allocates nodes and rewrites the header — even when + /// `key` turns out to be present, and even if the returned entry is dropped without + /// inserting anything. + /// + /// The cost is bounded and self-limiting: each node splits at most once, and that split + /// would have happened on the next insert through the same path regardless. Repeating the + /// same probes allocates nothing further. Still, prefer [`get`](Self::get) or + /// [`contains_key`](Self::contains_key) when you only mean to read. + /// + /// # Examples + /// + /// ```rust + /// use ic_stable_structures::{BTreeMap, DefaultMemoryImpl}; + /// + /// let mut map: BTreeMap = BTreeMap::new(DefaultMemoryImpl::default()); + /// + /// // Insert a default value when the key is absent, then read it back. + /// let val = map.entry(1).or_insert(0).get(); + /// assert_eq!(val, 0); + /// + /// // Increment a counter, seeding it to 1 when first seen. + /// for key in [1, 1, 2, 3, 1] { + /// map.entry(key).and_modify(|v| *v += 1).or_insert(1); + /// } + /// assert_eq!(map.get(&1), Some(3)); + /// assert_eq!(map.get(&2), Some(1)); + /// assert_eq!(map.get(&3), Some(1)); + /// ``` + /// + /// # See also + /// + /// - [`entry::Entry`] — the type returned by this method + /// - [`entry::OccupiedEntry`] — methods available when the key is present + /// - [`entry::VacantEntry`] — methods available when the key is absent + #[must_use = "`entry` writes to stable memory even when its result is unused; \ + for a read-only lookup use `get` or `contains_key`"] + pub fn entry(&mut self, key: K) -> entry::Entry { + // For an empty map the key is trivially absent. Avoid calling + // `find_node_for_insert` here because that eagerly allocates a root + // node and persists `root_addr` to the header, which would leave the + // map in an inconsistent state if the returned `VacantEntry` is dropped + // without calling `insert`. + if self.root_addr == NULL { + return entry::Entry::Vacant(entry::VacantEntry { + map: self, + key, + slot: None, + }); + } + + // Take the root from the node cache, or load it from memory. The node is handed + // to the returned entry, which reads and writes it without another load. It is + // not put back into the cache when the entry goes away: mutating an entry saves + // the node, and `save_node` invalidates the cache slot precisely because a + // just-saved node holds all of its keys and values materialized. In *node cache* + // terms an entry dropped without mutating therefore costs at most one cold slot, + // refilled by the next ordinary lookup. That is not the whole cost of a dropped + // entry though — see the "This method writes to stable memory" section above. + let mut root = self.take_or_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 + return entry::Entry::Occupied(entry::OccupiedEntry { + map: self, + key, + slot: entry::NodeSlot::new(root, idx, 0), + }); + } - // Split the old (full) root. - self.split_child(&mut new_root, 0, None); + root = self.split_root_if_full(root); - new_root - } else { - root + let (key, slot, key_exists) = + self.find_node_for_insert(root, key, 0, |_, node, idx, depth, key, key_exists| { + (key, entry::NodeSlot::new(node, idx, depth), key_exists) + }); + + if key_exists { + entry::Entry::Occupied(entry::OccupiedEntry { + map: self, + key, + slot, + }) + } else { + entry::Entry::Vacant(entry::VacantEntry { + map: self, + key, + slot: Some(slot), + }) + } + } + + /// Ensures the root node is not full, called before an insertion. + /// + /// If `root` is full, a new internal node is allocated, made the parent of the old + /// root, and the old root is split into two children. The new (non-full) root is + /// returned. If `root` is not full it is returned unchanged. + fn split_root_if_full(&mut self, root: Node) -> Node { + if !root.is_full() { + return root; + } + + // Allocate a new node that will become 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); + + // Persist the updated root address before splitting. + self.root_addr = new_root.address(); + self.save_header(); + + // Split the old (full) root into two children of new_root. + self.split_child(&mut new_root, 0, None); + + new_root + } + + /// Core B-tree insertion traversal used by [`BTreeMap::entry`]. + /// + /// Descends from `node` (which must not be full) to the leaf or internal node where + /// `key` either already exists or should be inserted, then invokes `callback` with: + /// + /// * `map` — mutable access to the map (for saving nodes, updating length, etc.) + /// * `node` — the target node + /// * `idx` — the relevant slot index within `node` + /// * `depth` — the depth of `node` (distance from the root) + /// * `key` — the key being inserted + /// * `key_exists` — `true` if `key` is already present at `idx`; `false` if `idx` is + /// the position where `key` should be inserted + /// + /// The callback's return value is propagated back to the caller. + /// + /// Unmodified nodes along the traversal path are returned to the node cache. The + /// target node is not: ownership passes to `callback`, which is responsible for it. + /// + /// NOTE: this is the same descent as [`BTreeMap::insert_nonfull`], which stops at the + /// target node instead of inserting into it. The two were deliberately left separate: + /// expressing `insert` in terms of this callback form benchmarked slower. Keep the + /// split and descent logic in the two in sync. + /// + /// PRECONDITION: `node` is not full. + fn find_node_for_insert( + &mut self, + mut node: Node, + key: K, + depth: u8, + callback: impl FnOnce(&mut Self, Node, usize, u8, K, bool) -> R, + ) -> R { + // 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. + callback(self, node, idx, depth, key, true) } - }; + Err(idx) => { + // The key isn't in the node. `idx` is where that key should be inserted. - self.insert_nonfull(root, key, value, 0) - .map(Cow::Owned) - .map(V::from_bytes) + match node.node_type() { + NodeType::Leaf => { + // The node is a non-full leaf. + callback(self, node, idx, depth, key, false) + } + NodeType::Internal => { + // The node is an internal node. + // Load the child that we should add the entry to. + let mut child = self.take_or_load_node(node.child(idx)); + let child_depth = depth.saturating_add(1); + + if child.is_full() { + // Check if the key already exists in the child. + if let Ok(idx) = child.search(&key, self.memory()) { + // Key found. The parent node is unmodified — return it to cache. + self.return_node(node, depth); + return callback(self, child, idx, child_depth, key, true); + } + + // The child is full. Split the child. + // Pass the already-loaded child to avoid a redundant load. + self.split_child(&mut node, idx, Some(child)); + + // The children have now changed. Search again for + // the child where we need to store the entry in. + // The median promoted by the split came out of `child`, which + // was just shown not to contain `key`, so the search cannot + // find it here; `Ok` would mean descending past a matching key + // and inserting a duplicate. + let search_result = node.search(&key, self.memory()); + debug_assert!( + search_result.is_err(), + "promoted median cannot equal the key" + ); + let idx = search_result.unwrap_or_else(|idx| idx); + child = self.load_node(node.child(idx)); + } else { + // Happy path: child is not full. The current node + // will not be modified — return it to cache. + self.return_node(node, depth); + } + + // The child should now be not full. + assert!(!child.is_full()); + + self.find_node_for_insert(child, key, child_depth, callback) + } + } + } + } } /// Inserts an entry into a node that is *not full*. + /// + /// NOTE: shares its descent and node-splitting logic with + /// [`BTreeMap::find_node_for_insert`]; keep the two in sync. fn insert_nonfull( &mut self, mut node: Node, @@ -728,16 +946,11 @@ where match node.node_type() { NodeType::Leaf => { - // The node is a non-full leaf. - // Insert the entry at the proper location. + // The node is a non-full leaf. Insert directly. 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 => { @@ -905,14 +1118,14 @@ where } #[inline(always)] - fn first_entry_inner(&self, node: &Node, depth: u8) -> Entry { + fn first_entry_inner(&self, node: &Node, depth: u8) -> node::Entry { self.find_first_or_last(node, true, depth, |n, i, m| { n.get_key_read_value_uncached(i, m) }) } #[inline(always)] - fn last_entry_inner(&self, node: &Node, depth: u8) -> Entry { + fn last_entry_inner(&self, node: &Node, depth: u8) -> node::Entry { self.find_first_or_last(node, false, depth, |n, i, m| { n.get_key_read_value_uncached(i, m) }) @@ -1041,10 +1254,10 @@ where if node.entries_len() == 0 { assert_eq!( - node.address(), self.root_addr, + 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; @@ -1401,7 +1614,7 @@ where /// Removes and returns the rightmost (maximum) entry in the subtree rooted /// at `node`, in a single top-down pass. This avoids the double traversal /// of the previous approach (get_max + remove_helper). - fn remove_rightmost(&mut self, mut node: Node, depth: u8) -> Option> { + fn remove_rightmost(&mut self, mut node: Node, depth: u8) -> Option> { match node.node_type() { NodeType::Leaf => { let entry = node.pop_entry(self.memory())?; @@ -1476,7 +1689,7 @@ where /// Removes and returns the leftmost (minimum) entry in the subtree rooted /// at `node`, in a single top-down pass. - fn remove_leftmost(&mut self, mut node: Node, depth: u8) -> Option> { + fn remove_leftmost(&mut self, mut node: Node, depth: u8) -> Option> { match node.node_type() { NodeType::Leaf => { if node.entries_len() == 0 { @@ -1667,7 +1880,7 @@ where /// Output: /// [1, 2, 3, 4, 5, 6, 7] (stored in the `into` node) /// `source` is deallocated. - fn merge(&mut self, source: Node, mut into: Node, median: Entry) -> Node { + fn merge(&mut self, source: Node, mut into: Node, median: node::Entry) -> Node { let source_addr = source.address(); into.merge(source, median, &mut self.allocator); // Node::merge saves `into` and deallocates `source` directly through diff --git a/src/btreemap/entry.rs b/src/btreemap/entry.rs new file mode 100644 index 00000000..e20afd5b --- /dev/null +++ b/src/btreemap/entry.rs @@ -0,0 +1,783 @@ +//! Entry API for [`BTreeMap`]. +//! +//! This module provides the [`Entry`] type, which gives efficient in-place access to a map's +//! entries, allowing inspection or modification without redundant key lookups. The API mirrors +//! [`std::collections::btree_map::Entry`] as closely as the stable-memory model allows. +//! +//! # Note on `or_insert` return type +//! +//! The standard library's `or_insert` returns `&mut V`, giving a direct reference into the +//! map. Because values in this [`BTreeMap`] live in stable memory, long-lived references are +//! not possible. Instead, `or_insert` (and its variants) return an [`OccupiedEntry`], which +//! lets you continue reading or modifying the entry without a second key lookup. +//! +//! For the same reason there is no equivalent of the standard library's `get_mut` or +//! `into_mut`. Use [`OccupiedEntry::and_modify`], which reads the value, hands it to your +//! closure, and writes it back. +//! +//! # `entry` writes to stable memory +//! +//! [`BTreeMap::entry`] splits full nodes on the way down so that a later insert can finish +//! in one pass. It therefore writes — allocating nodes and rewriting the header — even for +//! a key that turns out to be present, and even if the entry is dropped unused. The cost is +//! bounded (each node splits at most once, and that split was coming on the next insert +//! anyway), but prefer [`BTreeMap::get`] when you only mean to read. See +//! [`BTreeMap::entry`] for details. +//! +//! # Examples +//! +//! ```rust +//! use ic_stable_structures::{BTreeMap, DefaultMemoryImpl}; +//! +//! let mut map: BTreeMap = BTreeMap::new(DefaultMemoryImpl::default()); +//! +//! // Insert a value only when the key is absent. +//! map.entry(1).or_insert(42); +//! assert_eq!(map.get(&1), Some(42)); +//! +//! // Increment a counter, seeding it to 1 if absent. +//! map.entry(1).and_modify(|v| *v += 1).or_insert(1); +//! assert_eq!(map.get(&1), Some(43)); +//! ``` + +use crate::btreemap::node::{Node, NodeType}; +use crate::{BTreeMap, Memory, Storable}; +use std::borrow::Cow; +use std::marker::PhantomData; + +/// A loaded node together with the index of one slot inside it: for an +/// [`OccupiedEntry`] the slot holding the entry, for a [`VacantEntry`] the slot the key +/// would occupy. `depth` is the node's distance from the root, which the node cache's +/// eviction policy needs. +/// +/// The node itself is kept, rather than just its address, so that reading and writing +/// through an entry needs no further node loads. [`BTreeMap::entry`] takes the node out +/// of the map's node cache; it is put back only where that is both useful and cheap (see +/// [`OccupiedEntry::remove`]). Mutating an entry saves the node, which invalidates its +/// cache slot — a just-saved node holds every one of its keys and values materialized, +/// so returning it to the cache in that state would waste heap. +pub(crate) struct NodeSlot { + pub(crate) node: Node, + pub(crate) idx: usize, + pub(crate) depth: u8, +} + +impl NodeSlot { + pub(crate) fn new(node: Node, idx: usize, depth: u8) -> Self { + Self { node, idx, depth } + } +} + +/// A view into a single entry of a [`BTreeMap`], which may either be occupied or vacant. +/// +/// This type is returned by [`BTreeMap::entry`]. +pub enum Entry<'a, K: 'a + Storable + Ord + Clone, V: 'a + Storable, M: Memory> { + /// A vacant entry: the key is not present in the map. + Vacant(VacantEntry<'a, K, V, M>), + /// An occupied entry: the key is already present in the map. + Occupied(OccupiedEntry<'a, K, V, M>), +} + +/// A view into a vacant entry in a [`BTreeMap`]. +/// +/// Obtained from [`Entry::Vacant`]. +pub struct VacantEntry<'a, K: 'a + Storable + Ord + Clone, V: 'a + Storable, M: Memory> { + pub(crate) map: &'a mut BTreeMap, + pub(crate) key: K, + /// Pre-computed insertion point from [`BTreeMap::entry`]. + /// + /// `None` when the map was empty at the time `entry` was called — the root + /// node had not yet been allocated, so we defer the full insert to + /// [`VacantEntry::insert`] to avoid corrupting the map if this entry is + /// dropped without inserting. + pub(crate) slot: Option>, +} + +/// A view into an occupied entry in a [`BTreeMap`]. +/// +/// Obtained from [`Entry::Occupied`] or as the result of [`VacantEntry::insert`]. +pub struct OccupiedEntry<'a, K: 'a + Storable + Ord + Clone, V: 'a + Storable, M: Memory> { + pub(crate) map: &'a mut BTreeMap, + pub(crate) key: K, + pub(crate) slot: NodeSlot, +} + +/// A value returned by [`OccupiedEntry::insert`] or [`OccupiedEntry::remove`] that has not +/// yet been deserialized. +/// +/// Deserialization is deferred so that callers who do not need the previous value pay no +/// decode cost. Call [`into_value`](LazyValue::into_value) to obtain the concrete `T`. +/// +/// # Examples +/// +/// ```rust +/// use ic_stable_structures::{BTreeMap, DefaultMemoryImpl, btreemap::entry::Entry}; +/// +/// let mut map: BTreeMap = BTreeMap::new(DefaultMemoryImpl::default()); +/// map.insert(1, 10); +/// +/// if let Entry::Occupied(e) = map.entry(1) { +/// // Discard the old value without deserializing it. +/// let _old: _ = e.insert(99); +/// } +/// +/// if let Entry::Occupied(e) = map.entry(1) { +/// // Deserialize only when the value is actually needed. +/// let old_value = e.insert(0).into_value(); +/// assert_eq!(old_value, 99); +/// } +/// ``` +pub struct LazyValue { + bytes: Vec, + phantom_data: PhantomData, +} + +impl<'a, K: 'a + Storable + Ord + Clone, V: 'a + Storable, M: Memory> Entry<'a, K, V, M> { + /// Returns a reference to this entry's key. + pub fn key(&self) -> &K { + match self { + Entry::Occupied(entry) => entry.key(), + Entry::Vacant(entry) => entry.key(), + } + } + + /// Consumes the entry and returns its key. + pub fn into_key(self) -> K { + match self { + Entry::Occupied(entry) => entry.into_key(), + Entry::Vacant(entry) => entry.into_key(), + } + } + + /// Ensures a value is present by inserting `default` if the entry is vacant, then returns + /// an [`OccupiedEntry`] for further operations. + /// + /// # Examples + /// + /// ```rust + /// use ic_stable_structures::{BTreeMap, DefaultMemoryImpl}; + /// + /// let mut map: BTreeMap = BTreeMap::new(DefaultMemoryImpl::default()); + /// assert_eq!(map.entry(1).or_insert(10).get(), 10); + /// assert_eq!(map.entry(1).or_insert(99).get(), 10); // already present + /// ``` + pub fn or_insert(self, default: V) -> OccupiedEntry<'a, K, V, M> { + match self { + Entry::Occupied(entry) => entry, + Entry::Vacant(entry) => entry.insert(default), + } + } + + /// Ensures a value is present by inserting the result of `default` if the entry is vacant, + /// then returns an [`OccupiedEntry`]. + /// + /// # Examples + /// + /// ```rust + /// use ic_stable_structures::{BTreeMap, DefaultMemoryImpl}; + /// + /// let mut map: BTreeMap = BTreeMap::new(DefaultMemoryImpl::default()); + /// map.entry(1).or_insert_with(|| 42u32); + /// assert_eq!(map.get(&1), Some(42)); + /// ``` + pub fn or_insert_with(self, default: impl FnOnce() -> V) -> OccupiedEntry<'a, K, V, M> { + match self { + Entry::Occupied(entry) => entry, + Entry::Vacant(entry) => entry.insert(default()), + } + } + + /// Ensures a value is present by inserting the result of `default`, called with the + /// entry's key, if the entry is vacant. Returns an [`OccupiedEntry`]. + /// + /// # Examples + /// + /// ```rust + /// use ic_stable_structures::{BTreeMap, DefaultMemoryImpl}; + /// + /// let mut map: BTreeMap = BTreeMap::new(DefaultMemoryImpl::default()); + /// map.entry(7).or_insert_with_key(|&k| k * 2); + /// assert_eq!(map.get(&7), Some(14)); + /// ``` + pub fn or_insert_with_key(self, default: impl FnOnce(&K) -> V) -> OccupiedEntry<'a, K, V, M> { + match self { + Entry::Occupied(entry) => entry, + Entry::Vacant(entry) => { + let val = default(entry.key()); + entry.insert(val) + } + } + } + + /// Ensures a value is present by inserting `V::default()` if the entry is vacant, then + /// returns an [`OccupiedEntry`]. + /// + /// # Examples + /// + /// ```rust + /// use ic_stable_structures::{BTreeMap, DefaultMemoryImpl}; + /// + /// let mut map: BTreeMap = BTreeMap::new(DefaultMemoryImpl::default()); + /// map.entry(1).or_default(); + /// assert_eq!(map.get(&1), Some(0u32)); + /// ``` + pub fn or_default(self) -> OccupiedEntry<'a, K, V, M> + where + V: Default, + { + self.or_insert_with(V::default) + } + + /// Provides in-place mutable access to an occupied entry before any potential inserts + /// via `or_insert` and friends. + /// + /// If the entry is vacant the closure is not called and the entry is returned unchanged, + /// making it possible to chain with `or_insert` and friends. + /// + /// # Examples + /// + /// ```rust + /// use ic_stable_structures::{BTreeMap, DefaultMemoryImpl}; + /// + /// let mut map: BTreeMap = BTreeMap::new(DefaultMemoryImpl::default()); + /// map.insert(1, 10); + /// + /// // Increment existing value, or seed with 1 for a new key. + /// map.entry(1).and_modify(|v| *v += 1).or_insert(1); + /// assert_eq!(map.get(&1), Some(11)); + /// + /// map.entry(2).and_modify(|v| *v += 1).or_insert(1); + /// assert_eq!(map.get(&2), Some(1)); + /// ``` + pub fn and_modify(self, f: impl FnOnce(&mut V)) -> Self { + match self { + Entry::Occupied(entry) => Entry::Occupied(entry.and_modify(f)), + Entry::Vacant(entry) => Entry::Vacant(entry), + } + } +} + +impl<'a, K: 'a + Storable + Ord + Clone, V: 'a + Storable, M: Memory> VacantEntry<'a, K, V, M> { + /// Returns a reference to the entry's key. + pub fn key(&self) -> &K { + &self.key + } + + /// Consumes the entry and returns its key. + pub fn into_key(self) -> K { + self.key + } + + /// Inserts `value` into the map at this entry's key and returns an [`OccupiedEntry`] + /// pointing at the newly inserted value. + /// + /// # Panics + /// + /// Panics if the serialized key or value exceeds the maximum size of its type, exactly + /// as [`BTreeMap::insert`] does. Note that [`BTreeMap::entry`] may already have split + /// nodes by this point; the map is left valid and consistent, but not necessarily in + /// the shape it had before `entry` was called. + pub fn insert(self, value: V) -> OccupiedEntry<'a, K, V, M> { + let Self { map, key, slot } = self; + let slot = match slot { + Some(mut slot) => { + slot.node + .insert_entry(slot.idx, (key.clone(), value.into_bytes_checked())); + map.save_node(&mut slot.node); + map.length += 1; + map.save_header(); + slot + } + None => { + // The map was empty when `entry()` was called. Delegate to the regular + // insert path, which handles allocating the root, then point the new + // `OccupiedEntry` at the root node's first slot. + map.insert(key.clone(), value); + let node = map.take_or_load_node(map.root_addr); + NodeSlot::new(node, 0, 0) + } + }; + OccupiedEntry { map, key, slot } + } +} + +impl<'a, K: 'a + Storable + Ord + Clone, V: 'a + Storable, M: Memory> OccupiedEntry<'a, K, V, M> { + /// Returns a reference to the key stored in the map. + /// + /// As in the standard library, this is the key already held by the map, not the one + /// passed to [`BTreeMap::entry`]. The two compare equal, but they can differ in ways + /// `Ord` does not see — for example a `K` whose ordering ignores some of its fields. + pub fn key(&self) -> &K { + self.slot.node.key(self.slot.idx, self.map.memory()) + } + + /// Consumes the entry and returns the key stored in the map. + /// + /// Like [`key`](Self::key), this is the map's own key rather than the one passed to + /// [`BTreeMap::entry`]. + pub fn into_key(self) -> K { + self.key().clone() + } + + /// Returns the current value associated with this entry. + /// + /// Every call re-reads the value from stable memory and deserializes it; the result is + /// not memoized. Bind it to a local if you need it more than once. + pub fn get(&self) -> V { + // Read straight out of the node the entry already holds — no load required. + // The read is uncached so that repeated reads don't inflate the node with a + // materialized value buffer. + let value_bytes = self + .slot + .node + .read_value_uncached(self.slot.idx, self.map.memory()); + V::from_bytes(Cow::Owned(value_bytes)) + } + + /// Provides in-place mutable access to the value in this occupied entry. + /// + /// Reads the current value, calls `f` with a mutable reference to it, then writes the + /// modified value back. Returns `self` so the call can be chained. + /// + /// The write is unconditional: the value is re-serialized and the node saved even if + /// `f` left it untouched. + /// + /// # Examples + /// + /// ```rust + /// use ic_stable_structures::{BTreeMap, DefaultMemoryImpl, btreemap::entry::Entry}; + /// + /// let mut map: BTreeMap = BTreeMap::new(DefaultMemoryImpl::default()); + /// map.insert(1, 10); + /// + /// if let Entry::Occupied(e) = map.entry(1) { + /// e.and_modify(|v| *v *= 2); + /// } + /// assert_eq!(map.get(&1), Some(20)); + /// ``` + pub fn and_modify(mut self, f: impl FnOnce(&mut V)) -> Self { + let mut value = self.get(); + f(&mut value); + self.write_value(value); + self + } + + /// Overwrites the value in this entry's slot, returning the old bytes. + /// + /// `update_value` saves the node, which invalidates its cache slot. The node is + /// deliberately not put back into the cache: a just-saved node holds every one of + /// its keys and values materialized, so caching it in that state would waste heap. + fn write_value(&mut self, value: V) -> Vec { + self.map.update_value( + &mut self.slot.node, + self.slot.idx, + value.into_bytes_checked(), + ) + } + + /// Replaces the current value with `value` and returns the previous value as a + /// [`LazyValue`], which is only deserialized if you call [`LazyValue::into_value`]. + /// + /// # Panics + /// + /// Panics if the serialized `value` exceeds the maximum size of `V`, exactly as + /// [`BTreeMap::insert`] does. The map is left unmodified in that case. + /// + /// # Examples + /// + /// ```rust + /// use ic_stable_structures::{BTreeMap, DefaultMemoryImpl, btreemap::entry::Entry}; + /// + /// let mut map: BTreeMap = BTreeMap::new(DefaultMemoryImpl::default()); + /// map.insert(1, 10); + /// + /// if let Entry::Occupied(e) = map.entry(1) { + /// let old = e.insert(99).into_value(); + /// assert_eq!(old, 10); + /// } + /// assert_eq!(map.get(&1), Some(99)); + /// ``` + pub fn insert(mut self, value: V) -> LazyValue { + LazyValue::new(self.write_value(value)) + } + + /// Removes the entry from the map and returns the stored value as a [`LazyValue`], which + /// is only deserialized if you call [`LazyValue::into_value`]. + /// + /// # Examples + /// + /// ```rust + /// use ic_stable_structures::{BTreeMap, DefaultMemoryImpl, btreemap::entry::Entry}; + /// + /// let mut map: BTreeMap = BTreeMap::new(DefaultMemoryImpl::default()); + /// map.insert(1, 42); + /// + /// if let Entry::Occupied(e) = map.entry(1) { + /// assert_eq!(e.remove().into_value(), 42); + /// } + /// assert!(map.is_empty()); + /// ``` + pub fn remove(self) -> LazyValue { + let Self { map, key, slot } = self; + let NodeSlot { + mut node, + idx, + depth, + } = slot; + let bytes = match node.node_type() { + NodeType::Leaf if node.can_remove_entry_without_merging() => { + // Fast path: the leaf has enough entries to remove without merging. + let value = node.remove_entry(idx, map.memory()).1; + // `can_remove_entry_without_merging` guarantees the leaf held more than + // the minimum number of entries, so it cannot be empty after removal. + debug_assert!(node.entries_len() > 0); + map.save_node(&mut node); + map.length -= 1; + map.save_header(); + value + } + _ => { + // Slow path: the removal may require rebalancing/merging, so do a fresh + // traversal from the root. The node is unmodified, so return it to the + // cache first — the traversal can then take it from there rather than + // re-reading it from memory. + map.return_node(node, depth); + let root = map.take_or_load_node(map.root_addr); + map.remove_helper(root, &key, 0).expect("key must exist") + } + }; + LazyValue::new(bytes) + } +} + +impl LazyValue { + pub(crate) fn new(bytes: Vec) -> Self { + LazyValue { + bytes, + phantom_data: PhantomData, + } + } + + /// Deserializes and returns the value. + pub fn into_value(self) -> T { + T::from_bytes(Cow::Owned(self.bytes)) + } +} + +// The `Debug` impls below print keys only. Values are deliberately left out: reading one +// means a stable-memory read and a deserialization, which is not what anyone expects a +// `Debug` impl to do. + +impl std::fmt::Debug + for Entry<'_, K, V, M> +{ + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Entry::Vacant(entry) => f.debug_tuple("Entry::Vacant").field(entry).finish(), + Entry::Occupied(entry) => f.debug_tuple("Entry::Occupied").field(entry).finish(), + } + } +} + +impl std::fmt::Debug + for VacantEntry<'_, K, V, M> +{ + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("VacantEntry") + .field("key", self.key()) + .finish() + } +} + +impl std::fmt::Debug + for OccupiedEntry<'_, K, V, M> +{ + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("OccupiedEntry") + .field("key", self.key()) + .finish() + } +} + +impl std::fmt::Debug for LazyValue { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + // The value is intentionally not deserialized here. + f.debug_struct("LazyValue") + .field("num_bytes", &self.bytes.len()) + .finish() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::cell::RefCell; + use std::rc::Rc; + + fn new_map() -> BTreeMap>>> { + BTreeMap::new(Rc::new(RefCell::new(Vec::new()))) + } + + /// A key whose ordering ignores `tag`, so two keys can be `Ord`-equal while holding + /// different bytes. Used to tell the stored key apart from the one passed to `entry`. + #[derive(Clone, Debug)] + struct TaggedKey { + id: u32, + tag: u8, + } + + impl PartialEq for TaggedKey { + fn eq(&self, other: &Self) -> bool { + self.id == other.id + } + } + impl Eq for TaggedKey {} + impl PartialOrd for TaggedKey { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } + } + impl Ord for TaggedKey { + fn cmp(&self, other: &Self) -> std::cmp::Ordering { + self.id.cmp(&other.id) + } + } + impl Storable for TaggedKey { + fn to_bytes(&self) -> Cow<'_, [u8]> { + let mut bytes = self.id.to_be_bytes().to_vec(); + bytes.push(self.tag); + Cow::Owned(bytes) + } + fn into_bytes(self) -> Vec { + self.to_bytes().into_owned() + } + fn from_bytes(bytes: Cow<'_, [u8]>) -> Self { + Self { + id: u32::from_be_bytes(bytes[0..4].try_into().unwrap()), + tag: bytes[4], + } + } + const BOUND: crate::storable::Bound = crate::storable::Bound::Bounded { + max_size: 5, + is_fixed_size: true, + }; + } + + /// `OccupiedEntry::key` must report the key held by the map, like the std lib does, + /// not the `Ord`-equal one that was passed to `entry`. + #[test] + fn occupied_entry_reports_the_stored_key() { + let mut map: BTreeMap = BTreeMap::new(Rc::new(RefCell::new(Vec::new()))); + map.insert(TaggedKey { id: 1, tag: 7 }, 100); + + let Entry::Occupied(e) = map.entry(TaggedKey { id: 1, tag: 99 }) else { + panic!("key 1 is present"); + }; + assert_eq!(e.key().tag, 7, "`key` must return the stored key"); + assert_eq!(e.into_key().tag, 7, "`into_key` must return the stored key"); + + // Overwriting the value must leave the stored key alone, also matching the std lib. + if let Entry::Occupied(e) = map.entry(TaggedKey { id: 1, tag: 42 }) { + e.insert(101); + } + assert_eq!(map.iter().next().unwrap().key().tag, 7); + + // A vacant entry has no stored key, so it reports the one it was given. + let Entry::Vacant(e) = map.entry(TaggedKey { id: 2, tag: 55 }) else { + panic!("key 2 is absent"); + }; + assert_eq!(e.key().tag, 55); + } + + #[test] + fn entry_end_to_end() { + let mut map = new_map(); + + for i in 0u32..100 { + let Entry::Vacant(e) = map.entry(i) else { + panic!(); + }; + e.insert(i); + } + + for i in 0u32..100 { + let Entry::Occupied(e) = map.entry(i) else { + panic!(); + }; + assert_eq!(i, e.get()); + let old = e.insert(i + 1).into_value(); + assert_eq!(old, i); + } + + for i in 0u32..100 { + let Entry::Occupied(e) = map.entry(i) else { + panic!(); + }; + assert_eq!(i + 1, e.get()); + let removed = e.remove().into_value(); + assert_eq!(removed, i + 1); + } + + assert!(map.is_empty()); + } + + #[test] + fn or_insert_vacant() { + let mut map = new_map(); + assert_eq!(map.entry(1).or_insert(42).get(), 42); + assert_eq!(map.get(&1), Some(42)); + } + + #[test] + fn or_insert_occupied() { + let mut map = new_map(); + map.insert(1, 10); + assert_eq!(map.entry(1).or_insert(99).get(), 10); // default ignored + } + + #[test] + fn or_insert_with() { + let mut map = new_map(); + map.entry(1).or_insert_with(|| 7u32); + assert_eq!(map.get(&1), Some(7)); + // closure is not called when key is present + map.entry(1) + .or_insert_with(|| panic!("should not be called")); + assert_eq!(map.get(&1), Some(7)); + } + + #[test] + fn or_insert_with_key() { + let mut map = new_map(); + map.entry(6).or_insert_with_key(|&k| k * 3); + assert_eq!(map.get(&6), Some(18)); + } + + #[test] + fn or_default() { + let mut map = new_map(); + map.entry(1).or_default(); + assert_eq!(map.get(&1), Some(0u32)); + } + + #[test] + fn and_modify_occupied() { + let mut map = new_map(); + map.insert(1, 10); + map.entry(1).and_modify(|v| *v += 5); + assert_eq!(map.get(&1), Some(15)); + } + + #[test] + fn and_modify_vacant() { + let mut map = new_map(); + // closure must not be called; map must stay empty + map.entry(1).and_modify(|_| panic!("should not be called")); + assert_eq!(map.get(&1), None); + } + + #[test] + fn and_modify_then_or_insert() { + let mut map = new_map(); + map.insert(1, 10u32); + + map.entry(1).and_modify(|v| *v += 1).or_insert(1); + assert_eq!(map.get(&1), Some(11)); + + map.entry(2).and_modify(|v| *v += 1).or_insert(1); + assert_eq!(map.get(&2), Some(1)); + } + + #[test] + fn occupied_insert_returns_old_value() { + let mut map = new_map(); + map.insert(1, 10); + let Entry::Occupied(e) = map.entry(1) else { + panic!(); + }; + assert_eq!(e.insert(99).into_value(), 10); + assert_eq!(map.get(&1), Some(99)); + } + + #[test] + fn occupied_remove_returns_value() { + let mut map = new_map(); + map.insert(1, 42); + let Entry::Occupied(e) = map.entry(1) else { + panic!(); + }; + assert_eq!(e.remove().into_value(), 42); + assert!(map.is_empty()); + } + + #[test] + fn entry_uses_the_node_cache_along_the_path() { + let mut map = new_map().with_node_cache(32); + for i in 0u32..1000 { + map.insert(i, i); + } + + // Warm the cache along the path, then probe the same key again: the nodes on + // the way down are taken from the cache and put back by the traversal, so the + // second probe hits rather than re-reading them from memory. + let _ = map.entry(500); + map.node_cache_reset_metrics(); + let _ = map.entry(500); + assert!(map.node_cache_metrics().hits() > 0); + + // The same holds for a vacant entry dropped without inserting, and the map is + // unchanged by the probe. + let _ = map.entry(5000); + map.node_cache_reset_metrics(); + let _ = map.entry(5000); + assert!(map.node_cache_metrics().hits() > 0); + assert_eq!(map.get(&5000), None); + } + + #[test] + fn reading_and_writing_through_an_entry_needs_no_further_loads() { + let mut map = new_map().with_node_cache(32); + for i in 0u32..1000 { + map.insert(i, i); + } + + let Entry::Occupied(e) = map.entry(500) else { + panic!(); + }; + // The entry holds its node, so reading and writing through it touch the cache + // not at all — no lookups, hence no loads from memory. + e.map.node_cache_reset_metrics(); + assert_eq!(e.get(), 500); + let e = e.and_modify(|v| *v += 1); + assert_eq!(e.get(), 501); + assert_eq!(e.map.node_cache_metrics().total(), 0); + + assert_eq!(map.get(&500), Some(501)); + } + + #[test] + fn entry_mutations_do_not_leave_stale_nodes_in_cache() { + let mut map = new_map().with_node_cache(32); + for i in 0u32..1000 { + map.insert(i, i); + } + for i in 0u32..1000 { + // Warm the cache along the key's path, mutate through the entry API, + // then read back through the cached path. + assert_eq!(map.get(&i), Some(i)); + map.entry(i).and_modify(|v| *v += 1); + assert_eq!(map.get(&i), Some(i + 1)); + } + } + + #[test] + fn or_insert_on_empty_map_then_drop() { + // Dropping a VacantEntry on an empty map without inserting must not corrupt the map. + let mut map = new_map(); + map.entry(1).and_modify(|_| panic!("should not be called")); + assert_eq!(map.get(&1), None); + // The map must still be usable after the drop. + map.insert(1, 99); + assert_eq!(map.get(&1), Some(99)); + } +} diff --git a/src/btreemap/proptests.rs b/src/btreemap/proptests.rs index 8e193823..17a142ce 100644 --- a/src/btreemap/proptests.rs +++ b/src/btreemap/proptests.rs @@ -1,3 +1,4 @@ +use crate::btreemap::entry::Entry; use crate::{ btreemap::{ tests::{b, make_memory, run_btree_test}, @@ -9,7 +10,8 @@ use crate::{ use proptest::collection::btree_set as pset; use proptest::collection::vec as pvec; use proptest::prelude::*; -use std::collections::{BTreeMap as StdBTreeMap, BTreeSet}; +use std::collections::{btree_map, BTreeMap as StdBTreeMap, BTreeSet}; +use std::ops::BitXor; use test_strategy::proptest; #[derive(Debug, Clone)] @@ -21,6 +23,8 @@ enum Operation { Values { from: usize, len: usize }, Get(usize), Remove(usize), + EntryInsertOrXor { key: Vec, value: Vec }, + EntryRemove(usize), Range { from: usize, len: usize }, PopLast, PopFirst, @@ -43,6 +47,9 @@ fn operation_strategy() -> impl Strategy { .prop_map(|(from, len)| Operation::Values { from, len }), 50 => (any::()).prop_map(Operation::Get), 15 => (any::()).prop_map(Operation::Remove), + 10 => (any::>(), any::>()) + .prop_map(|(key, value)| Operation::EntryInsertOrXor { key, value }), + 10 => (any::()).prop_map(Operation::EntryRemove), 5 => (any::(), any::()) .prop_map(|(from, len)| Operation::Range { from, len }), 2 => Just(Operation::PopFirst), @@ -256,6 +263,72 @@ fn no_memory_leaks(#[strategy(pvec(pvec(0..u8::MAX, 100..10_000), 100))] keys: V assert_eq!(btree.allocator.num_allocated_chunks(), 0); } +// A node holds up to `CAPACITY` (11) entries, so the operation count has to comfortably +// exceed that for the tree to actually split and merge. This is also the only entry-API +// coverage of the V1 and V1-migrated-to-V2 layouts, via `run_btree_test`. +#[proptest] +fn entry( + #[strategy(pvec(0..255u8, 200))] keys: Vec, + #[strategy(pvec(0..3u8, 200))] operations: Vec, +) { + run_btree_test(|mut btree| { + let mut std_map = StdBTreeMap::new(); + + // Operations (if Occupied): + // 0 - insert + // 1 - increment + // 2 - remove + // + // Operations (if Vacant): + // - always insert + for (key, operation) in keys.iter().copied().zip(operations.iter().copied()) { + let entry = btree.entry(key); + let std_entry = std_map.entry(key); + let occupied = matches!(entry, Entry::Occupied(_)); + let std_occupied = matches!(std_entry, btree_map::Entry::Occupied(_)); + assert_eq!(occupied, std_occupied); + + match operation { + 0 => { + entry.and_modify(|v| *v = key).or_insert(key); + std_entry.and_modify(|v| *v = key).or_insert(key); + } + 1 => { + entry.and_modify(|v| *v = v.wrapping_add(1)).or_insert(key); + std_entry + .and_modify(|v| *v = v.wrapping_add(1)) + .or_insert(key); + } + 2 => { + match entry { + Entry::Occupied(e) => { + e.remove(); + } + Entry::Vacant(e) => { + e.insert(key); + } + } + match std_entry { + btree_map::Entry::Occupied(e) => { + e.remove(); + } + btree_map::Entry::Vacant(e) => { + e.insert(key); + } + } + } + _ => unreachable!(), + } + } + + let entries: Vec<_> = btree.iter().map(|e| (*e.key(), e.value())).collect(); + let std_entries: Vec<_> = std_map.into_iter().collect(); + + prop_assert_eq!(entries, std_entries); + Ok(()) + }); +} + // Given an operation, executes it on the given stable btreemap and standard btreemap, verifying // that the result of the operation is equal in both btrees. fn execute_operation( @@ -384,6 +457,57 @@ fn execute_operation( assert_eq!(btree.remove(&k), Some(v)); } } + Operation::EntryInsertOrXor { key, value } => { + std_btree + .entry(key.clone()) + .and_modify(|existing| { + *existing = existing + .iter() + .zip(value.clone()) + .map(|(l, r)| l.bitxor(r)) + .collect::>(); + }) + .or_insert(value.clone()); + + btree + .entry(key.clone()) + .and_modify(|existing| { + *existing = existing + .iter() + .zip(value.clone()) + .map(|(l, r)| l.bitxor(r)) + .collect::>(); + }) + .or_insert(value); + + assert_eq!(btree.get(&key).as_ref(), std_btree.get(&key)); + } + Operation::EntryRemove(idx) => { + assert_eq!(std_btree.len(), btree.len() as usize); + if std_btree.is_empty() { + return; + } + + let idx = idx % std_btree.len(); + + if let Some(k) = btree + .iter() + .skip(idx) + .take(1) + .next() + .map(|entry| entry.key().clone()) + { + eprintln!("EntryRemove({})", hex::encode(&k)); + let expected = std_btree.remove(&k).expect("key must exist in std map"); + match btree.entry(k) { + Entry::Occupied(e) => { + assert_eq!(e.remove().into_value(), expected); + } + Entry::Vacant(_) => panic!("entry must be occupied"), + } + } + } + Operation::Range { from, len } => { assert_eq!(std_btree.len(), btree.len() as usize); if std_btree.is_empty() {