Skip to content

Commit

Permalink
feat: implement DoubleEndedIterator for StableBTreeMap (#235)
Browse files Browse the repository at this point in the history
This adds the ability to iterate a `StableBTreeMap` in either direction.
We have been using this within OpenChat for a few weeks now with no
issues.
  • Loading branch information
hpeebles authored Nov 12, 2024
1 parent a206977 commit 6688005
Show file tree
Hide file tree
Showing 5 changed files with 455 additions and 135 deletions.
40 changes: 40 additions & 0 deletions benchmarks/src/btreemap.rs
Original file line number Diff line number Diff line change
Expand Up @@ -220,6 +220,26 @@ pub fn btreemap_insert_10mib_values() -> BenchResult {
})
}

#[bench(raw)]
pub fn btreemap_iter_small_values() -> BenchResult {
iter_helper(10_000, 0)
}

#[bench(raw)]
pub fn btreemap_iter_rev_small_values() -> BenchResult {
iter_rev_helper(10_000, 0)
}

#[bench(raw)]
pub fn btreemap_iter_10mib_values() -> BenchResult {
iter_helper(200, 10 * 1024)
}

#[bench(raw)]
pub fn btreemap_iter_rev_10mib_values() -> BenchResult {
iter_rev_helper(200, 10 * 1024)
}

#[bench(raw)]
pub fn btreemap_iter_count_small_values() -> BenchResult {
let mut btree = BTreeMap::new(DefaultMemoryImpl::default());
Expand Down Expand Up @@ -517,6 +537,26 @@ fn insert_helper<K: Clone + Ord + Storable + Random, V: Storable + Random>(
})
}

// Profiles iterating over a btreemap.
fn iter_helper(size: u32, value_size: u32) -> BenchResult {
let mut btree = BTreeMap::new(DefaultMemoryImpl::default());
for i in 0..size {
btree.insert(i, vec![0u8; value_size as usize]);
}

bench_fn(|| for _ in btree.iter() {})
}

// Profiles iterating in reverse over a btreemap.
fn iter_rev_helper(size: u32, value_size: u32) -> BenchResult {
let mut btree = BTreeMap::new(DefaultMemoryImpl::default());
for i in 0..size {
btree.insert(i, vec![0u8; value_size as usize]);
}

bench_fn(|| for _ in btree.iter().rev() {})
}

// Profiles getting a large number of random blobs from a btreemap.
fn get_blob_helper<const K: usize, const V: usize>() -> BenchResult {
let btree = BTreeMap::new_v1(DefaultMemoryImpl::default());
Expand Down
32 changes: 28 additions & 4 deletions canbench_results.yml
Original file line number Diff line number Diff line change
Expand Up @@ -365,15 +365,39 @@ benches:
heap_increase: 0
stable_memory_increase: 6
scopes: {}
btreemap_iter_10mib_values:
total:
instructions: 25583733
heap_increase: 0
stable_memory_increase: 0
scopes: {}
btreemap_iter_count_10mib_values:
total:
instructions: 525036
instructions: 544088
heap_increase: 0
stable_memory_increase: 0
scopes: {}
btreemap_iter_count_small_values:
total:
instructions: 10458516
instructions: 11007833
heap_increase: 0
stable_memory_increase: 0
scopes: {}
btreemap_iter_rev_10mib_values:
total:
instructions: 25585550
heap_increase: 0
stable_memory_increase: 0
scopes: {}
btreemap_iter_rev_small_values:
total:
instructions: 23878236
heap_increase: 0
stable_memory_increase: 0
scopes: {}
btreemap_iter_small_values:
total:
instructions: 23721014
heap_increase: 0
stable_memory_increase: 0
scopes: {}
Expand Down Expand Up @@ -517,13 +541,13 @@ benches:
scopes: {}
memory_manager_grow:
total:
instructions: 351687872
instructions: 350727867
heap_increase: 2
stable_memory_increase: 32000
scopes: {}
memory_manager_overhead:
total:
instructions: 1182143127
instructions: 1182141676
heap_increase: 0
stable_memory_increase: 8320
scopes: {}
Expand Down
73 changes: 4 additions & 69 deletions src/btreemap.rs
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,6 @@ use crate::{
};
use allocator::Allocator;
pub use iter::Iter;
use iter::{Cursor, Index};
use node::{DerivedPageSize, Entry, Node, NodeType, PageSize, Version};
use std::borrow::Cow;
use std::marker::PhantomData;
Expand Down Expand Up @@ -1029,74 +1028,10 @@ where
/// 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 self.root_addr == NULL {
// Map is empty.
return Iter::null(self);
}

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

let mut node = self.load_node(self.root_addr);
loop {
match node.search(bound) {
Ok(idx) | Err(idx) => {
match node.node_type() {
NodeType::Leaf => {
if idx == 0 {
// We descended into a leaf but didn't find a node less than
// the upper bound. Thus we unwind the cursor stack until we
// hit a cursor pointing to an element other than the first key,
// and we shift the position backward. If there is no such cursor,
// the bound must be <= min element, so we return an empty iterator.
while let Some(cursor) = cursors.pop() {
match cursor {
Cursor::Node {
node,
next: Index::Entry(n),
} => {
if n == 0 {
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;
}
}
_ => panic!("BUG: unexpected cursor shape"),
}
}
// If the cursors are empty, the iterator will be empty.
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_with_cursors(self, dummy_bounds, cursors);
}
NodeType::Internal => {
let child = self.load_node(node.child(idx));
// We push the node even if idx == node.entries_len()
// If we find the position in the child, the iterator will skip this
// cursor. But if the all keys in the child are greater than or equal to
// the bound, we will be able to use this cursor as a fallback.
cursors.push(Cursor::Node {
node,
next: Index::Entry(idx),
});
node = child;
}
}
}
}
if let Some((start_key, _)) = self.range(..bound).next_back() {
Iter::new_in_range(self, (Bound::Included(start_key), Bound::Unbounded))
} else {
Iter::null(self)
}
}

Expand Down
Loading

0 comments on commit 6688005

Please sign in to comment.