Skip to content

Commit

Permalink
feat: move initialization of StableBTreeMap cursors into Iter
Browse files Browse the repository at this point in the history
  • Loading branch information
@hpeebles
hpeebles committed Oct 16, 2024
1 parent 817a163 commit a862edf
Show file tree
Hide file tree
Showing 2 changed files with 123 additions and 115 deletions.
124 changes: 17 additions & 107 deletions src/btreemap.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1023,107 +1023,25 @@ 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.
/// Returns an empty iterator if there are no keys below the given bound.
pub fn iter_upper_bound(&self, bound: &K) -> Iter<K, V, M> {
if let Some(key) = self.find_by_upper_bound(bound) {
Iter::new_in_range(self, (Bound::Included(key), Bound::Unbounded))
} else {
Iter::null(self)
}
}

fn find_by_upper_bound(&self, bound: &K) -> Option<K> {
if self.root_addr == NULL {
// Map is empty.
return Iter::null(self);
return None;
}

let dummy_bounds = (Bound::Unbounded, Bound::Unbounded);
// INVARIANT: all cursors point to keys greater than or equal to bound.
let mut cursors = vec![];

Expand All @@ -1149,27 +1067,19 @@ where
debug_assert!(node.key(n) >= bound);
continue;
} else {
debug_assert!(node.key(n - 1) < bound);
cursors.push(Cursor::Node {
node,
next: Index::Entry(n - 1),
});
break;
let key: &K = node.key(n - 1);
debug_assert!(key < bound);
return Some(key.clone());
}
}
_ => panic!("BUG: unexpected cursor shape"),
}
}
// If the cursors are empty, the iterator will be empty.
return Iter::new_in_range(self, dummy_bounds, cursors);
return None;
}
debug_assert!(node.key(idx - 1) < bound);

cursors.push(Cursor::Node {
node,
next: Index::Entry(idx - 1),
});
return Iter::new_in_range(self, dummy_bounds, cursors);
let key = node.key(idx - 1);
debug_assert!(key < bound);
return Some(key.clone());
}
NodeType::Internal => {
let child = self.load_node(node.child(idx));
Expand Down
114 changes: 106 additions & 8 deletions src/btreemap/iter.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,9 @@ where
// A reference to the map being iterated on.
map: &'a BTreeMap<K, V, M>,

// A flag indicating if the cursors have been initialized yet.
cursors_initialized: bool,

// A stack of cursors indicating the current position in the tree.
cursors: Vec<Cursor<K>>,

Expand All @@ -45,8 +48,8 @@ where
pub(crate) fn new(map: &'a BTreeMap<K, V, M>) -> 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),
}
}
Expand All @@ -55,26 +58,121 @@ where
pub(crate) fn null(map: &'a BTreeMap<K, V, M>) -> Self {
Self {
map,
cursors_initialized: true,
cursors: vec![],
range: (Bound::Unbounded, Bound::Unbounded),
}
}

pub(crate) fn new_in_range(
map: &'a BTreeMap<K, V, M>,
range: (Bound<K>, Bound<K>),
cursors: Vec<Cursor<K>>,
) -> Self {
pub(crate) fn new_in_range(map: &'a BTreeMap<K, V, M>, range: (Bound<K>, Bound<K>)) -> Self {
Self {
map,
cursors,
cursors_initialized: false,
cursors: vec![],
range,
}
}

fn ensure_cursors_initialized(&mut self) {
if self.cursors_initialized {
return;
}

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<T, F: Fn(&Node<K>, usize) -> T>(&mut self, map: F) -> Option<T> {
self.ensure_cursors_initialized();

// If the cursors are empty. Iteration is complete.
match self.cursors.pop()? {
Cursor::Address(address) => {
Expand Down

0 comments on commit a862edf

Please sign in to comment.