From 7873f9ed5c07a8411725f575c31be84506e35b4f Mon Sep 17 00:00:00 2001
From: Gwilym Inzani <gw@ilym.me>
Date: Tue, 9 May 2023 22:43:06 +0100
Subject: [PATCH] More unsafe where needed

---
 agb-hashmap/src/lib.rs          | 21 ++++++++++++++++++---
 agb-hashmap/src/node.rs         | 24 ++++--------------------
 agb-hashmap/src/node_storage.rs | 11 +++++++++--
 3 files changed, 31 insertions(+), 25 deletions(-)

diff --git a/agb-hashmap/src/lib.rs b/agb-hashmap/src/lib.rs
index 4b4afbab..95554fb8 100644
--- a/agb-hashmap/src/lib.rs
+++ b/agb-hashmap/src/lib.rs
@@ -314,7 +314,13 @@ where
         let hash = self.hash(&key);
 
         if let Some(location) = self.nodes.location(&key, hash) {
-            Some(self.nodes.replace_at_location(location, key, value))
+            Some(
+                // SAFETY: location is valid due to the above
+                unsafe {
+                    self.nodes
+                        .replace_at_location_unchecked(location, key, value)
+                },
+            )
         } else {
             if self.nodes.capacity() <= self.len() {
                 self.resize(self.nodes.backing_vec_size() * 2);
@@ -330,7 +336,11 @@ where
         let hash = self.hash(&key);
 
         let location = if let Some(location) = self.nodes.location(&key, hash) {
-            self.nodes.replace_at_location(location, key, value);
+            // SAFETY: location is valid due to the above
+            unsafe {
+                self.nodes
+                    .replace_at_location_unchecked(location, key, value);
+            }
             location
         } else {
             if self.nodes.capacity() <= self.len() {
@@ -340,7 +350,12 @@ where
             self.nodes.insert_new(key, value, hash)
         };
 
-        self.nodes.node_at_mut(location).value_mut().unwrap()
+        // SAFETY: location is always valid
+        unsafe {
+            self.nodes
+                .node_at_unchecked_mut(location)
+                .value_mut_unchecked()
+        }
     }
 
     /// Returns `true` if the map contains a value for the specified key.
diff --git a/agb-hashmap/src/node.rs b/agb-hashmap/src/node.rs
index d9101def..0cbe22f5 100644
--- a/agb-hashmap/src/node.rs
+++ b/agb-hashmap/src/node.rs
@@ -38,17 +38,6 @@ impl<K, V> Node<K, V> {
         self.value.assume_init_ref()
     }
 
-    pub(crate) fn value_mut(&mut self) -> Option<&mut V> {
-        if self.has_value() {
-            Some(
-                // SAFETY: has a value
-                unsafe { self.value_mut_unchecked() },
-            )
-        } else {
-            None
-        }
-    }
-
     pub(crate) unsafe fn value_mut_unchecked(&mut self) -> &mut V {
         self.value.assume_init_mut()
     }
@@ -114,16 +103,11 @@ impl<K, V> Node<K, V> {
         old_value.assume_init()
     }
 
-    pub(crate) fn replace(&mut self, key: K, value: V) -> (K, V) {
-        if self.has_value() {
-            let old_key = mem::replace(&mut self.key, MaybeUninit::new(key));
-            let old_value = mem::replace(&mut self.value, MaybeUninit::new(value));
+    pub(crate) unsafe fn replace_unchecked(&mut self, key: K, value: V) -> (K, V) {
+        let old_key = mem::replace(&mut self.key, MaybeUninit::new(key));
+        let old_value = mem::replace(&mut self.value, MaybeUninit::new(value));
 
-            // SAFETY: has a value
-            unsafe { (old_key.assume_init(), old_value.assume_init()) }
-        } else {
-            panic!("Cannot replace an uninitialised node");
-        }
+        (old_key.assume_init(), old_value.assume_init())
     }
 
     pub(crate) fn increment_distance(&mut self) {
diff --git a/agb-hashmap/src/node_storage.rs b/agb-hashmap/src/node_storage.rs
index 43bc7d79..55a7d276 100644
--- a/agb-hashmap/src/node_storage.rs
+++ b/agb-hashmap/src/node_storage.rs
@@ -164,8 +164,15 @@ impl<K, V, ALLOCATOR: ClonableAllocator> NodeStorage<K, V, ALLOCATOR> {
         new_node_storage
     }
 
-    pub(crate) fn replace_at_location(&mut self, location: usize, key: K, value: V) -> V {
-        self.nodes[location].replace(key, value).1
+    pub(crate) unsafe fn replace_at_location_unchecked(
+        &mut self,
+        location: usize,
+        key: K,
+        value: V,
+    ) -> V {
+        self.node_at_unchecked_mut(location)
+            .replace_unchecked(key, value)
+            .1
     }
 
     pub(crate) fn iter_mut(&mut self) -> impl Iterator<Item = &mut Node<K, V>> {