diff --git a/lib/std/hash_map.zig b/lib/std/hash_map.zig index 15b62e0d40..2c0010211a 100644 --- a/lib/std/hash_map.zig +++ b/lib/std/hash_map.zig @@ -629,6 +629,13 @@ pub fn HashMap( return self.unmanaged.removeAdapted(key, ctx); } + /// Delete the entry with key pointed to by keyPtr from the hash map. + /// keyPtr is assumed to be a valid pointer to a key that is present + /// in the hash map. + pub fn removeByPtr(self: *Self, keyPtr: *K) void { + self.unmanaged.removeByPtr(keyPtr); + } + /// Creates a copy of this map, using the same allocator pub fn clone(self: Self) !Self { var other = try self.unmanaged.cloneContext(self.allocator, self.ctx); @@ -1377,6 +1384,14 @@ pub fn HashMapUnmanaged( return self.getIndex(key, ctx) != null; } + fn removeByIndex(self: *Self, idx: usize) void { + self.metadata.?[idx].remove(); + self.keys()[idx] = undefined; + self.values()[idx] = undefined; + self.size -= 1; + self.available += 1; + } + /// If there is an `Entry` with a matching key, it is deleted from /// the hash map, and this function returns true. Otherwise this /// function returns false. @@ -1390,17 +1405,29 @@ pub fn HashMapUnmanaged( } pub fn removeAdapted(self: *Self, key: anytype, ctx: anytype) bool { if (self.getIndex(key, ctx)) |idx| { - self.metadata.?[idx].remove(); - self.keys()[idx] = undefined; - self.values()[idx] = undefined; - self.size -= 1; - self.available += 1; + self.removeByIndex(idx); return true; } return false; } + /// Delete the entry with key pointed to by keyPtr from the hash map. + /// keyPtr is assumed to be a valid pointer to a key that is present + /// in the hash map. + pub fn removeByPtr(self: *Self, keyPtr: *K) void { + // TODO: replace with pointer subtraction once supported by zig + // if @sizeOf(K) == 0 then there is at most one item in the hash + // map, which is assumed to exist as keyPtr must be valid. This + // item must be at index 0. + const idx = if (@sizeOf(K) > 0) + (@ptrToInt(keyPtr) - @ptrToInt(self.keys())) / @sizeOf(K) + else + 0; + + self.removeByIndex(idx); + } + fn initMetadatas(self: *Self) void { @memset(@ptrCast([*]u8, self.metadata.?), 0, @sizeOf(Metadata) * self.capacity()); } @@ -2082,3 +2109,47 @@ test "std.hash_map ensureUnusedCapacity" { // should not change the capacity. try testing.expectEqual(capacity, map.capacity()); } + +test "std.hash_map removeByPtr" { + var map = AutoHashMap(i32, u64).init(testing.allocator); + defer map.deinit(); + + var i: i32 = undefined; + + i = 0; + while (i < 10) : (i += 1) { + try map.put(i, 0); + } + + try testing.expect(map.count() == 10); + + i = 0; + while (i < 10) : (i += 1) { + const keyPtr = map.getKeyPtr(i); + try testing.expect(keyPtr != null); + + if (keyPtr) |ptr| { + map.removeByPtr(ptr); + } + } + + try testing.expect(map.count() == 0); +} + +test "std.hash_map removeByPtr 0 sized key" { + var map = AutoHashMap(u0, u64).init(testing.allocator); + defer map.deinit(); + + try map.put(0, 0); + + try testing.expect(map.count() == 1); + + const keyPtr = map.getKeyPtr(0); + try testing.expect(keyPtr != null); + + if (keyPtr) |ptr| { + map.removeByPtr(ptr); + } + + try testing.expect(map.count() == 0); +}