Commit 11803a3a569205d640c7ec0b0aedba83f47a6e64 | zig

author	sentientwaffle <sentientwaffle@gmail.com> 2021-12-17 11:38:13-08:00
committer	Andrew Kelley <andrew@ziglang.org> 2021-12-17 15:21:41-08:00
log	11803a3a569205d640c7ec0b0aedba83f47a6e64
tree	5c45e90777f3993a162dbfbd9adb368df6ba75fc
parent	e8b39960bbd88d73cfd0d853e8ca7383e300c58c

std: optimize hash_map probe loop condition

See https://github.com/ziglang/zig/pull/10337 for context.

In #10337 the `available` tracking fix necessitated an additional condition on the probe loop in both `getOrPut` and `getIndex` to prevent an infinite loop. Previously, this condition was implicit thanks to the guaranteed presence of a free slot.

The new condition hurts the `HashMap` benchmarks (https://github.com/ziglang/zig/pull/10337#issuecomment-996432758).

This commit removes that extra condition on the loop. Instead, when probing, first check whether the "home" slot is the target key — if so, return it. Otherwise, save the home slot's metadata to the stack and temporarily "free" the slot (but don't touch its value). Then continue with the original loop. Once again, the loop will be implicitly broken by the new "free" slot. The original metadata is restored before the function returns.

`getOrPut` has one additional gotcha — if the home slot is a tombstone and `getOrPut` misses, then the home slot is is written with the new key; that is, its original metadata (the tombstone) is not restored.

Other changes:

- Test hash map misses.
- Test using `getOrPutAssumeCapacity` to get keys at the end (along with `get`).

1 files changed, 60 insertions(+), 13 deletions(-)

lib
- std
  - hash_map.zig

lib/std/hash_map.zig+60-13

 @@ -773,6 +773,11 @@ pub fn HashMapUnmanaged(
                 self.used = 0;
                 self.fingerprint = tombstone;
             }
+            pub fn reset(self: *Metadata) void {
+                self.used = 0;
+                self.fingerprint = free;
+            }
         };
         comptime {
 @@ -1110,12 +1115,24 @@ pub fn HashMapUnmanaged(
             }
             const mask = self.capacity() - 1;
             const fingerprint = Metadata.takeFingerprint(hash);
-            // Don't loop indefinitely when there are no empty slots.
+            var start = @truncate(usize, hash & mask);
-            var limit = self.capacity();
+            var idx = start;
-            var idx = @truncate(usize, hash & mask);
             var metadata = self.metadata.? + idx;
-            while ((metadata[0].isUsed() or metadata[0].isTombstone()) and limit != 0) {
+            if (metadata[0].isUsed() and metadata[0].fingerprint == fingerprint) {
+                const test_key = &self.keys()[idx];
+                const eql = ctx.eql(key, test_key.*);
+                if (eql) return idx;
+            }
+            // Temporarily mark the start position as "free" so that the probe loop
+            // doesn't infinitely loop when all other slots are filled or tombstones.
+            const saved = metadata[0];
+            metadata[0].reset();
+            defer self.metadata.?[start] = saved;
+            idx = (idx + 1) & mask; // initial idx was already checked
+            metadata = self.metadata.? + idx;
+            while (metadata[0].isUsed() or metadata[0].isTombstone()) {
                 if (metadata[0].isUsed() and metadata[0].fingerprint == fingerprint) {
                     const test_key = &self.keys()[idx];
                     // If you get a compile error on this line, it means that your generic eql
 @@ -1131,7 +1148,6 @@ pub fn HashMapUnmanaged(
                     }
                 }
-                limit -= 1;
                 idx = (idx + 1) & mask;
                 metadata = self.metadata.? + idx;
             }
 @@ -1289,12 +1305,39 @@ pub fn HashMapUnmanaged(
             }
             const mask = self.capacity() - 1;
             const fingerprint = Metadata.takeFingerprint(hash);
-            var limit = self.capacity();
+            var start = @truncate(usize, hash & mask);
-            var idx = @truncate(usize, hash & mask);
+            var idx = start;
             var first_tombstone_idx: usize = self.capacity(); // invalid index
             var metadata = self.metadata.? + idx;
-            while ((metadata[0].isUsed() or metadata[0].isTombstone()) and limit != 0) {
+            if (metadata[0].isUsed() and metadata[0].fingerprint == fingerprint) {
+                const test_key = &self.keys()[idx];
+                const eql = ctx.eql(key, test_key.*);
+                if (eql) {
+                    return GetOrPutResult{
+                        .key_ptr = test_key,
+                        .value_ptr = &self.values()[idx],
+                        .found_existing = true,
+                    };
+                }
+            } else if (metadata[0].isTombstone()) {
+                first_tombstone_idx = idx;
+            }
+            // Temporarily mark the start position as "free" so that the probe loop
+            // doesn't infinitely loop when all other slots are filled or tombstones.
+            const saved = metadata[0];
+            metadata[0].reset();
+            defer {
+                // Don't restore when the 'home' slot was originally a tombstone
+                // and the probe missed, since it is now occupied by the new key.
+                const home = &self.metadata.?[start];
+                assert(!home.isUsed() or home.fingerprint == fingerprint);
+                if (!home.isUsed()) home.* = saved;
+            }
+            idx = (idx + 1) & mask; // initial idx was already checked
+            metadata = self.metadata.? + idx;
+            while (metadata[0].isUsed() or metadata[0].isTombstone()) {
                 if (metadata[0].isUsed() and metadata[0].fingerprint == fingerprint) {
                     const test_key = &self.keys()[idx];
                     // If you get a compile error on this line, it means that your generic eql
 @@ -1316,7 +1359,6 @@ pub fn HashMapUnmanaged(
                     first_tombstone_idx = idx;
                 }
-                limit -= 1;
                 idx = (idx + 1) & mask;
                 metadata = self.metadata.? + idx;
             }
 @@ -1902,6 +1944,7 @@ test "std.hash_map repeat putAssumeCapacity/remove" {
     try map.ensureTotalCapacity(20);
     const limit = map.unmanaged.available;
+    const cycles = 100;
     var i: u32 = 0;
     while (i < limit) : (i += 1) {
 @@ -1911,7 +1954,7 @@ test "std.hash_map repeat putAssumeCapacity/remove" {
     // Repeatedly delete/insert an entry without resizing the map.
     // Put to different keys so entries don't land in the just-freed slot.
     i = 0;
-    while (i < 10 * limit) : (i += 1) {
+    while (i < cycles * limit) : (i += 1) {
         try testing.expect(map.remove(i));
         if (i % 2 == 0) {
             map.putAssumeCapacityNoClobber(limit + i, i);
 @@ -1920,9 +1963,13 @@ test "std.hash_map repeat putAssumeCapacity/remove" {
         }
     }
-    i = 9 * limit;
+    i = (cycles - 1) * limit;
-    while (i < 10 * limit) : (i += 1) {
+    while (i < cycles * limit) : (i += 1) {
-        try expectEqual(map.get(limit + i), i);
+        try expectEqual(map.get(i), null); // (removed) key miss
+        try expectEqual(map.get(limit + i), i); // key hit
+        const gop = map.getOrPutAssumeCapacity(limit + i);
+        try testing.expect(gop.found_existing);
+        try testing.expectEqual(gop.value_ptr.*, i);
     }
     try expectEqual(map.unmanaged.available, 0);
     try expectEqual(map.unmanaged.count(), limit);