From 023d665fb386bdd69f052a8a1fe96b4d30fd918b Mon Sep 17 00:00:00 2001
From: Mikayla Maki <mikayla.c.maki@gmail.com>
Date: Wed, 3 May 2023 08:39:25 -0700
Subject: [PATCH] Fix TreeMap retain

---
 crates/sum_tree/src/tree_map.rs | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/crates/sum_tree/src/tree_map.rs b/crates/sum_tree/src/tree_map.rs
index 2580b08783..1b97cbec9f 100644
--- a/crates/sum_tree/src/tree_map.rs
+++ b/crates/sum_tree/src/tree_map.rs
@@ -103,11 +103,11 @@ impl<K: Clone + Debug + Default + Ord, V: Clone + Debug> TreeMap<K, V> {
     }
 
     pub fn retain<F: FnMut(&K, &V) -> bool>(&mut self, mut predicate: F) {
-        let mut cursor = self.0.cursor::<MapKeyRef<'_, K>>();
-        cursor.seek(&MapKeyRef(None), Bias::Left, &());
-
         let mut new_map = SumTree::<MapEntry<K, V>>::default();
-        if let Some(item) = cursor.item() {
+
+        let mut cursor = self.0.cursor::<MapKeyRef<'_, K>>();
+        cursor.next(&());
+        while let Some(item) = cursor.item() {
             if predicate(&item.key, &item.value) {
                 new_map.push(item.clone(), &());
             }
@@ -265,5 +265,11 @@ mod tests {
         map.remove(&1);
         assert_eq!(map.get(&1), None);
         assert_eq!(map.iter().collect::<Vec<_>>(), vec![]);
+
+        map.insert(4, "d");
+        map.insert(5, "e");
+        map.insert(6, "f");
+        map.retain(|key, _| *key % 2 == 0);
+        assert_eq!(map.iter().collect::<Vec<_>>(), vec![(&4, &"d"), (&6, &"f")]);
     }
 }