From 98372cc85f00c70e76eaf42a70c574ddd87eccfa Mon Sep 17 00:00:00 2001
From: kayomn <kayomn@kayomn.net>
Date: Sat, 15 Oct 2022 21:08:52 +0100
Subject: [PATCH] Improve memory safety of hash table

---
 src/table.zig | 36 ++++++++++++++++++++----------------
 1 file changed, 20 insertions(+), 16 deletions(-)

diff --git a/src/table.zig b/src/table.zig
index aecd12d..5bdaa81 100644
--- a/src/table.zig
+++ b/src/table.zig
@@ -56,16 +56,18 @@ pub fn Hashed(comptime Key: type, comptime Value: type,
         }
 
         ///
-        /// Searches for `key` to delete it, returning the deleted value or `null` if no matching
-        /// key was found.
+        /// Searches for `key` and deletes it from `self.
         ///
-        pub fn remove(self: Self, key: Key) ?Value {
+        /// The removed value is returned or `null` if no key matching `key` was found.
+        ///
+        pub fn remove(self: *Self, key: Key) ?Value {
             var bucket = &(self.buckets[@mod(key_context.hash(key), self.buckets.len)]);
 
             if (bucket.maybe_entry) |*entry| if (key_context.equals(entry.key, key)) {
-                defer entry.value = null;
-
-                self.filled -= 1;
+                defer {
+                    bucket.maybe_entry = null;
+                    self.filled -= 1;
+                }
 
                 return entry.value;
             };
@@ -74,9 +76,10 @@ pub fn Hashed(comptime Key: type, comptime Value: type,
                 bucket = &(self.buckets[index]);
 
                 if (bucket.maybe_entry) |*entry| if (key_context.equals(entry.key, key)) {
-                    defer entry.value = null;
-
-                    self.filled -= 1;
+                    defer {
+                        bucket.maybe_entry = null;
+                        self.filled -= 1;
+                    }
 
                     return entry.value;
                 };
@@ -129,20 +132,21 @@ pub fn Hashed(comptime Key: type, comptime Value: type,
         }
 
         ///
-        /// Searches for a value indexed with `key` in `self`, returning it or `null` if no matching
-        /// entry was found.
+        /// Searches for a value indexed with `key` in `self`.
         ///
-        pub fn lookup(self: Self, key: Key) ?*Value {
+        /// The found value is returned or `null` if an key matching `key` failed to be found.
+        ///
+        pub fn lookup(self: Self, key: Key) ?Value {
             var bucket = &(self.buckets[@mod(key_context.hash(key), self.buckets.len)]);
 
-            if (bucket.maybe_entry) |*entry|
-                if (key_context.equals(entry.key, key)) return &entry.value;
+            if (bucket.maybe_entry) |entry|
+                if (key_context.equals(entry.key, key)) return entry.value;
 
             while (bucket.maybe_next_index) |index| {
                 bucket = &(self.buckets[index]);
 
-                if (bucket.maybe_entry) |*entry|
-                    if (key_context.equals(entry.key, key)) return &entry.value;
+                if (bucket.maybe_entry) |entry|
+                    if (key_context.equals(entry.key, key)) return entry.value;
             }
 
             return null;