diff --git a/canbench_results.yml b/canbench_results.yml index e3f4254d..f97d6a1a 100644 --- a/canbench_results.yml +++ b/canbench_results.yml @@ -367,13 +367,13 @@ benches: scopes: {} btreemap_iter_count_10mib_values: total: - instructions: 23609229 + instructions: 497381 heap_increase: 0 stable_memory_increase: 0 scopes: {} btreemap_iter_count_small_values: total: - instructions: 25579067 + instructions: 10116842 heap_increase: 0 stable_memory_increase: 0 scopes: {} diff --git a/src/btreemap/iter.rs b/src/btreemap/iter.rs index 5b1dc0a9..3949352f 100644 --- a/src/btreemap/iter.rs +++ b/src/btreemap/iter.rs @@ -71,19 +71,13 @@ where range, } } -} - -impl Iterator for Iter<'_, K, V, M> -where - K: Storable + Ord + Clone, - V: Storable, - M: Memory, -{ - type Item = (K, V); - fn next(&mut self) -> Option { - match self.cursors.pop() { - Some(Cursor::Address(address)) => { + // 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 the cursors are empty. Iteration is complete. + match self.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); @@ -97,13 +91,13 @@ where node, }); } - self.next() + self.next_map(map) } - Some(Cursor::Node { + Cursor::Node { node, next: Index::Child(child_idx), - }) => { + } => { let child_address = node.child(child_idx); // After iterating on the child, iterate on the next _entry_ in this node. @@ -116,19 +110,26 @@ where // Add the child to the top of the cursors to be iterated on first. self.cursors.push(Cursor::Address(child_address)); - self.next() + self.next_map(map) } - Some(Cursor::Node { + Cursor::Node { node, next: Index::Entry(entry_idx), - }) => { + } => { if entry_idx >= node.entries_len() { // No more entries to iterate on in this node. - return self.next(); + return self.next_map(map); } - let (key, encoded_value) = node.entry(entry_idx, self.map.memory()); + // 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![]; + return None; + } + + let res = map(&node, entry_idx); // Add to the cursors the next element to be traversed. self.cursors.push(Cursor::Node { @@ -141,23 +142,39 @@ where node, }); - // If the key does not belong to the range, iteration stops. - if !self.range.contains(&key) { - // Clear all cursors to avoid needless work in subsequent calls. - self.cursors = vec![]; - return None; - } - - Some((key, V::from_bytes(Cow::Owned(encoded_value)))) - } - None => { - // The cursors are empty. Iteration is complete. - None + Some(res) } } } } +impl Iterator for Iter<'_, K, V, M> +where + K: Storable + Ord + Clone, + V: Storable, + M: Memory, +{ + type Item = (K, V); + + fn next(&mut self) -> Option { + self.next_map(|node, entry_idx| { + let (key, encoded_value) = node.entry(entry_idx, self.map.memory()); + (key, V::from_bytes(Cow::Owned(encoded_value))) + }) + } + + fn count(mut self) -> usize + where + Self: Sized, + { + let mut cnt = 0; + while self.next_map(|_, _| ()).is_some() { + cnt += 1; + } + cnt + } +} + #[cfg(test)] mod test { use super::*; diff --git a/src/btreemap/proptests.rs b/src/btreemap/proptests.rs index 175af763..144d97e9 100644 --- a/src/btreemap/proptests.rs +++ b/src/btreemap/proptests.rs @@ -144,6 +144,21 @@ fn map_upper_bound_iter(#[strategy(pvec(0u64..u64::MAX -1 , 10..100))] keys: Vec }); } +#[proptest(cases = 10)] +fn iter_count_test(#[strategy(0..250u8)] start: u8, #[strategy(#start..255u8)] end: u8) { + btree_test(|mut btree| { + for i in start..end { + assert_eq!(btree.insert(b(&[i]), b(&[])), None); + } + + for i in start..end { + for j in i..end { + assert_eq!(btree.range(b(&[i])..b(&[j])).count(), (j - i) as usize); + } + } + }); +} + // Given an operation, executes it on the given stable btreemap and standard btreemap, verifying // that the result of the operation is equal in both btrees. fn execute_operation(