perf: BTreeMap: make range().count() more efficient

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: {}

src/btreemap/iter.rs

@@ -71,19 +71,13 @@ where
             range,
         }
     }
-}
-
-impl<K, V, M> 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::Item> {
-        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<T, F: Fn(&Node<K>, usize) -> T>(&mut self, map: F) -> Option<T> {
+        // 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<K, V, M> 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::Item> {
+        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::*;

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<M: Memory>(