authorgravatar for prokop@rdck.devProkop Randáček <prokop@rdck.dev> 2026-04-28 19:49:51+02:00
committergravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2026-05-02 20:15:10+02:00
log02097dff7049cecb16e301d37cd270159ef0f371
tree6806b9233b09088bfc487e833cf34cc541242401
parent04481c76cb5f2ec78bf4b10c0cdb090e8a64240e

setKey shouldn't recompute the entire index

setKey used to just call rebuildIndex which allocates a new index and inserts into it all entries. This is the only allocation that setKey needed to do. Instead I propose that setKey sets the key in the entries and then does a remove of the old key from the index and insert of the new key into the index. This eliminates the need for setKey to allocate and makes it infallible To achieve this I extracted a helper function for inserting a single entry into the index. Function for removing one entry from the index already exists. setKey now just calls these two functions.

1 files changed, 94 insertions(+), 39 deletions(-)

lib/std/array_hash_map.zig+94-39
......@@ -887,17 +887,25 @@ pub fn Custom(
887887 }
888888
889889 /// Modify an entry's key without reordering any entries.
890 pub fn setKey(self: *Self, gpa: Allocator, index: usize, new_key: K) Oom!void {
890 pub fn setKey(self: *Self, index: usize, new_key: K) void {
891891 if (@sizeOf(ByIndexContext) != 0)
892892 @compileError("Cannot infer context " ++ @typeName(Context) ++ ", call setKeyContext instead.");
893 return setKeyContext(self, gpa, index, new_key, undefined);
893 return setKeyContext(self, index, new_key, undefined);
894894 }
895895
896 pub fn setKeyContext(self: *Self, gpa: Allocator, index: usize, new_key: K, ctx: Context) Oom!void {
897 const key_ptr = &self.entries.items(.key)[index];
898 key_ptr.* = new_key;
899 if (store_hash) self.entries.items(.hash)[index] = checkedHash(ctx, key_ptr.*);
900 try rebuildIndex(self, gpa, undefined);
896 pub fn setKeyContext(self: *Self, index: usize, new_key: K, ctx: Context) void {
897 if (self.index_header) |header| {
898 self.removeFromIndexByIndex(index, if (store_hash) {} else ctx, header);
899
900 self.entries.items(.key)[index] = new_key;
901 const h = checkedHash(ctx, new_key);
902 if (store_hash) self.entries.items(.hash)[index] = h;
903
904 insertEntryIntoNewHeader(header, h, index);
905 } else {
906 self.entries.items(.key)[index] = new_key;
907 if (store_hash) self.entries.items(.hash)[index] = checkedHash(ctx, new_key);
908 }
901909 }
902910
903911 fn rebuildIndex(self: *Self, gpa: Allocator, ctx: Context) Oom!void {
......@@ -1420,39 +1428,50 @@ pub fn Custom(
14201428 fn insertAllEntriesIntoNewHeaderGeneric(self: *Self, ctx: ByIndexContext, header: *IndexHeader, comptime I: type) void {
14211429 const slice = self.entries.slice();
14221430 const items = if (store_hash) slice.items(.hash) else slice.items(.key);
1423 const indexes = header.indexes(I);
14241431
1425 entry_loop: for (items, 0..) |key, i| {
1426 const h = if (store_hash) key else checkedHash(ctx, key);
1427 const start_index = safeTruncate(usize, h);
1428 const end_index = start_index +% indexes.len;
1429 var index = start_index;
1430 var entry_index = @as(I, @intCast(i));
1431 var distance_from_start_index: I = 0;
1432 while (index != end_index) : ({
1433 index +%= 1;
1434 distance_from_start_index += 1;
1435 }) {
1436 const slot = header.constrainIndex(index);
1437 const next_index = indexes[slot];
1438 if (next_index.isEmpty()) {
1439 indexes[slot] = .{
1440 .distance_from_start_index = distance_from_start_index,
1441 .entry_index = entry_index,
1442 };
1443 continue :entry_loop;
1444 }
1445 if (next_index.distance_from_start_index < distance_from_start_index) {
1446 indexes[slot] = .{
1447 .distance_from_start_index = distance_from_start_index,
1448 .entry_index = entry_index,
1449 };
1450 distance_from_start_index = next_index.distance_from_start_index;
1451 entry_index = next_index.entry_index;
1452 }
1432 for (items, 0..) |hash_or_key, i| {
1433 const h = if (store_hash) hash_or_key else checkedHash(ctx, hash_or_key);
1434 insertEntryIntoNewHeaderGeneric(header, h, i, I);
1435 }
1436 }
1437
1438 fn insertEntryIntoNewHeader(header: *IndexHeader, h: u32, i: usize) void {
1439 switch (header.capacityIndexType()) {
1440 .u8 => insertEntryIntoNewHeaderGeneric(header, h, i, u8),
1441 .u16 => insertEntryIntoNewHeaderGeneric(header, h, i, u16),
1442 .u32 => insertEntryIntoNewHeaderGeneric(header, h, i, u32),
1443 }
1444 }
1445 fn insertEntryIntoNewHeaderGeneric(header: *IndexHeader, h: u32, i: usize, comptime I: type) void {
1446 const indexes = header.indexes(I);
1447 const start_index = safeTruncate(usize, h);
1448 const end_index = start_index +% indexes.len;
1449 var index = start_index;
1450 var entry_index: I = @intCast(i);
1451 var distance_from_start_index: I = 0;
1452 while (index != end_index) : ({
1453 index +%= 1;
1454 distance_from_start_index += 1;
1455 }) {
1456 const slot = header.constrainIndex(index);
1457 const next_index = indexes[slot];
1458 if (next_index.isEmpty()) {
1459 indexes[slot] = .{
1460 .distance_from_start_index = distance_from_start_index,
1461 .entry_index = entry_index,
1462 };
1463 return;
1464 }
1465 if (next_index.distance_from_start_index < distance_from_start_index) {
1466 indexes[slot] = .{
1467 .distance_from_start_index = distance_from_start_index,
1468 .entry_index = entry_index,
1469 };
1470 distance_from_start_index = next_index.distance_from_start_index;
1471 entry_index = next_index.entry_index;
14531472 }
1454 unreachable;
14551473 }
1474 unreachable;
14561475 }
14571476
14581477 fn checkedHash(ctx: anytype, key: anytype) u32 {
......@@ -2118,7 +2137,7 @@ test "setKey storehash true" {
21182137 try map.put(gpa, 12, 34);
21192138 try map.put(gpa, 56, 78);
21202139
2121 try map.setKey(gpa, 0, 42);
2140 map.setKey(0, 42);
21222141 try testing.expectEqual(2, map.count());
21232142 try testing.expectEqual(false, map.contains(12));
21242143 try testing.expectEqual(34, map.get(42));
......@@ -2134,13 +2153,49 @@ test "setKey storehash false" {
21342153 try map.put(gpa, 12, 34);
21352154 try map.put(gpa, 56, 78);
21362155
2137 try map.setKey(gpa, 0, 42);
2156 map.setKey(0, 42);
21382157 try testing.expectEqual(2, map.count());
21392158 try testing.expectEqual(false, map.contains(12));
21402159 try testing.expectEqual(34, map.get(42));
21412160 try testing.expectEqual(78, map.get(56));
21422161}
21432162
2163test "setKey storehash false with index" {
2164 const gpa = std.testing.allocator;
2165
2166 const T = ArrayHashMap(usize, usize, AutoContext(usize), false);
2167
2168 var map: T = .empty;
2169 defer map.deinit(gpa);
2170
2171 for (0..T.linear_scan_max + 1) |i| try map.put(gpa, i, i);
2172
2173 map.setKey(0, 42);
2174 try testing.expectEqual(T.linear_scan_max + 1, map.count());
2175 try testing.expectEqual(false, map.contains(0));
2176 try testing.expectEqual(0, map.get(42));
2177
2178 for (1..T.linear_scan_max + 1) |i| try testing.expectEqual(i, map.get(i));
2179}
2180
2181test "setKey storehash true with index" {
2182 const gpa = std.testing.allocator;
2183
2184 const T = ArrayHashMap(usize, usize, AutoContext(usize), false);
2185
2186 var map: ArrayHashMap(usize, usize, AutoContext(usize), true) = .empty;
2187 defer map.deinit(gpa);
2188
2189 for (0..T.linear_scan_max + 1) |i| try map.put(gpa, i, i);
2190
2191 map.setKey(0, 42);
2192 try testing.expectEqual(T.linear_scan_max + 1, map.count());
2193 try testing.expectEqual(false, map.contains(0));
2194 try testing.expectEqual(0, map.get(42));
2195
2196 for (1..T.linear_scan_max + 1) |i| try testing.expectEqual(i, map.get(i));
2197}
2198
21442199pub fn getHashPtrAddrFn(comptime K: type, comptime Context: type) (fn (Context, K) u32) {
21452200 return struct {
21462201 fn hash(ctx: Context, key: K) u32 {