From c658b6afe5708ff17a42eb2c3c347bb6e0ce0dd6 Mon Sep 17 00:00:00 2001
From: Gwilym Kuiper <email@gwilym.dev>
Date: Sun, 2 Oct 2022 18:55:21 +0100
Subject: [PATCH] Add implementation of retain to HashMap

---
 agb/src/hash_map.rs | 45 +++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 45 insertions(+)

diff --git a/agb/src/hash_map.rs b/agb/src/hash_map.rs
index 11cf4519..9a73c4a8 100644
--- a/agb/src/hash_map.rs
+++ b/agb/src/hash_map.rs
@@ -196,6 +196,14 @@ impl<K, V, ALLOCATOR: ClonableAllocator> HashMap<K, V, ALLOCATOR> {
         self.nodes.nodes.iter_mut().filter_map(Node::key_value_mut)
     }
 
+    /// Retains only the elements specified by the predicate `f`.
+    pub fn retain<F>(&mut self, f: F)
+    where
+        F: FnMut(&K, &mut V) -> bool,
+    {
+        self.nodes.retain(f);
+    }
+
     /// Returns `true` if the map contains no elements
     #[must_use]
     pub fn is_empty(&self) -> bool {
@@ -737,6 +745,27 @@ impl<K, V, ALLOCATOR: ClonableAllocator> NodeStorage<K, V, ALLOCATOR> {
         inserted_location
     }
 
+    fn retain<F>(&mut self, mut f: F)
+    where
+        F: FnMut(&K, &mut V) -> bool,
+    {
+        let num_nodes = self.nodes.len();
+        let mut i = 0;
+
+        while i < num_nodes {
+            let node = &mut self.nodes[i];
+
+            if let Some((k, v)) = node.key_value_mut() {
+                if !f(k, v) {
+                    self.remove_from_location(i);
+                    continue;
+                }
+            }
+
+            i += 1;
+        }
+    }
+
     fn remove_from_location(&mut self, location: usize) -> V {
         let mut current_location = location;
         self.number_of_items -= 1;
@@ -1441,5 +1470,21 @@ mod test {
 
             assert_eq!(map[&2], 1);
         }
+
+        #[test_case]
+        fn test_retain(_gba: &mut Gba) {
+            let mut map = HashMap::new();
+
+            for i in 0..100 {
+                map.insert(i, i);
+            }
+
+            map.retain(|k, _| k % 2 == 0);
+
+            assert_eq!(map[&2], 2);
+            assert_eq!(map.get(&3), None);
+
+            assert_eq!(map.iter().count(), 50); // force full iteration
+        }
     }
 }