fixups

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);
     }
 };

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;