From a9115c23e859839b02f04ba393865cb2d110d297 Mon Sep 17 00:00:00 2001
From: Gwilym Kuiper <gw@ilym.me>
Date: Fri, 18 Mar 2022 00:18:39 +0000
Subject: [PATCH] Allow for resizing

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

diff --git a/agb/src/hash_map.rs b/agb/src/hash_map.rs
index 47ecb487..87f36d0e 100644
--- a/agb/src/hash_map.rs
+++ b/agb/src/hash_map.rs
@@ -37,9 +37,9 @@ struct NodeStorage<K, V>(Vec<Option<Node<K, V>>>);
 
 impl<K, V> NodeStorage<K, V> {
     fn with_size(capacity: usize) -> Self {
-        let next_power_of_2 = find_next_power_of_2(capacity);
+        assert!(capacity.is_power_of_two(), "Capacity must be a power of 2");
 
-        Self(iter::repeat_with(|| None).take(next_power_of_2).collect())
+        Self(iter::repeat_with(|| None).take(capacity).collect())
     }
 
     fn len(&self) -> usize {
@@ -140,6 +140,35 @@ impl<K, V> HashMap<K, V> {
     pub fn len(&self) -> usize {
         self.number_of_elements
     }
+
+    pub fn resize(&mut self, new_size: usize) {
+        assert!(
+            new_size >= self.nodes.len(),
+            "Can only increase the size of a hash map"
+        );
+        if new_size == self.nodes.len() {
+            return;
+        }
+
+        let mut new_node_storage = NodeStorage::with_size(new_size);
+        let mut new_max_distance_to_initial_bucket = 0;
+        let number_of_elements = self.number_of_elements;
+
+        for node in self.nodes.0.drain(..) {
+            if let Some(node) = node {
+                new_max_distance_to_initial_bucket = new_node_storage.insert_new(
+                    node.key,
+                    node.value,
+                    node.hash,
+                    number_of_elements,
+                    new_max_distance_to_initial_bucket,
+                );
+            }
+        }
+
+        self.nodes = new_node_storage;
+        self.max_distance_to_initial_bucket = new_max_distance_to_initial_bucket;
+    }
 }
 
 const fn fast_mod(len: usize, hash: HashType) -> usize {
@@ -147,15 +176,6 @@ const fn fast_mod(len: usize, hash: HashType) -> usize {
     (hash as usize) & (len - 1)
 }
 
-const fn find_next_power_of_2(n: usize) -> usize {
-    let mut next_power_of_2 = 1;
-    while next_power_of_2 <= n {
-        next_power_of_2 *= 2;
-    }
-
-    next_power_of_2
-}
-
 impl<K, V> HashMap<K, V>
 where
     K: Eq,
@@ -193,6 +213,10 @@ where
             return Some(old_value);
         }
 
+        if self.nodes.len() * 85 / 100 <= self.number_of_elements {
+            self.resize(self.nodes.len() * 2);
+        }
+
         self.max_distance_to_initial_bucket = self.nodes.insert_new(
             key,
             value,
@@ -345,4 +369,17 @@ mod test {
         assert_eq!(num_found, 8);
         assert_eq!(max_found, 7);
     }
+
+    #[test_case]
+    fn can_insert_more_than_initial_capacity(_gba: &mut Gba) {
+        let mut map = HashMap::new();
+
+        for i in 0..65 {
+            map.put(i, i % 4);
+        }
+
+        for i in 0..65 {
+            assert_eq!(map.get(&i), Some(&(i % 4)));
+        }
+    }
 }