perf: improve contains_key performance

benchmarks/src/btreemap.rs

@@ -707,6 +707,54 @@ fn get_helper<K: Clone + Ord + Storable + Random, V: Storable + Random>(
     })
 }
 
+/// Benchmarks `contains_key` of a BTreeMap.
+#[bench(raw)]
+pub fn btreemap_contains_key_blob_4_1024() -> BenchResult {
+    contains_key_blob_helper::<4, 1024>()
+}
+
+#[bench(raw)]
+pub fn btreemap_contains_key_blob_4_1024_v2() -> BenchResult {
+    contains_key_blob_helper_v2::<4, 1024>()
+}
+
+// Profiles `contains_key` on a large number of random blobs from a btreemap.
+fn contains_key_blob_helper<const K: usize, const V: usize>() -> BenchResult {
+    let btree = BTreeMap::new_v1(DefaultMemoryImpl::default());
+    contains_key_helper::<Blob<K>, Blob<V>>(btree)
+}
+
+fn contains_key_blob_helper_v2<const K: usize, const V: usize>() -> BenchResult {
+    let btree = BTreeMap::new(DefaultMemoryImpl::default());
+    contains_key_helper::<Blob<K>, Blob<V>>(btree)
+}
+
+fn contains_key_helper<K: Clone + Ord + Storable + Random, V: Storable + Random>(
+    mut btree: BTreeMap<K, V, impl Memory>,
+) -> BenchResult {
+    let num_keys = 10_000;
+    let mut rng = Rng::from_seed(0);
+    let mut random_keys = Vec::with_capacity(num_keys);
+    let mut random_values = Vec::with_capacity(num_keys);
+
+    for _ in 0..num_keys {
+        random_keys.push(K::random(&mut rng));
+        random_values.push(V::random(&mut rng));
+    }
+
+    // Insert the keys into the btree.
+    for (k, v) in random_keys.iter().zip(random_values.into_iter()) {
+        btree.insert(k.clone(), v);
+    }
+
+    // Checks if the keys are in the map.
+    bench_fn(|| {
+        for k in random_keys.into_iter() {
+            btree.contains_key(&k);
+        }
+    })
+}
+
 // Inserts a large number of random blobs into a btreemap, then profiles removing them.
 fn remove_blob_helper<const K: usize, const V: usize>() -> BenchResult {
     let btree = BTreeMap::new_v1(DefaultMemoryImpl::default());

canbench_results.yml

@@ -1,145 +1,157 @@
 benches:
+  btreemap_contains_key_blob_4_1024:
+    total:
+      instructions: 166001914
+      heap_increase: 0
+      stable_memory_increase: 0
+    scopes: {}
+  btreemap_contains_key_blob_4_1024_v2:
+    total:
+      instructions: 246601950
+      heap_increase: 0
+      stable_memory_increase: 0
+    scopes: {}
   btreemap_get_blob_128_1024:
     total:
-      instructions: 871309202
+      instructions: 871377876
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   btreemap_get_blob_128_1024_v2:
     total:
-      instructions: 952323558
+      instructions: 952392232
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   btreemap_get_blob_16_1024:
     total:
-      instructions: 246794283
+      instructions: 246862999
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   btreemap_get_blob_16_1024_v2:
     total:
-      instructions: 322817530
+      instructions: 322886246
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   btreemap_get_blob_256_1024:
     total:
-      instructions: 1438305373
+      instructions: 1438374032
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   btreemap_get_blob_256_1024_v2:
     total:
-      instructions: 1522441414
+      instructions: 1522510073
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   btreemap_get_blob_32_1024:
     total:
-      instructions: 280038069
+      instructions: 280106753
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   btreemap_get_blob_32_1024_v2:
     total:
-      instructions: 357990649
+      instructions: 358059333
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   btreemap_get_blob_4_1024:
     total:
-      instructions: 183602777
+      instructions: 183878301
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   btreemap_get_blob_4_1024_v2:
     total:
-      instructions: 268129985
+      instructions: 268405509
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   btreemap_get_blob_512_1024:
     total:
-      instructions: 2573953601
+      instructions: 2574022258
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   btreemap_get_blob_512_1024_v2:
     total:
-      instructions: 2652827669
+      instructions: 2652896326
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   btreemap_get_blob_512_1024_v2_mem_manager:
     total:
-      instructions: 2758867134
+      instructions: 2758935791
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   btreemap_get_blob_64_1024:
     total:
-      instructions: 522384319
+      instructions: 522452995
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   btreemap_get_blob_64_1024_v2:
     total:
-      instructions: 602402025
+      instructions: 602470701
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   btreemap_get_blob_8_1024:
     total:
-      instructions: 217332289
+      instructions: 217400958
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   btreemap_get_blob_8_1024_v2:
     total:
-      instructions: 299640561
+      instructions: 299709230
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   btreemap_get_blob_8_u64:
     total:
-      instructions: 201109665
+      instructions: 201178258
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   btreemap_get_blob_8_u64_v2:
     total:
-      instructions: 297999015
+      instructions: 298067608
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   btreemap_get_u64_blob_8:
     total:
-      instructions: 175387736
+      instructions: 175651041
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   btreemap_get_u64_blob_8_v2:
     total:
-      instructions: 249704264
+      instructions: 249967569
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   btreemap_get_u64_u64:
     total:
-      instructions: 176095721
+      instructions: 176359016
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   btreemap_get_u64_u64_v2:
     total:
-      instructions: 256392975
+      instructions: 256656270
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   btreemap_get_u64_u64_v2_mem_manager:
     total:
-      instructions: 335916653
+      instructions: 336179948
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
@@ -731,4 +743,4 @@ benches:
       heap_increase: 0
       stable_memory_increase: 1
     scopes: {}
-version: 0.1.9
+version: 0.1.11

src/btreemap.rs

@@ -509,38 +509,40 @@ where
         node.save(self.allocator_mut());
     }
 
