Merge pull request #318 from gwilymk/add-hashmap-retain-method

Add implementation of `HashMap.retain` method

CHANGELOG.md

@@ -12,9 +12,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Support for the blend mode of the GBA. Blending allows for alpha blending between layers and fading to black and white.
 - Added a new agb::sync module that contains GBA-specific synchronization primitives.
 - Added support for save files.
+- Added implementation of `HashMap.retain()`.
 
 ### Changes
 - Many of the places that originally disabled IRQs now use the `sync` module, reducing the chance of missed interrupts.
+- HashMap iterators now implement `size_hint` which should result in slightly better generation of code using those iterators.
 
 ### Fixed
 - Fixed the fast magnitude function in agb_fixnum. This is also used in fast_normalise. Previously only worked for positive (x, y).

agb/src/hash_map.rs

@@ -188,7 +188,11 @@ impl<K, V, ALLOCATOR: ClonableAllocator> HashMap<K, V, ALLOCATOR> {
 
     /// An iterator visiting all key-value pairs in an arbitrary order
     pub fn iter(&self) -> impl Iterator<Item = (&'_ K, &'_ V)> {
-        self.nodes.nodes.iter().filter_map(Node::key_value_ref)
+        Iter {
+            map: self,
+            at: 0,
+            num_found: 0,
+        }
     }
 
     /// An iterator visiting all key-value pairs in an arbitrary order, with mutable references to the values
@@ -196,6 +200,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 {
@@ -332,6 +344,7 @@ where
 pub struct Iter<'a, K: 'a, V: 'a, ALLOCATOR: ClonableAllocator> {
     map: &'a HashMap<K, V, ALLOCATOR>,
     at: usize,
+    num_found: usize,
 }
 
 impl<'a, K, V, ALLOCATOR: ClonableAllocator> Iterator for Iter<'a, K, V, ALLOCATOR> {
@@ -347,10 +360,18 @@ impl<'a, K, V, ALLOCATOR: ClonableAllocator> Iterator for Iter<'a, K, V, ALLOCAT
             self.at += 1;
 
             if node.has_value() {
+                self.num_found += 1;
                 return Some((node.key_ref().unwrap(), node.value_ref().unwrap()));
             }
         }
     }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        (
+            self.map.len() - self.num_found,
+            Some(self.map.len() - self.num_found),
+        )
+    }
 }
 
 impl<'a, K, V, ALLOCATOR: ClonableAllocator> IntoIterator for &'a HashMap<K, V, ALLOCATOR> {
@@ -358,7 +379,11 @@ impl<'a, K, V, ALLOCATOR: ClonableAllocator> IntoIterator for &'a HashMap<K, V,
     type IntoIter = Iter<'a, K, V, ALLOCATOR>;
 
     fn into_iter(self) -> Self::IntoIter {
-        Iter { map: self, at: 0 }
+        Iter {
+            map: self,
+            at: 0,
+            num_found: 0,
+        }
     }
 }
 
@@ -369,6 +394,7 @@ impl<'a, K, V, ALLOCATOR: ClonableAllocator> IntoIterator for &'a HashMap<K, V,
 pub struct IterOwned<K, V, ALLOCATOR: Allocator = Global> {
     map: HashMap<K, V, ALLOCATOR>,
     at: usize,
+    num_found: usize,
 }
 
 impl<K, V, ALLOCATOR: ClonableAllocator> Iterator for IterOwned<K, V, ALLOCATOR> {
@@ -384,10 +410,18 @@ impl<K, V, ALLOCATOR: ClonableAllocator> Iterator for IterOwned<K, V, ALLOCATOR>
             self.at += 1;
 
             if let Some((k, v, _)) = maybe_kv {
+                self.num_found += 1;
                 return Some((k, v));
             }
         }
     }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        (
+            self.map.len() - self.num_found,
+            Some(self.map.len() - self.num_found),
+        )
+    }
 }
 
 /// An iterator over entries of a [`HashMap`]
@@ -399,7 +433,11 @@ impl<K, V, ALLOCATOR: ClonableAllocator> IntoIterator for HashMap<K, V, ALLOCATO
     type IntoIter = IterOwned<K, V, ALLOCATOR>;
 
     fn into_iter(self) -> Self::IntoIter {
-        IterOwned { map: self, at: 0 }
+        IterOwned {
+            map: self,
+            at: 0,
+            num_found: 0,
+        }
     }
 }
 
@@ -729,6 +767,31 @@ 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);
+
+                    // Need to continue before adding 1 to i because remove from location could
+                    // put the element which was next into the ith location in the nodes array,
+                    // so we need to check if that one needs removing too.
+                    continue;
+                }
+            }
+
+            i += 1;
+        }
+    }
+
     fn remove_from_location(&mut self, location: usize) -> V {
         let mut current_location = location;
         self.number_of_items -= 1;
@@ -1222,6 +1285,54 @@ mod test {
         drop_registry.assert_dropped_n_times(id1, 2);
     }
 
+    #[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
+    }
+
+    #[test_case]
+    fn test_size_hint_iter(_gba: &mut Gba) {
+        let mut map = HashMap::new();
+
+        for i in 0..100 {
+            map.insert(i, i);
+        }
+
+        let mut iter = map.iter();
+        assert_eq!(iter.size_hint(), (100, Some(100)));
+
+        iter.next();
+
+        assert_eq!(iter.size_hint(), (99, Some(99)));
+    }
+
+    #[test_case]
+    fn test_size_hint_into_iter(_gba: &mut Gba) {
+        let mut map = HashMap::new();
+
+        for i in 0..100 {
+            map.insert(i, i);
+        }
+
+        let mut iter = map.into_iter();
+        assert_eq!(iter.size_hint(), (100, Some(100)));
+
+        iter.next();
+
+        assert_eq!(iter.size_hint(), (99, Some(99)));
+    }
+
     // Following test cases copied from the rust source
     // https://github.com/rust-lang/rust/blob/master/library/std/src/collections/hash/map/tests.rs
     mod rust_std_tests {