perf: optimize value swap in BTreeMap

src/btreemap.rs

@@ -506,10 +506,10 @@ where
 
             // Check if the key already exists in the root.
             if let Ok(idx) = root.search(&key, self.memory()) {
-                // The key exists. Overwrite it and return the previous value.
-                let (_, previous_value) = root.swap_entry(idx, (key, value), self.memory());
-                self.save_node(&mut root);
-                return Some(V::from_bytes(Cow::Owned(previous_value)));
+                // Key found, replace its value and return the old one.
+                return Some(V::from_bytes(Cow::Owned(
+                    self.update_value(&mut root, idx, value),
+                )));
             }
 
             // If the root is full, we need to introduce a new node as the root.
@@ -550,12 +550,8 @@ where
         // Look for the key in the node.
         match node.search(&key, self.memory()) {
             Ok(idx) => {
-                // The key is already in the node.
-                // Overwrite it and return the previous value.
-                let (_, previous_value) = node.swap_entry(idx, (key, value), self.memory());
-
-                self.save_node(&mut node);
-                Some(previous_value)
+                // Key found, replace its value and return the old one.
+                Some(self.update_value(&mut node, idx, value))
             }
             Err(idx) => {
                 // The key isn't in the node. `idx` is where that key should be inserted.
@@ -582,11 +578,8 @@ where
                         if child.is_full() {
                             // Check if the key already exists in the child.
                             if let Ok(idx) = child.search(&key, self.memory()) {
-                                // The key exists. Overwrite it and return the previous value.
-                                let (_, previous_value) =
-                                    child.swap_entry(idx, (key, value), self.memory());
-                                self.save_node(&mut child);
-                                return Some(previous_value);
+                                // Key found, replace its value and return the old one.
+                                return Some(self.update_value(&mut child, idx, value));
                             }
 
                             // The child is full. Split the child.
@@ -1281,6 +1274,13 @@ where
         node.save(self.allocator_mut());
     }
 
+    /// Replaces the value at `idx` in the node, saves the node, and returns the old value.
+    fn update_value(&mut self, node: &mut Node<K>, idx: usize, new_value: Vec<u8>) -> Vec<u8> {
+        let old_value = node.swap_value(idx, new_value, self.memory());
+        self.save_node(node);
+        old_value
+    }
+
     /// Saves the map to memory.
     fn save_header(&self) {
         let header = BTreeHeader {

src/btreemap/node.rs

@@ -156,6 +156,15 @@ impl<K: Storable + Ord + Clone> Node<K> {
         self.keys_and_encoded_values.len() >= CAPACITY
     }
 
+    /// Replaces the value at `idx` and returns the old one.
+    pub fn swap_value<M: Memory>(&mut self, idx: usize, new: Vec<u8>, memory: &M) -> Vec<u8> {
+        let old = core::mem::replace(
+            &mut self.keys_and_encoded_values[idx].1,
+            LazyValue::by_value(new),
+        );
+        self.extract_value(old, memory)
+    }
+
     /// Swaps the entry at index `idx` with the given entry, returning the old entry.
     pub fn swap_entry<M: Memory>(
         &mut self,