Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

feat: move initialization of StableBTreeMap cursors into Iter #231

Open
wants to merge 9 commits into
base: main
Choose a base branch
from
+185 −160
Open

feat: move initialization of StableBTreeMap cursors into Iter #231

Show file tree
Hide file tree
Changes from 5 commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
a862edf
feat: move initialization of `StableBTreeMap` cursors into `Iter`
hpeebles Oct 16, 2024
a720764
Add `new_with_cursors` ctor
hpeebles Oct 16, 2024
d2306c0
Not sure why the formatting went weird there
hpeebles Oct 16, 2024
03d78a8
Add comment
hpeebles Oct 16, 2024
db8a197
Tiny speedup
hpeebles Oct 17, 2024
b38d79d
Add comment
hpeebles Oct 25, 2024
7b388e8
Merge branch 'main' into iter
hpeebles Oct 25, 2024
2c2d7eb
Merge branch 'main' into iter
dsarlis Oct 29, 2024
b323b0a
Update canbench results
hpeebles Nov 5, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
95 changes: 3 additions & 92 deletions src/btreemap.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -1161,15 +1072,15 @@ 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);

cursors.push(Cursor::Node {
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));
Expand Down
118 changes: 115 additions & 3 deletions src/btreemap/iter.rs
View file
Open in desktop
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,
dsarlis marked this conversation as resolved.
Show resolved Hide resolved

// 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,135 @@ 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(
pub(crate) fn new_in_range(map: &'a BTreeMap<K, V, M>, range: (Bound<K>, Bound<K>)) -> 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<K, V, M>,
range: (Bound<K>, Bound<K>),
cursors: Vec<Cursor<K>>,
) -> 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<T, F: Fn(&Node<K>, usize) -> T>(&mut self, map: F) -> Option<T> {
if !self.cursors_initialized {
self.initialize_cursors();
}

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