Merge branch 'suirad-windows-wide-imports'

Author: andrewrk

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();
     }
 
+    /// 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;
+    }
+
+    /// `key` and `value` are copied into the BufMap.
     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);
+        // 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: *const BufMap, key: []const u8) ?[]const u8 {
+    pub fn get(self: BufMap, key: []const u8) ?[]const u8 {
         const entry = self.hash_map.get(key) orelse return null;
         return entry.value;
     }
@@ -47,19 +65,19 @@ 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();
     }
 
     pub fn iterator(self: *const BufMap) BufMapHashMap.Iterator {
         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

@@ -702,8 +702,8 @@ pub fn getEnvMap(allocator: *Allocator) !BufMap {
     errdefer result.deinit();
 
     if (is_windows) {
-        const ptr = windows.GetEnvironmentStringsA() orelse return error.OutOfMemory;
-        defer assert(windows.FreeEnvironmentStringsA(ptr) != 0);
+        const ptr = windows.GetEnvironmentStringsW() orelse return error.OutOfMemory;
+        defer assert(windows.FreeEnvironmentStringsW(ptr) != 0);
 
         var i: usize = 0;
         while (true) {
@@ -712,17 +712,21 @@ pub fn getEnvMap(allocator: *Allocator) !BufMap {
             const key_start = i;
 
             while (ptr[i] != 0 and ptr[i] != '=') : (i += 1) {}
-            const key = ptr[key_start..i];
+            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 = ptr[value_start..i];
+            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);
+            try result.setMove(key, value);
         }
     } else {
         for (posix_environ_raw) |ptr| {
@@ -740,6 +744,11 @@ pub fn getEnvMap(allocator: *Allocator) !BufMap {
     }
 }
 
+test "os.getEnvMap" {
+    var env = try getEnvMap(std.debug.global_allocator);
+    defer env.deinit();
+}
+
 /// TODO make this go through libc when we have it
 pub fn getEnvPosix(key: []const u8) ?[]const u8 {
     for (posix_environ_raw) |ptr| {
@@ -760,21 +769,24 @@ pub fn getEnvPosix(key: []const u8) ?[]const u8 {
 pub const GetEnvVarOwnedError = error{
     OutOfMemory,
     EnvironmentVariableNotFound,
+
+    /// See https://github.com/ziglang/zig/issues/1774
+    InvalidUtf8,
 };
 
 /// Caller must free returned memory.
 /// TODO make this go through libc when we have it
 pub fn getEnvVarOwned(allocator: *mem.Allocator, key: []const u8) GetEnvVarOwnedError![]u8 {
     if (is_windows) {
-        const key_with_null = try cstr.addNullByte(allocator, key);
+        const key_with_null = try std.unicode.utf8ToUtf16LeWithNull(allocator, key);
         defer allocator.free(key_with_null);
 
-        var buf = try allocator.alloc(u8, 256);
-        errdefer allocator.free(buf);
+        var buf = try allocator.alloc(u16, 256);
+        defer allocator.free(buf);
 
         while (true) {
             const windows_buf_len = math.cast(windows.DWORD, buf.len) catch return error.OutOfMemory;
-            const result = windows.GetEnvironmentVariableA(key_with_null.ptr, buf.ptr, windows_buf_len);
+            const result = windows.GetEnvironmentVariableW(key_with_null.ptr, buf.ptr, windows_buf_len);
 
             if (result == 0) {
                 const err = windows.GetLastError();
@@ -788,18 +800,28 @@ pub fn getEnvVarOwned(allocator: *mem.Allocator, key: []const u8) GetEnvVarOwned
             }
 
             if (result > buf.len) {
-                buf = try allocator.realloc(u8, buf, result);
+                buf = try allocator.realloc(u16, buf, result);
                 continue;
             }
 
-            return allocator.shrink(u8, buf, result);
+            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;
         return mem.dupe(allocator, u8, result);
     }
 }
 
+test "os.getEnvVarOwned" {
+    var ga = debug.global_allocator;
+    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;

std/os/windows/kernel32.zig

@@ -50,7 +50,7 @@ pub extern "kernel32" stdcallcc fn FindFirstFileW(lpFileName: [*]const u16, lpFi
 pub extern "kernel32" stdcallcc fn FindClose(hFindFile: HANDLE) BOOL;
 pub extern "kernel32" stdcallcc fn FindNextFileW(hFindFile: HANDLE, lpFindFileData: *WIN32_FIND_DATAW) BOOL;
 
-pub extern "kernel32" stdcallcc fn FreeEnvironmentStringsA(penv: [*]u8) BOOL;
+pub extern "kernel32" stdcallcc fn FreeEnvironmentStringsW(penv: [*]u16) BOOL;
 
 pub extern "kernel32" stdcallcc fn GetCommandLineA() LPSTR;
 
@@ -63,9 +63,9 @@ pub extern "kernel32" stdcallcc fn GetCurrentDirectoryW(nBufferLength: DWORD, lp
 pub extern "kernel32" stdcallcc fn GetCurrentThread() HANDLE;
 pub extern "kernel32" stdcallcc fn GetCurrentThreadId() DWORD;
 
-pub extern "kernel32" stdcallcc fn GetEnvironmentStringsA() ?[*]u8;
+pub extern "kernel32" stdcallcc fn GetEnvironmentStringsW() ?[*]u16;
 
-pub extern "kernel32" stdcallcc fn GetEnvironmentVariableA(lpName: LPCSTR, lpBuffer: LPSTR, nSize: DWORD) DWORD;
+pub extern "kernel32" stdcallcc fn GetEnvironmentVariableW(lpName: LPWSTR, lpBuffer: LPWSTR, nSize: DWORD) DWORD;
 
 pub extern "kernel32" stdcallcc fn GetExitCodeProcess(hProcess: HANDLE, lpExitCode: *DWORD) BOOL;