-    /// Returns the value associated with the given key if it exists.
+    /// Returns the value for the given key, if it exists.
     pub fn get(&self, key: &K) -> Option<V> {
         if self.root_addr == NULL {
             return None;
         }
+        self.traverse(self.root_addr, key, |node, idx| {
+            node.into_entry(idx, self.memory()).1 // Extract value.
+        })
+        .map(Cow::Owned)
+        .map(V::from_bytes)
+    }
 
-        self.get_helper(self.root_addr, key)
-            .map(Cow::Owned)
-            .map(V::from_bytes)
+    /// Returns true if the key exists.
+    pub fn contains_key(&self, key: &K) -> bool {
+        // An empty closure returns Some(()) if the key is found.
+        self.root_addr != NULL && self.traverse(self.root_addr, key, |_, _| ()).is_some()
     }
 
-    fn get_helper(&self, node_addr: Address, key: &K) -> Option<Vec<u8>> {
+    /// Recursively traverses from `node_addr`, invoking `f` if `key` is found. Stops at a leaf if not.
+    fn traverse<F, R>(&self, node_addr: Address, key: &K, f: F) -> Option<R>
+    where
+        F: Fn(Node<K>, usize) -> R + Clone,
+    {
         let node = self.load_node(node_addr);
+        // Look for the key in the current node.
         match node.search(key) {
-            Ok(idx) => Some(node.into_entry(idx, self.memory()).1),
-            Err(idx) => {
-                match node.node_type() {
-                    NodeType::Leaf => None, // Key not found.
-                    NodeType::Internal => {
-                        // The key isn't in the node. Look for the key in the child.
-                        self.get_helper(node.child(idx), key)
-                    }
-                }
-            }
+            Ok(idx) => Some(f(node, idx)), // Key found: apply `f`.
+            Err(idx) => match node.node_type() {
+                NodeType::Leaf => None, // At a leaf: key not present.
+                NodeType::Internal => self.traverse(node.child(idx), key, f), // Continue search in child.
+            },
         }
     }
 
-    /// Returns `true` if the key exists in the map, `false` otherwise.
-    pub fn contains_key(&self, key: &K) -> bool {
-        self.get(key).is_some()
-    }
-
     /// Returns `true` if the map contains no elements.
     pub fn is_empty(&self) -> bool {
         self.length == 0