From 02097dff7049cecb16e301d37cd270159ef0f371 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Prokop=20Rand=C3=A1=C4=8Dek?= Date: Tue, 28 Apr 2026 19:49:51 +0200 Subject: [PATCH] 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. --- lib/std/array_hash_map.zig | 133 ++++++++++++++++++++++++++----------- 1 file changed, 94 insertions(+), 39 deletions(-) diff --git a/lib/std/array_hash_map.zig b/lib/std/array_hash_map.zig index 05e5e0c4da0d5fe39f12eecc255419aa8ac9e814..af209c71be78796708f8d95cf7e2fb818f528b44 100644 --- a/lib/std/array_hash_map.zig +++ b/lib/std/array_hash_map.zig @@ -887,17 +887,25 @@ pub fn Custom( } /// Modify an entry's key without reordering any entries. - pub fn setKey(self: *Self, gpa: Allocator, index: usize, new_key: K) Oom!void { + pub fn setKey(self: *Self, index: usize, new_key: K) void { if (@sizeOf(ByIndexContext) != 0) @compileError("Cannot infer context " ++ @typeName(Context) ++ ", call setKeyContext instead."); - return setKeyContext(self, gpa, index, new_key, undefined); + return setKeyContext(self, index, new_key, undefined); } - pub fn setKeyContext(self: *Self, gpa: Allocator, index: usize, new_key: K, ctx: Context) Oom!void { - const key_ptr = &self.entries.items(.key)[index]; - key_ptr.* = new_key; - if (store_hash) self.entries.items(.hash)[index] = checkedHash(ctx, key_ptr.*); - try rebuildIndex(self, gpa, undefined); + pub fn setKeyContext(self: *Self, index: usize, new_key: K, ctx: Context) void { + if (self.index_header) |header| { + self.removeFromIndexByIndex(index, if (store_hash) {} else ctx, header); + + self.entries.items(.key)[index] = new_key; + const h = checkedHash(ctx, new_key); + if (store_hash) self.entries.items(.hash)[index] = h; + + insertEntryIntoNewHeader(header, h, index); + } else { + self.entries.items(.key)[index] = new_key; + if (store_hash) self.entries.items(.hash)[index] = checkedHash(ctx, new_key); + } } fn rebuildIndex(self: *Self, gpa: Allocator, ctx: Context) Oom!void { @@ -1420,39 +1428,50 @@ pub fn Custom( fn insertAllEntriesIntoNewHeaderGeneric(self: *Self, ctx: ByIndexContext, header: *IndexHeader, comptime I: type) void { const slice = self.entries.slice(); const items = if (store_hash) slice.items(.hash) else slice.items(.key); + + for (items, 0..) |hash_or_key, i| { + const h = if (store_hash) hash_or_key else checkedHash(ctx, hash_or_key); + insertEntryIntoNewHeaderGeneric(header, h, i, I); + } + } + + fn insertEntryIntoNewHeader(header: *IndexHeader, h: u32, i: usize) void { + switch (header.capacityIndexType()) { + .u8 => insertEntryIntoNewHeaderGeneric(header, h, i, u8), + .u16 => insertEntryIntoNewHeaderGeneric(header, h, i, u16), + .u32 => insertEntryIntoNewHeaderGeneric(header, h, i, u32), + } + } + fn insertEntryIntoNewHeaderGeneric(header: *IndexHeader, h: u32, i: usize, comptime I: type) void { const indexes = header.indexes(I); - - entry_loop: for (items, 0..) |key, i| { - const h = if (store_hash) key else checkedHash(ctx, key); - const start_index = safeTruncate(usize, h); - const end_index = start_index +% indexes.len; - var index = start_index; - var entry_index = @as(I, @intCast(i)); - var distance_from_start_index: I = 0; - while (index != end_index) : ({ - index +%= 1; - distance_from_start_index += 1; - }) { - const slot = header.constrainIndex(index); - const next_index = indexes[slot]; - if (next_index.isEmpty()) { - indexes[slot] = .{ - .distance_from_start_index = distance_from_start_index, - .entry_index = entry_index, - }; - continue :entry_loop; - } - if (next_index.distance_from_start_index < distance_from_start_index) { - indexes[slot] = .{ - .distance_from_start_index = distance_from_start_index, - .entry_index = entry_index, - }; - distance_from_start_index = next_index.distance_from_start_index; - entry_index = next_index.entry_index; - } + const start_index = safeTruncate(usize, h); + const end_index = start_index +% indexes.len; + var index = start_index; + var entry_index: I = @intCast(i); + var distance_from_start_index: I = 0; + while (index != end_index) : ({ + index +%= 1; + distance_from_start_index += 1; + }) { + const slot = header.constrainIndex(index); + const next_index = indexes[slot]; + if (next_index.isEmpty()) { + indexes[slot] = .{ + .distance_from_start_index = distance_from_start_index, + .entry_index = entry_index, + }; + return; + } + if (next_index.distance_from_start_index < distance_from_start_index) { + indexes[slot] = .{ + .distance_from_start_index = distance_from_start_index, + .entry_index = entry_index, + }; + distance_from_start_index = next_index.distance_from_start_index; + entry_index = next_index.entry_index; } - unreachable; } + unreachable; } fn checkedHash(ctx: anytype, key: anytype) u32 { @@ -2118,7 +2137,7 @@ test "setKey storehash true" { try map.put(gpa, 12, 34); try map.put(gpa, 56, 78); - try map.setKey(gpa, 0, 42); + map.setKey(0, 42); try testing.expectEqual(2, map.count()); try testing.expectEqual(false, map.contains(12)); try testing.expectEqual(34, map.get(42)); @@ -2134,13 +2153,49 @@ test "setKey storehash false" { try map.put(gpa, 12, 34); try map.put(gpa, 56, 78); - try map.setKey(gpa, 0, 42); + map.setKey(0, 42); try testing.expectEqual(2, map.count()); try testing.expectEqual(false, map.contains(12)); try testing.expectEqual(34, map.get(42)); try testing.expectEqual(78, map.get(56)); } +test "setKey storehash false with index" { + const gpa = std.testing.allocator; + + const T = ArrayHashMap(usize, usize, AutoContext(usize), false); + + var map: T = .empty; + defer map.deinit(gpa); + + for (0..T.linear_scan_max + 1) |i| try map.put(gpa, i, i); + + map.setKey(0, 42); + try testing.expectEqual(T.linear_scan_max + 1, map.count()); + try testing.expectEqual(false, map.contains(0)); + try testing.expectEqual(0, map.get(42)); + + for (1..T.linear_scan_max + 1) |i| try testing.expectEqual(i, map.get(i)); +} + +test "setKey storehash true with index" { + const gpa = std.testing.allocator; + + const T = ArrayHashMap(usize, usize, AutoContext(usize), false); + + var map: ArrayHashMap(usize, usize, AutoContext(usize), true) = .empty; + defer map.deinit(gpa); + + for (0..T.linear_scan_max + 1) |i| try map.put(gpa, i, i); + + map.setKey(0, 42); + try testing.expectEqual(T.linear_scan_max + 1, map.count()); + try testing.expectEqual(false, map.contains(0)); + try testing.expectEqual(0, map.get(42)); + + for (1..T.linear_scan_max + 1) |i| try testing.expectEqual(i, map.get(i)); +} + pub fn getHashPtrAddrFn(comptime K: type, comptime Context: type) (fn (Context, K) u32) { return struct { fn hash(ctx: Context, key: K) u32 { -- 2.54.0