diff --git a/benchmarks/src/btreemap.rs b/benchmarks/src/btreemap.rs index 5e1e92f3..d280a167 100644 --- a/benchmarks/src/btreemap.rs +++ b/benchmarks/src/btreemap.rs @@ -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()); @@ -517,6 +537,26 @@ fn insert_helper( }) } +// 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() -> BenchResult { let btree = BTreeMap::new_v1(DefaultMemoryImpl::default()); diff --git a/canbench_results.yml b/canbench_results.yml index 62a7066d..830313c6 100644 --- a/canbench_results.yml +++ b/canbench_results.yml @@ -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: {} @@ -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: {} diff --git a/src/btreemap.rs b/src/btreemap.rs index 47c6ab58..3bfa1bfd 100644 --- a/src/btreemap.rs +++ b/src/btreemap.rs @@ -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; @@ -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 { - 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) } } diff --git a/src/btreemap/iter.rs b/src/btreemap/iter.rs index 2492aca7..848757c2 100644 --- a/src/btreemap/iter.rs +++ b/src/btreemap/iter.rs @@ -29,13 +29,15 @@ 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 + // Flags indicating if the cursors have been initialized yet. These are 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, + // finished (in both cases the cursors will be empty). + forward_cursors_initialized: bool, + backward_cursors_initialized: bool, - // A stack of cursors indicating the current position in the tree. - cursors: Vec>, + // Stacks of cursors indicating the current iteration positions in the tree. + forward_cursors: Vec>, + backward_cursors: Vec>, // The range of keys we want to traverse. range: (Bound, Bound), @@ -50,8 +52,10 @@ where pub(crate) fn new(map: &'a BTreeMap) -> Self { Self { map, - cursors_initialized: false, - cursors: vec![], + forward_cursors_initialized: false, + backward_cursors_initialized: false, + forward_cursors: vec![], + backward_cursors: vec![], range: (Bound::Unbounded, Bound::Unbounded), } } @@ -60,8 +64,10 @@ where pub(crate) fn null(map: &'a BTreeMap) -> Self { Self { map, - cursors_initialized: true, - cursors: vec![], + forward_cursors_initialized: true, + backward_cursors_initialized: true, + forward_cursors: vec![], + backward_cursors: vec![], range: (Bound::Unbounded, Bound::Unbounded), } } @@ -69,32 +75,21 @@ where pub(crate) fn new_in_range(map: &'a BTreeMap, range: (Bound, Bound)) -> Self { Self { map, - cursors_initialized: false, - cursors: vec![], + forward_cursors_initialized: false, + backward_cursors_initialized: false, + forward_cursors: vec![], + backward_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); + fn initialize_forward_cursors(&mut self) { + debug_assert!(!self.forward_cursors_initialized); match self.range.start_bound() { Bound::Unbounded => { - self.cursors.push(Cursor::Address(self.map.root_addr)); + self.forward_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); @@ -104,7 +99,7 @@ where 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 { + self.forward_cursors.push(Cursor::Node { node, next: Index::Entry(idx), }); @@ -119,16 +114,16 @@ where NodeType::Leaf => None, }; - if idx + 1 != node.entries_len() + if idx + 1 < node.entries_len() && self.range.contains(node.key(idx + 1)) { - self.cursors.push(Cursor::Node { + self.forward_cursors.push(Cursor::Node { node, next: Index::Entry(idx + 1), }); } if let Some(right_child) = right_child { - self.cursors.push(Cursor::Address(right_child)); + self.forward_cursors.push(Cursor::Address(right_child)); } break; } @@ -158,10 +153,101 @@ where }; if idx < node.entries_len() && self.range.contains(node.key(idx)) { - self.cursors.push(Cursor::Node { + self.forward_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.forward_cursors_initialized = true; + } + + fn initialize_backward_cursors(&mut self) { + debug_assert!(!self.backward_cursors_initialized); + + match self.range.end_bound() { + Bound::Unbounded => { + self.backward_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.end_bound() { + // We found the key exactly matching the right bound. + // Here is where we'll start the iteration. + self.backward_cursors.push(Cursor::Node { node, next: Index::Entry(idx), }); + break; + } else { + // We found the key that we must + // exclude. We add its left neighbor + // to the stack and start iterating + // from its left child. + let left_child = match node.node_type() { + NodeType::Internal => Some(node.child(idx)), + NodeType::Leaf => None, + }; + + if idx > 0 && self.range.contains(node.key(idx - 1)) { + self.backward_cursors.push(Cursor::Node { + node, + next: Index::Entry(idx - 1), + }); + } + if let Some(left_child) = left_child { + self.backward_cursors.push(Cursor::Address(left_child)); + } + break; + } + } + Err(idx) => { + // The `idx` variable points to the first + // key that is greater than the right + // bound. + // + // If the index points to a valid node, we + // will visit its left subtree. + // + // 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 > 0 && self.range.contains(node.key(idx - 1)) { + self.backward_cursors.push(Cursor::Node { + node, + next: Index::Entry(idx - 1), + }); } match child { @@ -179,23 +265,23 @@ where } } } - self.cursors_initialized = true; + self.backward_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 !self.forward_cursors_initialized { + self.initialize_forward_cursors(); } // If the cursors are empty. Iteration is complete. - match self.cursors.pop()? { + match self.forward_cursors.pop()? { Cursor::Address(address) => { if address != NULL { // Load the node at the given address, and add it to the cursors. let node = self.map.load_node(address); - self.cursors.push(Cursor::Node { + self.forward_cursors.push(Cursor::Node { next: match node.node_type() { // Iterate on internal nodes starting from the first child. NodeType::Internal => Index::Child(0), @@ -214,15 +300,17 @@ where } => { let child_address = node.child(child_idx); - // After iterating on the child, iterate on the next _entry_ in this node. - // The entry immediately after the child has the same index as the child's. - self.cursors.push(Cursor::Node { - node, - next: Index::Entry(child_idx), - }); + if child_idx < node.entries_len() { + // After iterating on the child, iterate on the next _entry_ in this node. + // The entry immediately after the child has the same index as the child's. + self.forward_cursors.push(Cursor::Node { + node, + next: Index::Entry(child_idx), + }); + } // Add the child to the top of the cursors to be iterated on first. - self.cursors.push(Cursor::Address(child_address)); + self.forward_cursors.push(Cursor::Address(child_address)); self.next_map(map) } @@ -231,30 +319,110 @@ where node, next: Index::Entry(entry_idx), } => { - if entry_idx >= node.entries_len() { - // No more entries to iterate on in this node. - return self.next_map(map); - } - // If the key does not belong to the range, iteration stops. if !self.range.contains(node.key(entry_idx)) { // Clear all cursors to avoid needless work in subsequent calls. - self.cursors = vec![]; + self.forward_cursors = vec![]; + self.backward_cursors = vec![]; return None; } let res = map(&node, entry_idx); + self.range.0 = Bound::Excluded(node.key(entry_idx).clone()); + + let next = match node.node_type() { + // If this is an internal node, add the next child to the cursors. + NodeType::Internal => Index::Child(entry_idx + 1), + // If this is a leaf node, and it contains another entry, add the + // next entry to the cursors. + NodeType::Leaf if entry_idx + 1 < node.entries_len() => { + Index::Entry(entry_idx + 1) + } + _ => return Some(res), + }; // Add to the cursors the next element to be traversed. - self.cursors.push(Cursor::Node { - next: match node.node_type() { - // If this is an internal node, add the next child to the cursors. - NodeType::Internal => Index::Child(entry_idx + 1), - // If this is a leaf node, add the next entry to the cursors. - NodeType::Leaf => Index::Entry(entry_idx + 1), - }, - node, - }); + self.forward_cursors.push(Cursor::Node { next, node }); + Some(res) + } + } + } + + // Iterates to find the next back element in the requested range. + // If it exists, `map` is applied to that element and the result is returned. + fn next_back_map, usize) -> T>(&mut self, map: F) -> Option { + if !self.backward_cursors_initialized { + self.initialize_backward_cursors(); + } + + // If the cursors are empty. Iteration is complete. + match self.backward_cursors.pop()? { + Cursor::Address(address) => { + if address != NULL { + // Load the node at the given address, and add it to the cursors. + let node = self.map.load_node(address); + if let Some(next) = match node.node_type() { + // Iterate on internal nodes starting from the last child. + NodeType::Internal if node.children_len() > 0 => { + Some(Index::Child(node.children_len() - 1)) + } + // Iterate on leaf nodes starting from the last entry. + NodeType::Leaf if node.entries_len() > 0 => { + Some(Index::Entry(node.entries_len() - 1)) + } + _ => None, + } { + self.backward_cursors.push(Cursor::Node { next, node }); + } + } + self.next_back_map(map) + } + + Cursor::Node { + node, + next: Index::Child(child_idx), + } => { + let child_address = node.child(child_idx); + + if 0 < child_idx { + // After iterating on the child, iterate on the previous _entry_ in this node. + self.backward_cursors.push(Cursor::Node { + node, + next: Index::Entry(child_idx - 1), + }); + } + + // Add the child to the top of the cursors to be iterated on first. + self.backward_cursors.push(Cursor::Address(child_address)); + + self.next_back_map(map) + } + + Cursor::Node { + node, + next: Index::Entry(entry_idx), + } => { + // If the key does not belong to the range, iteration stops. + if !self.range.contains(node.key(entry_idx)) { + // Clear all cursors to avoid needless work in subsequent calls. + self.forward_cursors = vec![]; + self.backward_cursors = vec![]; + return None; + } + + let res = map(&node, entry_idx); + self.range.1 = Bound::Excluded(node.key(entry_idx).clone()); + + if let Some(next) = match node.node_type() { + // If this is an internal node, add the previous child to the cursors. + NodeType::Internal => Some(Index::Child(entry_idx)), + // If this is a leaf node, add the previous entry to the cursors. + NodeType::Leaf if entry_idx > 0 => Some(Index::Entry(entry_idx - 1)), + _ => None, + } { + // Add to the cursors the next element to be traversed. + self.backward_cursors.push(Cursor::Node { next, node }); + } Some(res) } @@ -289,6 +457,20 @@ where } } +impl DoubleEndedIterator for Iter<'_, K, V, M> +where + K: Storable + Ord + Clone, + V: Storable, + M: Memory, +{ + fn next_back(&mut self) -> Option { + self.next_back_map(|node, entry_idx| { + let (key, encoded_value) = node.entry(entry_idx, self.map.memory()); + (key, V::from_bytes(Cow::Owned(encoded_value))) + }) + } +} + #[cfg(test)] mod test { use super::*; @@ -338,4 +520,119 @@ mod test { assert_eq!(i, 100); } + + #[test] + fn iterate_range() { + let mem = make_memory(); + let mut btree = BTreeMap::new(mem); + + // Insert the elements in reverse order. + for i in (0..100u64).rev() { + btree.insert(i, i + 1); + } + + // Iteration should be in ascending order. + let mut i = 10; + for (key, value) in btree.range(10..90) { + assert_eq!(key, i); + assert_eq!(value, i + 1); + i += 1; + } + + assert_eq!(i, 90); + } + + #[test] + fn iterate_reverse() { + let mem = make_memory(); + let mut btree = BTreeMap::new(mem); + + // Insert the elements in reverse order. + for i in (0..100u64).rev() { + btree.insert(i, i + 1); + } + + // Iteration should be in descending order. + let mut i = 100; + for (key, value) in btree.iter().rev() { + i -= 1; + assert_eq!(key, i); + assert_eq!(value, i + 1); + } + + assert_eq!(i, 0); + } + + #[test] + fn iterate_range_reverse() { + let mem = make_memory(); + let mut btree = BTreeMap::new(mem); + + // Insert the elements in reverse order. + for i in (0..100u64).rev() { + btree.insert(i, i + 1); + } + + // Iteration should be in descending order. + let mut i = 80; + for (key, value) in btree.range(20..80).rev() { + i -= 1; + assert_eq!(key, i); + assert_eq!(value, i + 1); + } + + assert_eq!(i, 20); + } + + #[test] + fn iterate_from_both_ends() { + let mem = make_memory(); + let mut btree = BTreeMap::new(mem); + + // Insert the elements in reverse order. + for i in (0..100u64).rev() { + btree.insert(i, i + 1); + } + + let mut iter = btree.iter(); + + for i in 0..50 { + let (key, value) = iter.next().unwrap(); + assert_eq!(key, i); + assert_eq!(value, i + 1); + + let (key, value) = iter.next_back().unwrap(); + assert_eq!(key, 99 - i); + assert_eq!(value, 100 - i); + } + + assert!(iter.next().is_none()); + assert!(iter.next_back().is_none()); + } + + #[test] + fn iterate_range_from_both_ends() { + let mem = make_memory(); + let mut btree = BTreeMap::new(mem); + + // Insert the elements in reverse order. + for i in (0..100u64).rev() { + btree.insert(i, i + 1); + } + + let mut iter = btree.range(30..70); + + for i in 0..20 { + let (key, value) = iter.next().unwrap(); + assert_eq!(key, 30 + i); + assert_eq!(value, 31 + i); + + let (key, value) = iter.next_back().unwrap(); + assert_eq!(key, 69 - i); + assert_eq!(value, 70 - i); + } + + assert!(iter.next().is_none()); + assert!(iter.next_back().is_none()); + } } diff --git a/src/btreemap/proptests.rs b/src/btreemap/proptests.rs index 550b67db..f800fde6 100644 --- a/src/btreemap/proptests.rs +++ b/src/btreemap/proptests.rs @@ -16,6 +16,7 @@ use test_strategy::proptest; enum Operation { Insert { key: Vec, value: Vec }, Iter { from: usize, len: usize }, + IterRev { from: usize, len: usize }, Get(usize), Remove(usize), Range { from: usize, len: usize }, @@ -32,6 +33,8 @@ fn operation_strategy() -> impl Strategy { .prop_map(|(key, value)| Operation::Insert { key, value }), 5 => (any::(), any::()) .prop_map(|(from, len)| Operation::Iter { from, len }), + 5 => (any::(), any::()) + .prop_map(|(from, len)| Operation::IterRev { from, len }), 50 => (any::()).prop_map(Operation::Get), 15 => (any::()).prop_map(Operation::Remove), 5 => (any::(), any::()) @@ -204,11 +207,32 @@ fn execute_operation( eprintln!("Iterate({}, {})", from, len); let std_iter = std_btree.iter().skip(from).take(len); - let stable_iter = btree.iter().skip(from).take(len); - for ((k1, v1), (k2, v2)) in std_iter.zip(stable_iter) { + let mut stable_iter = btree.iter().skip(from).take(len); + for (k1, v1) in std_iter { + let (k2, v2) = stable_iter.next().unwrap(); + assert_eq!(k1, &k2); + assert_eq!(v1, &v2); + } + assert!(stable_iter.next().is_none()); + } + Operation::IterRev { from, len } => { + assert_eq!(std_btree.len(), btree.len() as usize); + if std_btree.is_empty() { + return; + } + + let from = from % std_btree.len(); + let len = len % std_btree.len(); + + eprintln!("Iterate({}, {})", from, len); + let std_iter = std_btree.iter().rev().skip(from).take(len); + let mut stable_iter = btree.iter().rev().skip(from).take(len); + for (k1, v1) in std_iter { + let (k2, v2) = stable_iter.next().unwrap(); assert_eq!(k1, &k2); assert_eq!(v1, &v2); } + assert!(stable_iter.next().is_none()); } Operation::Get(idx) => { assert_eq!(std_btree.len(), btree.len() as usize);