From 4bf3e0d59749a0febbcb90c72a4179a285333043 Mon Sep 17 00:00:00 2001
From: Gwilym Kuiper <email@gwilym.dev>
Date: Sun, 2 Oct 2022 19:10:50 +0100
Subject: [PATCH] Add tests for size_hint and ensure .iter() returns a better
 iterator

---
 agb/src/hash_map.rs | 80 +++++++++++++++++++++++++++++++++------------
 1 file changed, 59 insertions(+), 21 deletions(-)

diff --git a/agb/src/hash_map.rs b/agb/src/hash_map.rs
index d38e60c2..d537b4cf 100644
--- a/agb/src/hash_map.rs
+++ b/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
@@ -340,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> {
@@ -355,13 +360,17 @@ 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(), Some(self.map.len()))
+        (
+            self.map.len() - self.num_found,
+            Some(self.map.len() - self.num_found),
+        )
     }
 }
 
@@ -370,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,
+        }
     }
 }
 
@@ -381,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> {
@@ -396,13 +410,17 @@ 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(), Some(self.map.len()))
+        (
+            self.map.len() - self.num_found,
+            Some(self.map.len() - self.num_found),
+        )
     }
 }
 
@@ -415,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,
+        }
     }
 }
 
@@ -1279,6 +1301,38 @@ mod test {
         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 {
@@ -1490,21 +1544,5 @@ 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
-        }
     }
 }