From b6350a2b3f5c92ff4f0db152684f9c55bedb0bf6 Mon Sep 17 00:00:00 2001
From: Vexu <git@vexu.eu>
Date: Wed, 11 Nov 2020 00:01:40 +0200
Subject: [PATCH] std: fix HashMap.putAssumeCapacity

---
 lib/std/hash_map.zig | 69 ++++++++++++++++++++++----------------------
 1 file changed, 35 insertions(+), 34 deletions(-)

diff --git a/lib/std/hash_map.zig b/lib/std/hash_map.zig
index 458706d71f..8454bf8fc6 100644
--- a/lib/std/hash_map.zig
+++ b/lib/std/hash_map.zig
@@ -468,41 +468,12 @@ pub fn HashMapUnmanaged(
             self.putAssumeCapacityNoClobber(key, value);
         }
 
+        /// Asserts there is enough capacity to store the new key-value pair.
+        /// Clobbers any existing data. To detect if a put would clobber
+        /// existing data, see `getOrPutAssumeCapacity`.
         pub fn putAssumeCapacity(self: *Self, key: K, value: V) void {
-            const hash = hashFn(key);
-            const mask = self.capacity() - 1;
-            const fingerprint = Metadata.takeFingerprint(hash);
-            var idx = @truncate(usize, hash & mask);
-
-            var first_tombstone_idx: usize = self.capacity(); // invalid index
-            var metadata = self.metadata.? + idx;
-            while (metadata[0].isUsed() or metadata[0].isTombstone()) {
-                if (metadata[0].isUsed() and metadata[0].fingerprint == fingerprint) {
-                    const entry = &self.entries()[idx];
-                    if (eqlFn(entry.key, key)) {
-                        return;
-                    }
-                } else if (first_tombstone_idx == self.capacity() and metadata[0].isTombstone()) {
-                    first_tombstone_idx = idx;
-                }
-
-                idx = (idx + 1) & mask;
-                metadata = self.metadata.? + idx;
-            }
-
-            if (first_tombstone_idx < self.capacity()) {
-                // Cheap try to lower probing lengths after deletions. Recycle a tombstone.
-                idx = first_tombstone_idx;
-                metadata = self.metadata.? + idx;
-            } else {
-                // We're using a slot previously free.
-                self.available -= 1;
-            }
-
-            metadata[0].fill(fingerprint);
-            const entry = &self.entries()[idx];
-            entry.* = .{ .key = key, .value = undefined };
-            self.size += 1;
+            const gop = self.getOrPutAssumeCapacity(key);
+            gop.entry.value = value;
         }
 
         /// Insert an entry in the map. Assumes it is not already present,
@@ -1148,6 +1119,36 @@ test "std.hash_map put" {
     }
 }
 
+test "std.hash_map putAssumeCapacity" {
+    var map = AutoHashMap(u32, u32).init(std.testing.allocator);
+    defer map.deinit();
+    
+    try map.ensureCapacity(20);
+    var i: u32 = 0;
+    while (i < 20) : (i += 1) {
+        map.putAssumeCapacityNoClobber(i, i);
+    }
+
+    i = 0;
+    var sum = i;
+    while (i < 20) : (i += 1) {
+        sum += map.get(i).?;
+    }
+    expectEqual(sum, 190);
+
+    i = 0;
+    while (i < 20) : (i += 1) {
+        map.putAssumeCapacity(i, 1);
+    }
+
+    i = 0;
+    sum = i;
+    while (i < 20) : (i += 1) {
+        sum += map.get(i).?;
+    }
+    expectEqual(sum, 20);
+}
+
 test "std.hash_map getOrPut" {
     var map = AutoHashMap(u32, u32).init(std.testing.allocator);
     defer map.deinit();