Merge pull request #10796 from marler8997/envmap

Envmap

6 files changed, 231 insertions, 37 deletions

diff --git a/lib/std/buf_map.zig b/lib/std/buf_map.zig
index 5b26ae9684..2a6239c490 100644
--- a/lib/std/buf_map.zig
+++ b/lib/std/buf_map.zig
@@ -82,7 +82,7 @@ pub const BufMap = struct {
     }
 
     /// Returns the number of KV pairs stored in the map.
-    pub fn count(self: BufMap) usize {
+    pub fn count(self: BufMap) BufMapHashMap.Size {
         return self.hash_map.count();
     }
 
diff --git a/lib/std/build.zig b/lib/std/build.zig
index b0200a928b..0f196cc9e9 100644
--- a/lib/std/build.zig
+++ b/lib/std/build.zig
@@ -12,7 +12,7 @@ const StringHashMap = std.StringHashMap;
 const Allocator = mem.Allocator;
 const process = std.process;
 const BufSet = std.BufSet;
-const BufMap = std.BufMap;
+const EnvMap = std.process.EnvMap;
 const fmt_lib = std.fmt;
 const File = std.fs.File;
 const CrossTarget = std.zig.CrossTarget;
@@ -48,7 +48,7 @@ pub const Builder = struct {
     invalid_user_input: bool,
     zig_exe: []const u8,
     default_step: *Step,
-    env_map: *BufMap,
+    env_map: *EnvMap,
     top_level_steps: ArrayList(*TopLevelStep),
     install_prefix: []const u8,
     dest_dir: ?[]const u8,
@@ -167,7 +167,7 @@ pub const Builder = struct {
         cache_root: []const u8,
         global_cache_root: []const u8,
     ) !*Builder {
-        const env_map = try allocator.create(BufMap);
+        const env_map = try allocator.create(EnvMap);
         env_map.* = try process.getEnvMap(allocator);
 
         const host = try NativeTargetInfo.detect(allocator, .{});
@@ -963,7 +963,7 @@ pub const Builder = struct {
         warn("\n", .{});
     }
 
-    pub fn spawnChildEnvMap(self: *Builder, cwd: ?[]const u8, env_map: *const BufMap, argv: []const []const u8) !void {
+    pub fn spawnChildEnvMap(self: *Builder, cwd: ?[]const u8, env_map: *const EnvMap, argv: []const []const u8) !void {
         if (self.verbose) {
             printCmd(cwd, argv);
         }
diff --git a/lib/std/build/RunStep.zig b/lib/std/build/RunStep.zig
index e00fe3deb6..e8cf87a441 100644
--- a/lib/std/build/RunStep.zig
+++ b/lib/std/build/RunStep.zig
@@ -9,7 +9,7 @@ const fs = std.fs;
 const mem = std.mem;
 const process = std.process;
 const ArrayList = std.ArrayList;
-const BufMap = std.BufMap;
+const EnvMap = process.EnvMap;
 const Allocator = mem.Allocator;
 const ExecError = build.Builder.ExecError;
 
@@ -29,7 +29,7 @@ argv: ArrayList(Arg),
 cwd: ?[]const u8,
 
 /// Override this field to modify the environment, or use setEnvironmentVariable
-env_map: ?*BufMap,
+env_map: ?*EnvMap,
 
 stdout_action: StdIoAction = .inherit,
 stderr_action: StdIoAction = .inherit,
@@ -91,27 +91,16 @@ pub fn addArgs(self: *RunStep, args: []const []const u8) void {
 }
 
 pub fn clearEnvironment(self: *RunStep) void {
-    const new_env_map = self.builder.allocator.create(BufMap) catch unreachable;
-    new_env_map.* = BufMap.init(self.builder.allocator);
+    const new_env_map = self.builder.allocator.create(EnvMap) catch unreachable;
+    new_env_map.* = EnvMap.init(self.builder.allocator);
     self.env_map = new_env_map;
 }
 
 pub fn addPathDir(self: *RunStep, search_path: []const u8) void {
     const env_map = self.getEnvMap();
 
-    var key: []const u8 = undefined;
-    var prev_path: ?[]const u8 = undefined;
-    if (builtin.os.tag == .windows) {
-        key = "Path";
-        prev_path = env_map.get(key);
-        if (prev_path == null) {
-            key = "PATH";
-            prev_path = env_map.get(key);
-        }
-    } else {
-        key = "PATH";
-        prev_path = env_map.get(key);
-    }
+    const key = "PATH";
+    var prev_path = env_map.get(key);
 
     if (prev_path) |pp| {
         const new_path = self.builder.fmt("{s}" ++ [1]u8{fs.path.delimiter} ++ "{s}", .{ pp, search_path });
@@ -121,9 +110,9 @@ pub fn addPathDir(self: *RunStep, search_path: []const u8) void {
     }
 }
 
-pub fn getEnvMap(self: *RunStep) *BufMap {
+pub fn getEnvMap(self: *RunStep) *EnvMap {
     return self.env_map orelse {
-        const env_map = self.builder.allocator.create(BufMap) catch unreachable;
+        const env_map = self.builder.allocator.create(EnvMap) catch unreachable;
         env_map.* = process.getEnvMap(self.builder.allocator) catch unreachable;
         self.env_map = env_map;
         return env_map;
diff --git a/lib/std/child_process.zig b/lib/std/child_process.zig
index 5f01ed01dd..0bb737decb 100644
--- a/lib/std/child_process.zig
+++ b/lib/std/child_process.zig
@@ -12,7 +12,7 @@ const linux = os.linux;
 const mem = std.mem;
 const math = std.math;
 const debug = std.debug;
-const BufMap = std.BufMap;
+const EnvMap = process.EnvMap;
 const Os = std.builtin.Os;
 const TailQueue = std.TailQueue;
 const maxInt = std.math.maxInt;
@@ -34,7 +34,7 @@ pub const ChildProcess = struct {
     argv: []const []const u8,
 
     /// Leave as null to use the current env map using the supplied allocator.
-    env_map: ?*const BufMap,
+    env_map: ?*const EnvMap,
 
     stdin_behavior: StdIo,
     stdout_behavior: StdIo,
@@ -375,7 +375,7 @@ pub const ChildProcess = struct {
         argv: []const []const u8,
         cwd: ?[]const u8 = null,
         cwd_dir: ?fs.Dir = null,
-        env_map: ?*const BufMap = null,
+        env_map: ?*const EnvMap = null,
         max_output_bytes: usize = 50 * 1024,
         expand_arg0: Arg0Expand = .no_expand,
     }) !ExecResult {
@@ -1237,7 +1237,7 @@ fn readIntFd(fd: i32) !ErrInt {
 }
 
 /// Caller must free result.
-pub fn createWindowsEnvBlock(allocator: mem.Allocator, env_map: *const BufMap) ![]u16 {
+pub fn createWindowsEnvBlock(allocator: mem.Allocator, env_map: *const EnvMap) ![]u16 {
     // count bytes needed
     const max_chars_needed = x: {
         var max_chars_needed: usize = 4; // 4 for the final 4 null bytes
@@ -1273,7 +1273,7 @@ pub fn createWindowsEnvBlock(allocator: mem.Allocator, env_map: *const BufMap) !
     return allocator.shrink(result, i);
 }
 
-pub fn createNullDelimitedEnvMap(arena: mem.Allocator, env_map: *const std.BufMap) ![:null]?[*:0]u8 {
+pub fn createNullDelimitedEnvMap(arena: mem.Allocator, env_map: *const EnvMap) ![:null]?[*:0]u8 {
     const envp_count = env_map.count();
     const envp_buf = try arena.allocSentinel(?[*:0]u8, envp_count, null);
     {
@@ -1294,7 +1294,7 @@ pub fn createNullDelimitedEnvMap(arena: mem.Allocator, env_map: *const std.BufMa
 test "createNullDelimitedEnvMap" {
     const testing = std.testing;
     const allocator = testing.allocator;
-    var envmap = BufMap.init(allocator);
+    var envmap = EnvMap.init(allocator);
     defer envmap.deinit();
 
     try envmap.put("HOME", "/home/ifreund");
diff --git a/lib/std/os/windows/ntdll.zig b/lib/std/os/windows/ntdll.zig
index e3e590b094..bf9dc9bd2f 100644
--- a/lib/std/os/windows/ntdll.zig
+++ b/lib/std/os/windows/ntdll.zig
@@ -229,6 +229,10 @@ pub extern "ntdll" fn RtlEqualUnicodeString(
     CaseInSensitive: BOOLEAN,
 ) callconv(WINAPI) BOOLEAN;
 
+pub extern "ntdll" fn RtlUpcaseUnicodeChar(
+    SourceCharacter: u16,
+) callconv(WINAPI) u16;
+
 pub extern "ntdll" fn NtLockFile(
     FileHandle: HANDLE,
     Event: ?HANDLE,
diff --git a/lib/std/process.zig b/lib/std/process.zig
index c0f11b22ce..0b64b5910d 100644
--- a/lib/std/process.zig
+++ b/lib/std/process.zig
@@ -2,7 +2,6 @@ const std = @import("std.zig");
 const builtin = @import("builtin");
 const os = std.os;
 const fs = std.fs;
-const BufMap = std.BufMap;
 const mem = std.mem;
 const math = std.math;
 const Allocator = mem.Allocator;
@@ -53,9 +52,205 @@ test "getCwdAlloc" {
     testing.allocator.free(cwd);
 }
 
-/// Caller owns resulting `BufMap`.
-pub fn getEnvMap(allocator: Allocator) !BufMap {
-    var result = BufMap.init(allocator);
+pub const EnvMap = struct {
+    hash_map: HashMap,
+
+    const HashMap = std.HashMap(
+        []const u8,
+        []const u8,
+        EnvNameHashContext,
+        std.hash_map.default_max_load_percentage,
+    );
+
+    pub const Size = HashMap.Size;
+
+    pub const EnvNameHashContext = struct {
+        fn upcase(c: u21) u21 {
+            if (c <= std.math.maxInt(u16))
+                return std.os.windows.ntdll.RtlUpcaseUnicodeChar(@intCast(u16, c));
+            return c;
+        }
+
+        pub fn hash(self: @This(), s: []const u8) u64 {
+            _ = self;
+            if (builtin.os.tag == .windows) {
+                var h = std.hash.Wyhash.init(0);
+                var it = std.unicode.Utf8View.initUnchecked(s).iterator();
+                while (it.nextCodepoint()) |cp| {
+                    const cp_upper = upcase(cp);
+                    h.update(&[_]u8{
+                        @intCast(u8, (cp_upper >> 16) & 0xff),
+                        @intCast(u8, (cp_upper >> 8) & 0xff),
+                        @intCast(u8, (cp_upper >> 0) & 0xff),
+                    });
+                }
+                return h.final();
+            }
+            return std.hash_map.hashString(s);
+        }
+
+        pub fn eql(self: @This(), a: []const u8, b: []const u8) bool {
+            _ = self;
+            if (builtin.os.tag == .windows) {
+                var it_a = std.unicode.Utf8View.initUnchecked(a).iterator();
+                var it_b = std.unicode.Utf8View.initUnchecked(b).iterator();
+                while (true) {
+                    const c_a = it_a.nextCodepoint() orelse break;
+                    const c_b = it_b.nextCodepoint() orelse return false;
+                    if (upcase(c_a) != upcase(c_b))
+                        return false;
+                }
+                return if (it_b.nextCodepoint()) |_| false else true;
+            }
+            return std.hash_map.eqlString(a, b);
+        }
+    };
+
+    /// Create a EnvMap backed by a specific allocator.
+    /// That allocator will be used for both backing allocations
+    /// and string deduplication.
+    pub fn init(allocator: Allocator) EnvMap {
+        return EnvMap{ .hash_map = HashMap.init(allocator) };
+    }
+
+    /// Free the backing storage of the map, as well as all
+    /// of the stored keys and values.
+    pub fn deinit(self: *EnvMap) void {
+        var it = self.hash_map.iterator();
+        while (it.next()) |entry| {
+            self.free(entry.key_ptr.*);
+            self.free(entry.value_ptr.*);
+        }
+
+        self.hash_map.deinit();
+    }
+
+    /// Same as `put` but the key and value become owned by the EnvMap rather
+    /// than being copied.
+    /// If `putMove` fails, the ownership of key and value does not transfer.
+    /// On Windows `key` must be a valid UTF-8 string.
+    pub fn putMove(self: *EnvMap, 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.key_ptr.*);
+            self.free(get_or_put.value_ptr.*);
+            get_or_put.key_ptr.* = key;
+        }
+        get_or_put.value_ptr.* = value;
+    }
+
+    /// `key` and `value` are copied into the EnvMap.
+    /// On Windows `key` must be a valid UTF-8 string.
+    pub fn put(self: *EnvMap, key: []const u8, value: []const u8) !void {
+        const value_copy = try self.copy(value);
+        errdefer self.free(value_copy);
+        const get_or_put = try self.hash_map.getOrPut(key);
+        if (get_or_put.found_existing) {
+            self.free(get_or_put.value_ptr.*);
+        } else {
+            get_or_put.key_ptr.* = self.copy(key) catch |err| {
+                _ = self.hash_map.remove(key);
+                return err;
+            };
+        }
+        get_or_put.value_ptr.* = value_copy;
+    }
+
+    /// Find the address of the value associated with a key.
+    /// The returned pointer is invalidated if the map resizes.
+    /// On Windows `key` must be a valid UTF-8 string.
+    pub fn getPtr(self: EnvMap, key: []const u8) ?*[]const u8 {
+        return self.hash_map.getPtr(key);
+    }
+
+    /// Return the map's copy of the value associated with
+    /// a key.  The returned string is invalidated if this
+    /// key is removed from the map.
+    /// On Windows `key` must be a valid UTF-8 string.
+    pub fn get(self: EnvMap, key: []const u8) ?[]const u8 {
+        return self.hash_map.get(key);
+    }
+
+    /// Removes the item from the map and frees its value.
+    /// This invalidates the value returned by get() for this key.
+    /// On Windows `key` must be a valid UTF-8 string.
+    pub fn remove(self: *EnvMap, key: []const u8) void {
+        const kv = self.hash_map.fetchRemove(key) orelse return;
+        self.free(kv.key);
+        self.free(kv.value);
+    }
+
+    /// Returns the number of KV pairs stored in the map.
+    pub fn count(self: EnvMap) HashMap.Size {
+        return self.hash_map.count();
+    }
+
+    /// Returns an iterator over entries in the map.
+    pub fn iterator(self: *const EnvMap) HashMap.Iterator {
+        return self.hash_map.iterator();
+    }
+
+    fn free(self: EnvMap, value: []const u8) void {
+        self.hash_map.allocator.free(value);
+    }
+
+    fn copy(self: EnvMap, value: []const u8) ![]u8 {
+        return self.hash_map.allocator.dupe(u8, value);
+    }
+};
+
+test "EnvMap" {
+    var env = EnvMap.init(testing.allocator);
+    defer env.deinit();
+
+    try env.put("SOMETHING_NEW", "hello");
+    try testing.expectEqualStrings("hello", env.get("SOMETHING_NEW").?);
+    try testing.expectEqual(@as(EnvMap.Size, 1), env.count());
+
+    // overwrite
+    try env.put("SOMETHING_NEW", "something");
+    try testing.expectEqualStrings("something", env.get("SOMETHING_NEW").?);
+    try testing.expectEqual(@as(EnvMap.Size, 1), env.count());
+
+    // a new longer name to test the Windows-specific conversion buffer
+    try env.put("SOMETHING_NEW_AND_LONGER", "1");
+    try testing.expectEqualStrings("1", env.get("SOMETHING_NEW_AND_LONGER").?);
+    try testing.expectEqual(@as(EnvMap.Size, 2), env.count());
+
+    // case insensitivity on Windows only
+    if (builtin.os.tag == .windows) {
+        try testing.expectEqualStrings("1", env.get("something_New_aNd_LONGER").?);
+    } else {
+        try testing.expect(null == env.get("something_New_aNd_LONGER"));
+    }
+
+    var it = env.iterator();
+    var count: EnvMap.Size = 0;
+    while (it.next()) |entry| {
+        const is_an_expected_name = std.mem.eql(u8, "SOMETHING_NEW", entry.key_ptr.*) or std.mem.eql(u8, "SOMETHING_NEW_AND_LONGER", entry.key_ptr.*);
+        try testing.expect(is_an_expected_name);
+        count += 1;
+    }
+    try testing.expectEqual(@as(EnvMap.Size, 2), count);
+
+    env.remove("SOMETHING_NEW");
+    try testing.expect(env.get("SOMETHING_NEW") == null);
+
+    try testing.expectEqual(@as(EnvMap.Size, 1), env.count());
+
+    // test Unicode case-insensitivity on Windows
+    if (builtin.os.tag == .windows) {
+        try env.put("КИРиллИЦА", "something else");
+        try testing.expectEqualStrings("something else", env.get("кириллица").?);
+    }
+}
+
+/// Returns a snapshot of the environment variables of the current process.
+/// Any modifications to the resulting EnvMap will not be not reflected in the environment, and
+/// likewise, any future modifications to the environment will not be reflected in the EnvMap.
+/// Caller owns resulting `EnvMap` and should call its `deinit` fn when done.
+pub fn getEnvMap(allocator: Allocator) !EnvMap {
+    var result = EnvMap.init(allocator);
     errdefer result.deinit();
 
     if (builtin.os.tag == .windows) {
@@ -65,6 +260,12 @@ pub fn getEnvMap(allocator: Allocator) !BufMap {
         while (ptr[i] != 0) {
             const key_start = i;
 
+            // There are some special environment variables that start with =,
+            // so we need a special case to not treat = as a key/value separator
+            // if it's the first character.
+            // https://devblogs.microsoft.com/oldnewthing/20100506-00/?p=14133
+            if (ptr[key_start] == '=') i += 1;
+
             while (ptr[i] != 0 and ptr[i] != '=') : (i += 1) {}
             const key_w = ptr[key_start..i];
             const key = try std.unicode.utf16leToUtf8Alloc(allocator, key_w);
@@ -140,8 +341,8 @@ pub fn getEnvMap(allocator: Allocator) !BufMap {
     }
 }
 
-test "os.getEnvMap" {
-    var env = try getEnvMap(std.testing.allocator);
+test "getEnvMap" {
+    var env = try getEnvMap(testing.allocator);
     defer env.deinit();
 }
 
@@ -985,7 +1186,7 @@ pub fn execv(allocator: mem.Allocator, argv: []const []const u8) ExecvError {
 pub fn execve(
     allocator: mem.Allocator,
     argv: []const []const u8,
-    env_map: ?*const std.BufMap,
+    env_map: ?*const EnvMap,
 ) ExecvError {
     if (!can_execv) @compileError("The target OS does not support execv");