| ... | ... | @@ -2,7 +2,6 @@ const std = @import("std.zig"); |
| 2 | 2 | const builtin = @import("builtin"); |
| 3 | 3 | const os = std.os; |
| 4 | 4 | const fs = std.fs; |
| 5 | | const BufMap = std.BufMap; |
| 6 | 5 | const mem = std.mem; |
| 7 | 6 | const math = std.math; |
| 8 | 7 | const Allocator = mem.Allocator; |
| ... | ... | @@ -53,9 +52,205 @@ test "getCwdAlloc" { |
| 53 | 52 | testing.allocator.free(cwd); |
| 54 | 53 | } |
| 55 | 54 | |
| 56 | | /// Caller owns resulting `BufMap`. |
| 57 | | pub fn getEnvMap(allocator: Allocator) !BufMap { |
| 58 | | var result = BufMap.init(allocator); |
| 55 | pub const EnvMap = struct { |
| 56 | hash_map: HashMap, |
| 57 | |
| 58 | const HashMap = std.HashMap( |
| 59 | []const u8, |
| 60 | []const u8, |
| 61 | EnvNameHashContext, |
| 62 | std.hash_map.default_max_load_percentage, |
| 63 | ); |
| 64 | |
| 65 | pub const Size = HashMap.Size; |
| 66 | |
| 67 | pub const EnvNameHashContext = struct { |
| 68 | fn upcase(c: u21) u21 { |
| 69 | if (c <= std.math.maxInt(u16)) |
| 70 | return std.os.windows.ntdll.RtlUpcaseUnicodeChar(@intCast(u16, c)); |
| 71 | return c; |
| 72 | } |
| 73 | |
| 74 | pub fn hash(self: @This(), s: []const u8) u64 { |
| 75 | _ = self; |
| 76 | if (builtin.os.tag == .windows) { |
| 77 | var h = std.hash.Wyhash.init(0); |
| 78 | var it = std.unicode.Utf8View.initUnchecked(s).iterator(); |
| 79 | while (it.nextCodepoint()) |cp| { |
| 80 | const cp_upper = upcase(cp); |
| 81 | h.update(&[_]u8{ |
| 82 | @intCast(u8, (cp_upper >> 16) & 0xff), |
| 83 | @intCast(u8, (cp_upper >> 8) & 0xff), |
| 84 | @intCast(u8, (cp_upper >> 0) & 0xff), |
| 85 | }); |
| 86 | } |
| 87 | return h.final(); |
| 88 | } |
| 89 | return std.hash_map.hashString(s); |
| 90 | } |
| 91 | |
| 92 | pub fn eql(self: @This(), a: []const u8, b: []const u8) bool { |
| 93 | _ = self; |
| 94 | if (builtin.os.tag == .windows) { |
| 95 | var it_a = std.unicode.Utf8View.initUnchecked(a).iterator(); |
| 96 | var it_b = std.unicode.Utf8View.initUnchecked(b).iterator(); |
| 97 | while (true) { |
| 98 | const c_a = it_a.nextCodepoint() orelse break; |
| 99 | const c_b = it_b.nextCodepoint() orelse return false; |
| 100 | if (upcase(c_a) != upcase(c_b)) |
| 101 | return false; |
| 102 | } |
| 103 | return if (it_b.nextCodepoint()) |_| false else true; |
| 104 | } |
| 105 | return std.hash_map.eqlString(a, b); |
| 106 | } |
| 107 | }; |
| 108 | |
| 109 | /// Create a EnvMap backed by a specific allocator. |
| 110 | /// That allocator will be used for both backing allocations |
| 111 | /// and string deduplication. |
| 112 | pub fn init(allocator: Allocator) EnvMap { |
| 113 | return EnvMap{ .hash_map = HashMap.init(allocator) }; |
| 114 | } |
| 115 | |
| 116 | /// Free the backing storage of the map, as well as all |
| 117 | /// of the stored keys and values. |
| 118 | pub fn deinit(self: *EnvMap) void { |
| 119 | var it = self.hash_map.iterator(); |
| 120 | while (it.next()) |entry| { |
| 121 | self.free(entry.key_ptr.*); |
| 122 | self.free(entry.value_ptr.*); |
| 123 | } |
| 124 | |
| 125 | self.hash_map.deinit(); |
| 126 | } |
| 127 | |
| 128 | /// Same as `put` but the key and value become owned by the EnvMap rather |
| 129 | /// than being copied. |
| 130 | /// If `putMove` fails, the ownership of key and value does not transfer. |
| 131 | /// On Windows `key` must be a valid UTF-8 string. |
| 132 | pub fn putMove(self: *EnvMap, key: []u8, value: []u8) !void { |
| 133 | const get_or_put = try self.hash_map.getOrPut(key); |
| 134 | if (get_or_put.found_existing) { |
| 135 | self.free(get_or_put.key_ptr.*); |
| 136 | self.free(get_or_put.value_ptr.*); |
| 137 | get_or_put.key_ptr.* = key; |
| 138 | } |
| 139 | get_or_put.value_ptr.* = value; |
| 140 | } |
| 141 | |
| 142 | /// `key` and `value` are copied into the EnvMap. |
| 143 | /// On Windows `key` must be a valid UTF-8 string. |
| 144 | pub fn put(self: *EnvMap, key: []const u8, value: []const u8) !void { |
| 145 | const value_copy = try self.copy(value); |
| 146 | errdefer self.free(value_copy); |
| 147 | const get_or_put = try self.hash_map.getOrPut(key); |
| 148 | if (get_or_put.found_existing) { |
| 149 | self.free(get_or_put.value_ptr.*); |
| 150 | } else { |
| 151 | get_or_put.key_ptr.* = self.copy(key) catch |err| { |
| 152 | _ = self.hash_map.remove(key); |
| 153 | return err; |
| 154 | }; |
| 155 | } |
| 156 | get_or_put.value_ptr.* = value_copy; |
| 157 | } |
| 158 | |
| 159 | /// Find the address of the value associated with a key. |
| 160 | /// The returned pointer is invalidated if the map resizes. |
| 161 | /// On Windows `key` must be a valid UTF-8 string. |
| 162 | pub fn getPtr(self: EnvMap, key: []const u8) ?*[]const u8 { |
| 163 | return self.hash_map.getPtr(key); |
| 164 | } |
| 165 | |
| 166 | /// Return the map's copy of the value associated with |
| 167 | /// a key. The returned string is invalidated if this |
| 168 | /// key is removed from the map. |
| 169 | /// On Windows `key` must be a valid UTF-8 string. |
| 170 | pub fn get(self: EnvMap, key: []const u8) ?[]const u8 { |
| 171 | return self.hash_map.get(key); |
| 172 | } |
| 173 | |
| 174 | /// Removes the item from the map and frees its value. |
| 175 | /// This invalidates the value returned by get() for this key. |
| 176 | /// On Windows `key` must be a valid UTF-8 string. |
| 177 | pub fn remove(self: *EnvMap, key: []const u8) void { |
| 178 | const kv = self.hash_map.fetchRemove(key) orelse return; |
| 179 | self.free(kv.key); |
| 180 | self.free(kv.value); |
| 181 | } |
| 182 | |
| 183 | /// Returns the number of KV pairs stored in the map. |
| 184 | pub fn count(self: EnvMap) HashMap.Size { |
| 185 | return self.hash_map.count(); |
| 186 | } |
| 187 | |
| 188 | /// Returns an iterator over entries in the map. |
| 189 | pub fn iterator(self: *const EnvMap) HashMap.Iterator { |
| 190 | return self.hash_map.iterator(); |
| 191 | } |
| 192 | |
| 193 | fn free(self: EnvMap, value: []const u8) void { |
| 194 | self.hash_map.allocator.free(value); |
| 195 | } |
| 196 | |
| 197 | fn copy(self: EnvMap, value: []const u8) ![]u8 { |
| 198 | return self.hash_map.allocator.dupe(u8, value); |
| 199 | } |
| 200 | }; |
| 201 | |
| 202 | test "EnvMap" { |
| 203 | var env = EnvMap.init(testing.allocator); |
| 204 | defer env.deinit(); |
| 205 | |
| 206 | try env.put("SOMETHING_NEW", "hello"); |
| 207 | try testing.expectEqualStrings("hello", env.get("SOMETHING_NEW").?); |
| 208 | try testing.expectEqual(@as(EnvMap.Size, 1), env.count()); |
| 209 | |
| 210 | // overwrite |
| 211 | try env.put("SOMETHING_NEW", "something"); |
| 212 | try testing.expectEqualStrings("something", env.get("SOMETHING_NEW").?); |
| 213 | try testing.expectEqual(@as(EnvMap.Size, 1), env.count()); |
| 214 | |
| 215 | // a new longer name to test the Windows-specific conversion buffer |
| 216 | try env.put("SOMETHING_NEW_AND_LONGER", "1"); |
| 217 | try testing.expectEqualStrings("1", env.get("SOMETHING_NEW_AND_LONGER").?); |
| 218 | try testing.expectEqual(@as(EnvMap.Size, 2), env.count()); |
| 219 | |
| 220 | // case insensitivity on Windows only |
| 221 | if (builtin.os.tag == .windows) { |
| 222 | try testing.expectEqualStrings("1", env.get("something_New_aNd_LONGER").?); |
| 223 | } else { |
| 224 | try testing.expect(null == env.get("something_New_aNd_LONGER")); |
| 225 | } |
| 226 | |
| 227 | var it = env.iterator(); |
| 228 | var count: EnvMap.Size = 0; |
| 229 | while (it.next()) |entry| { |
| 230 | 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.*); |
| 231 | try testing.expect(is_an_expected_name); |
| 232 | count += 1; |
| 233 | } |
| 234 | try testing.expectEqual(@as(EnvMap.Size, 2), count); |
| 235 | |
| 236 | env.remove("SOMETHING_NEW"); |
| 237 | try testing.expect(env.get("SOMETHING_NEW") == null); |
| 238 | |
| 239 | try testing.expectEqual(@as(EnvMap.Size, 1), env.count()); |
| 240 | |
| 241 | // test Unicode case-insensitivity on Windows |
| 242 | if (builtin.os.tag == .windows) { |
| 243 | try env.put("КИРиллИЦА", "something else"); |
| 244 | try testing.expectEqualStrings("something else", env.get("кириллица").?); |
| 245 | } |
| 246 | } |
| 247 | |
| 248 | /// Returns a snapshot of the environment variables of the current process. |
| 249 | /// Any modifications to the resulting EnvMap will not be not reflected in the environment, and |
| 250 | /// likewise, any future modifications to the environment will not be reflected in the EnvMap. |
| 251 | /// Caller owns resulting `EnvMap` and should call its `deinit` fn when done. |
| 252 | pub fn getEnvMap(allocator: Allocator) !EnvMap { |
| 253 | var result = EnvMap.init(allocator); |
| 59 | 254 | errdefer result.deinit(); |
| 60 | 255 | |
| 61 | 256 | if (builtin.os.tag == .windows) { |
| ... | ... | @@ -65,6 +260,12 @@ pub fn getEnvMap(allocator: Allocator) !BufMap { |
| 65 | 260 | while (ptr[i] != 0) { |
| 66 | 261 | const key_start = i; |
| 67 | 262 | |
| 263 | // There are some special environment variables that start with =, |
| 264 | // so we need a special case to not treat = as a key/value separator |
| 265 | // if it's the first character. |
| 266 | // https://devblogs.microsoft.com/oldnewthing/20100506-00/?p=14133 |
| 267 | if (ptr[key_start] == '=') i += 1; |
| 268 | |
| 68 | 269 | while (ptr[i] != 0 and ptr[i] != '=') : (i += 1) {} |
| 69 | 270 | const key_w = ptr[key_start..i]; |
| 70 | 271 | const key = try std.unicode.utf16leToUtf8Alloc(allocator, key_w); |
| ... | ... | @@ -140,8 +341,8 @@ pub fn getEnvMap(allocator: Allocator) !BufMap { |
| 140 | 341 | } |
| 141 | 342 | } |
| 142 | 343 | |
| 143 | | test "os.getEnvMap" { |
| 144 | | var env = try getEnvMap(std.testing.allocator); |
| 344 | test "getEnvMap" { |
| 345 | var env = try getEnvMap(testing.allocator); |
| 145 | 346 | defer env.deinit(); |
| 146 | 347 | } |
| 147 | 348 | |
| ... | ... | @@ -985,7 +1186,7 @@ pub fn execv(allocator: mem.Allocator, argv: []const []const u8) ExecvError { |
| 985 | 1186 | pub fn execve( |
| 986 | 1187 | allocator: mem.Allocator, |
| 987 | 1188 | argv: []const []const u8, |
| 988 | | env_map: ?*const std.BufMap, |
| 1189 | env_map: ?*const EnvMap, |
| 989 | 1190 | ) ExecvError { |
| 990 | 1191 | if (!can_execv) @compileError("The target OS does not support execv"); |
| 991 | 1192 | |