diff --git a/std/buf_map.zig b/std/buf_map.zig index a82d1b731a..6de0d20cdb 100644 --- a/std/buf_map.zig +++ b/std/buf_map.zig @@ -16,7 +16,7 @@ pub const BufMap = struct { return self; } - pub fn deinit(self: *const BufMap) void { + pub fn deinit(self: *BufMap) void { var it = self.hash_map.iterator(); while (true) { const entry = it.next() orelse break; @@ -27,16 +27,34 @@ pub const BufMap = struct { self.hash_map.deinit(); } - pub fn set(self: *BufMap, key: []const u8, value: []const u8) !void { - self.delete(key); - const key_copy = try self.copy(key); - errdefer self.free(key_copy); - const value_copy = try self.copy(value); - errdefer self.free(value_copy); - _ = try self.hash_map.put(key_copy, value_copy); + /// Same as `set` but the key and value become owned by the BufMap rather + /// than being copied. + /// If `setMove` fails, the ownership of key and value does not transfer. + pub fn setMove(self: *BufMap, key: []u8, value: []u8) !void { + const get_or_put = try self.hash_map.getOrPut(key); + if (get_or_put.found_existing) { + self.free(get_or_put.kv.key); + get_or_put.kv.key = key; + } + get_or_put.kv.value = value; } - pub fn get(self: *const BufMap, key: []const u8) ?[]const u8 { + /// `key` and `value` are copied into the BufMap. + pub fn set(self: *BufMap, key: []const u8, value: []const u8) !void { + const value_copy = try self.copy(value); + errdefer self.free(value_copy); + // Avoid copying key if it already exists + const get_or_put = try self.hash_map.getOrPut(key); + if (!get_or_put.found_existing) { + get_or_put.kv.key = self.copy(key) catch |err| { + _ = self.hash_map.remove(key); + return err; + }; + } + get_or_put.kv.value = value_copy; + } + + pub fn get(self: BufMap, key: []const u8) ?[]const u8 { const entry = self.hash_map.get(key) orelse return null; return entry.value; } @@ -47,7 +65,7 @@ pub const BufMap = struct { self.free(entry.value); } - pub fn count(self: *const BufMap) usize { + pub fn count(self: BufMap) usize { return self.hash_map.count(); } @@ -55,11 +73,11 @@ pub const BufMap = struct { return self.hash_map.iterator(); } - fn free(self: *const BufMap, value: []const u8) void { + fn free(self: BufMap, value: []const u8) void { self.hash_map.allocator.free(value); } - fn copy(self: *const BufMap, value: []const u8) ![]const u8 { + fn copy(self: BufMap, value: []const u8) ![]u8 { return mem.dupe(self.hash_map.allocator, u8, value); } }; diff --git a/std/os/index.zig b/std/os/index.zig index 779fa8f8c8..be82ad4716 100644 --- a/std/os/index.zig +++ b/std/os/index.zig @@ -705,54 +705,28 @@ pub fn getEnvMap(allocator: *Allocator) !BufMap { const ptr = windows.GetEnvironmentStringsW() orelse return error.OutOfMemory; defer assert(windows.FreeEnvironmentStringsW(ptr) != 0); - var buf: [100]u8 = undefined; - var i: usize = 0; while (true) { if (ptr[i] == 0) return result; const key_start = i; - var fallocator = &std.heap.FixedBufferAllocator.init(buf[0..]).allocator; while (ptr[i] != 0 and ptr[i] != '=') : (i += 1) {} - - const key_slice = ptr[key_start..i]; - var key: []u8 = undefined; - var heap_key = false; - - key = std.unicode.utf16leToUtf8Alloc(fallocator, key_slice) catch undefined; - - if (key.len == 0) { - key = try std.unicode.utf16leToUtf8Alloc(allocator, key_slice); - heap_key = true; - } + const key_w = ptr[key_start..i]; + const key = try std.unicode.utf16leToUtf8Alloc(allocator, key_w); + errdefer allocator.free(key); if (ptr[i] == '=') i += 1; const value_start = i; while (ptr[i] != 0) : (i += 1) {} - - const value_slice = ptr[value_start..i]; - var value: []u8 = undefined; - var heap_value = false; - - value = std.unicode.utf16leToUtf8Alloc(fallocator, value_slice) catch undefined; - - if (value.len == 0) { - value = try std.unicode.utf16leToUtf8Alloc(allocator, value_slice); - heap_value = true; - } + const value_w = ptr[value_start..i]; + const value = try std.unicode.utf16leToUtf8Alloc(allocator, value_w); + errdefer allocator.free(value); i += 1; // skip over null byte - try result.set(key, value); - - if (heap_key) { - allocator.free(key); - } - if (heap_value) { - allocator.free(value); - } + try result.setMove(key, value); } } else { for (posix_environ_raw) |ptr| { @@ -795,9 +769,6 @@ pub fn getEnvPosix(key: []const u8) ?[]const u8 { pub const GetEnvVarOwnedError = error{ OutOfMemory, EnvironmentVariableNotFound, - DanglingSurrogateHalf, - ExpectedSecondSurrogateHalf, - UnexpectedSecondSurrogateHalf, /// See https://github.com/ziglang/zig/issues/1774 InvalidUtf8, @@ -833,7 +804,12 @@ pub fn getEnvVarOwned(allocator: *mem.Allocator, key: []const u8) GetEnvVarOwned continue; } - return try std.unicode.utf16leToUtf8Alloc(allocator, buf); + return std.unicode.utf16leToUtf8Alloc(allocator, buf) catch |err| switch (err) { + error.DanglingSurrogateHalf => return error.InvalidUtf8, + error.ExpectedSecondSurrogateHalf => return error.InvalidUtf8, + error.UnexpectedSecondSurrogateHalf => return error.InvalidUtf8, + error.OutOfMemory => return error.OutOfMemory, + }; } } else { const result = getEnvPosix(key) orelse return error.EnvironmentVariableNotFound; @@ -846,7 +822,6 @@ test "os.getEnvVarOwned" { debug.assertError(getEnvVarOwned(ga, "BADENV"), error.EnvironmentVariableNotFound); } - /// Caller must free the returned memory. pub fn getCwdAlloc(allocator: *Allocator) ![]u8 { var buf: [MAX_PATH_BYTES]u8 = undefined;