authorgravatar for sentientwaffle@gmail.comsentientwaffle <sentientwaffle@gmail.com> 2021-12-17 11:38:13-08:00
committergravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2021-12-17 15:21:41-08:00
log11803a3a569205d640c7ec0b0aedba83f47a6e64
tree5c45e90777f3993a162dbfbd9adb368df6ba75fc
parente8b39960bbd88d73cfd0d853e8ca7383e300c58c

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+60-13
......@@ -773,6 +773,11 @@ pub fn HashMapUnmanaged(
773773 self.used = 0;
774774 self.fingerprint = tombstone;
775775 }
776
777 pub fn reset(self: *Metadata) void {
778 self.used = 0;
779 self.fingerprint = free;
780 }
776781 };
777782
778783 comptime {
......@@ -1110,12 +1115,24 @@ pub fn HashMapUnmanaged(
11101115 }
11111116 const mask = self.capacity() - 1;
11121117 const fingerprint = Metadata.takeFingerprint(hash);
1113 // Don't loop indefinitely when there are no empty slots.
1114 var limit = self.capacity();
1115 var idx = @truncate(usize, hash & mask);
1118 var start = @truncate(usize, hash & mask);
1119 var idx = start;
11161120
11171121 var metadata = self.metadata.? + idx;
1118 while ((metadata[0].isUsed() or metadata[0].isTombstone()) and limit != 0) {
1122 if (metadata[0].isUsed() and metadata[0].fingerprint == fingerprint) {
1123 const test_key = &self.keys()[idx];
1124 const eql = ctx.eql(key, test_key.*);
1125 if (eql) return idx;
1126 }
1127 // Temporarily mark the start position as "free" so that the probe loop
1128 // doesn't infinitely loop when all other slots are filled or tombstones.
1129 const saved = metadata[0];
1130 metadata[0].reset();
1131 defer self.metadata.?[start] = saved;
1132 idx = (idx + 1) & mask; // initial idx was already checked
1133 metadata = self.metadata.? + idx;
1134
1135 while (metadata[0].isUsed() or metadata[0].isTombstone()) {
11191136 if (metadata[0].isUsed() and metadata[0].fingerprint == fingerprint) {
11201137 const test_key = &self.keys()[idx];
11211138 // If you get a compile error on this line, it means that your generic eql
......@@ -1131,7 +1148,6 @@ pub fn HashMapUnmanaged(
11311148 }
11321149 }
11331150
1134 limit -= 1;
11351151 idx = (idx + 1) & mask;
11361152 metadata = self.metadata.? + idx;
11371153 }
......@@ -1289,12 +1305,39 @@ pub fn HashMapUnmanaged(
12891305 }
12901306 const mask = self.capacity() - 1;
12911307 const fingerprint = Metadata.takeFingerprint(hash);
1292 var limit = self.capacity();
1293 var idx = @truncate(usize, hash & mask);
1308 var start = @truncate(usize, hash & mask);
1309 var idx = start;
12941310
12951311 var first_tombstone_idx: usize = self.capacity(); // invalid index
12961312 var metadata = self.metadata.? + idx;
1297 while ((metadata[0].isUsed() or metadata[0].isTombstone()) and limit != 0) {
1313 if (metadata[0].isUsed() and metadata[0].fingerprint == fingerprint) {
1314 const test_key = &self.keys()[idx];
1315 const eql = ctx.eql(key, test_key.*);
1316 if (eql) {
1317 return GetOrPutResult{
1318 .key_ptr = test_key,
1319 .value_ptr = &self.values()[idx],
1320 .found_existing = true,
1321 };
1322 }
1323 } else if (metadata[0].isTombstone()) {
1324 first_tombstone_idx = idx;
1325 }
1326 // Temporarily mark the start position as "free" so that the probe loop
1327 // doesn't infinitely loop when all other slots are filled or tombstones.
1328 const saved = metadata[0];
1329 metadata[0].reset();
1330 defer {
1331 // Don't restore when the 'home' slot was originally a tombstone
1332 // and the probe missed, since it is now occupied by the new key.
1333 const home = &self.metadata.?[start];
1334 assert(!home.isUsed() or home.fingerprint == fingerprint);
1335 if (!home.isUsed()) home.* = saved;
1336 }
1337 idx = (idx + 1) & mask; // initial idx was already checked
1338 metadata = self.metadata.? + idx;
1339
1340 while (metadata[0].isUsed() or metadata[0].isTombstone()) {
12981341 if (metadata[0].isUsed() and metadata[0].fingerprint == fingerprint) {
12991342 const test_key = &self.keys()[idx];
13001343 // If you get a compile error on this line, it means that your generic eql
......@@ -1316,7 +1359,6 @@ pub fn HashMapUnmanaged(
13161359 first_tombstone_idx = idx;
13171360 }
13181361
1319 limit -= 1;
13201362 idx = (idx + 1) & mask;
13211363 metadata = self.metadata.? + idx;
13221364 }
......@@ -1902,6 +1944,7 @@ test "std.hash_map repeat putAssumeCapacity/remove" {
19021944
19031945 try map.ensureTotalCapacity(20);
19041946 const limit = map.unmanaged.available;
1947 const cycles = 100;
19051948
19061949 var i: u32 = 0;
19071950 while (i < limit) : (i += 1) {
......@@ -1911,7 +1954,7 @@ test "std.hash_map repeat putAssumeCapacity/remove" {
19111954 // Repeatedly delete/insert an entry without resizing the map.
19121955 // Put to different keys so entries don't land in the just-freed slot.
19131956 i = 0;
1914 while (i < 10 * limit) : (i += 1) {
1957 while (i < cycles * limit) : (i += 1) {
19151958 try testing.expect(map.remove(i));
19161959 if (i % 2 == 0) {
19171960 map.putAssumeCapacityNoClobber(limit + i, i);
......@@ -1920,9 +1963,13 @@ test "std.hash_map repeat putAssumeCapacity/remove" {
19201963 }
19211964 }
19221965
1923 i = 9 * limit;
1924 while (i < 10 * limit) : (i += 1) {
1925 try expectEqual(map.get(limit + i), i);
1966 i = (cycles - 1) * limit;
1967 while (i < cycles * limit) : (i += 1) {
1968 try expectEqual(map.get(i), null); // (removed) key miss
1969 try expectEqual(map.get(limit + i), i); // key hit
1970 const gop = map.getOrPutAssumeCapacity(limit + i);
1971 try testing.expect(gop.found_existing);
1972 try testing.expectEqual(gop.value_ptr.*, i);
19261973 }
19271974 try expectEqual(map.unmanaged.available, 0);
19281975 try expectEqual(map.unmanaged.count(), limit);