From aaaaac3cc155aaa66ac56eaf2b19389036b2091c Mon Sep 17 00:00:00 2001 From: Hamish Peebles Date: Mon, 4 Nov 2024 22:09:52 +0000 Subject: [PATCH] feat: move initialization of `StableBTreeMap` cursors into `Iter` --- src/btreemap.rs | 95 ++-------------------------------- src/btreemap/iter.rs | 120 +++++++++++++++++++++++++++++++++++++++++-- 2 files changed, 120 insertions(+), 95 deletions(-) diff --git a/src/btreemap.rs b/src/btreemap.rs index 716eaf77..47c6ab58 100644 --- a/src/btreemap.rs +++ b/src/btreemap.rs @@ -1023,96 +1023,7 @@ where key_range.end_bound().cloned(), ); - let mut cursors = vec![]; - - match key_range.start_bound() { - Bound::Unbounded => { - cursors.push(Cursor::Address(self.root_addr)); - Iter::new_in_range(self, range, cursors) - } - Bound::Included(key) | Bound::Excluded(key) => { - let mut node = self.load_node(self.root_addr); - loop { - match node.search(key) { - Ok(idx) => { - if let Bound::Included(_) = key_range.start_bound() { - // We found the key exactly matching the left bound. - // Here is where we'll start the iteration. - cursors.push(Cursor::Node { - node, - next: Index::Entry(idx), - }); - return Iter::new_in_range(self, range, cursors); - } else { - // We found the key that we must - // exclude. We add its right neighbor - // to the stack and start iterating - // from its right child. - let right_child = match node.node_type() { - NodeType::Internal => Some(node.child(idx + 1)), - NodeType::Leaf => None, - }; - - if idx + 1 != node.entries_len() - && key_range.contains(node.key(idx + 1)) - { - cursors.push(Cursor::Node { - node, - next: Index::Entry(idx + 1), - }); - } - if let Some(right_child) = right_child { - cursors.push(Cursor::Address(right_child)); - } - return Iter::new_in_range(self, range, cursors); - } - } - Err(idx) => { - // The `idx` variable points to the first - // key that is greater than the left - // bound. - // - // If the index points to a valid node, we - // will visit its left subtree and then - // return to this key. - // - // If the index points at the end of - // array, we'll continue with the right - // child of the last key. - - // Load the left child of the node to visit if it exists. - // This is done first to avoid cloning the node. - let child = match node.node_type() { - NodeType::Internal => { - // Note that loading a child node cannot fail since - // len(children) = len(entries) + 1 - Some(self.load_node(node.child(idx))) - } - NodeType::Leaf => None, - }; - - if idx < node.entries_len() && key_range.contains(node.key(idx)) { - cursors.push(Cursor::Node { - node, - next: Index::Entry(idx), - }); - } - - match child { - None => { - // Leaf node. Return an iterator with the found cursors. - return Iter::new_in_range(self, range, cursors); - } - Some(child) => { - // Iterate over the child node. - node = child; - } - } - } - } - } - } - } + Iter::new_in_range(self, range) } /// Returns an iterator pointing to the first element below the given bound. @@ -1161,7 +1072,7 @@ where } } // If the cursors are empty, the iterator will be empty. - return Iter::new_in_range(self, dummy_bounds, cursors); + return Iter::new_with_cursors(self, dummy_bounds, cursors); } debug_assert!(node.key(idx - 1) < bound); @@ -1169,7 +1080,7 @@ where node, next: Index::Entry(idx - 1), }); - return Iter::new_in_range(self, dummy_bounds, cursors); + return Iter::new_with_cursors(self, dummy_bounds, cursors); } NodeType::Internal => { let child = self.load_node(node.child(idx)); diff --git a/src/btreemap/iter.rs b/src/btreemap/iter.rs index 3949352f..2492aca7 100644 --- a/src/btreemap/iter.rs +++ b/src/btreemap/iter.rs @@ -29,6 +29,11 @@ where // A reference to the map being iterated on. map: &'a BTreeMap, + // A flag indicating if the cursors have been initialized yet. This is needed to distinguish + // between the case where the iteration hasn't started yet and the case where the iteration has + // finished (in both cases the `cursors` field will be empty). + cursors_initialized: bool, + // A stack of cursors indicating the current position in the tree. cursors: Vec>, @@ -45,8 +50,8 @@ where pub(crate) fn new(map: &'a BTreeMap) -> Self { Self { map, - // Initialize the cursors with the address of the root of the map. - cursors: vec![Cursor::Address(map.root_addr)], + cursors_initialized: false, + cursors: vec![], range: (Bound::Unbounded, Bound::Unbounded), } } @@ -55,26 +60,135 @@ where pub(crate) fn null(map: &'a BTreeMap) -> Self { Self { map, + cursors_initialized: true, cursors: vec![], range: (Bound::Unbounded, Bound::Unbounded), } } - pub(crate) fn new_in_range( + pub(crate) fn new_in_range(map: &'a BTreeMap, range: (Bound, Bound)) -> Self { + Self { + map, + cursors_initialized: false, + cursors: vec![], + range, + } + } + + // This can be used as an optimisation if the cursors have already been calculated + pub(crate) fn new_with_cursors( map: &'a BTreeMap, range: (Bound, Bound), cursors: Vec>, ) -> Self { Self { map, + cursors_initialized: true, cursors, range, } } + fn initialize_cursors(&mut self) { + debug_assert!(!self.cursors_initialized); + + match self.range.start_bound() { + Bound::Unbounded => { + self.cursors.push(Cursor::Address(self.map.root_addr)); + } + Bound::Included(key) | Bound::Excluded(key) => { + let mut node = self.map.load_node(self.map.root_addr); + loop { + match node.search(key) { + Ok(idx) => { + if let Bound::Included(_) = self.range.start_bound() { + // We found the key exactly matching the left bound. + // Here is where we'll start the iteration. + self.cursors.push(Cursor::Node { + node, + next: Index::Entry(idx), + }); + break; + } else { + // We found the key that we must + // exclude. We add its right neighbor + // to the stack and start iterating + // from its right child. + let right_child = match node.node_type() { + NodeType::Internal => Some(node.child(idx + 1)), + NodeType::Leaf => None, + }; + + if idx + 1 != node.entries_len() + && self.range.contains(node.key(idx + 1)) + { + self.cursors.push(Cursor::Node { + node, + next: Index::Entry(idx + 1), + }); + } + if let Some(right_child) = right_child { + self.cursors.push(Cursor::Address(right_child)); + } + break; + } + } + Err(idx) => { + // The `idx` variable points to the first + // key that is greater than the left + // bound. + // + // If the index points to a valid node, we + // will visit its left subtree and then + // return to this key. + // + // If the index points at the end of + // array, we'll continue with the right + // child of the last key. + + // Load the left child of the node to visit if it exists. + // This is done first to avoid cloning the node. + let child = match node.node_type() { + NodeType::Internal => { + // Note that loading a child node cannot fail since + // len(children) = len(entries) + 1 + Some(self.map.load_node(node.child(idx))) + } + NodeType::Leaf => None, + }; + + if idx < node.entries_len() && self.range.contains(node.key(idx)) { + self.cursors.push(Cursor::Node { + node, + next: Index::Entry(idx), + }); + } + + match child { + None => { + // Leaf node. Return an iterator with the found cursors. + break; + } + Some(child) => { + // Iterate over the child node. + node = child; + } + } + } + } + } + } + } + self.cursors_initialized = true; + } + // Iterates to find the next element in the requested range. // If it exists, `map` is applied to that element and the result is returned. fn next_map, usize) -> T>(&mut self, map: F) -> Option { + if !self.cursors_initialized { + self.initialize_cursors(); + } + // If the cursors are empty. Iteration is complete. match self.cursors.pop()? { Cursor::Address(address) => {