Map

pub const Map = struct {
    array_hash_map: ArrayHashMap,
    allocator: Allocator,

    const ArrayHashMap = std.array_hash_map.Custom([]const u8, []const u8, EnvNameHashContext, false);

    pub const Size = usize;

    pub const EnvNameHashContext = struct {
        pub fn hash(self: @This(), s: []const u8) u32 {
            _ = self;
            switch (native_os) {
                else => return std.array_hash_map.hashString(s),
                .windows => {
                    var h = std.hash.Wyhash.init(0);
                    var it = unicode.Wtf8View.initUnchecked(s).iterator();
                    while (it.nextCodepoint()) |cp| {
                        const cp_upper = if (std.math.cast(u16, cp)) |wtf16|
                            std.os.windows.toUpperWtf16(wtf16)
                        else
                            cp;
                        h.update(&[_]u8{
                            @truncate(cp_upper >> 0),
                            @truncate(cp_upper >> 8),
                            @truncate(cp_upper >> 16),
                        });
                    }
                    return @truncate(h.final());
                },
            }
        }

        pub fn eql(self: @This(), a: []const u8, b: []const u8, b_index: usize) bool {
            _ = self;
            _ = b_index;
            return eqlKeys(a, b);
        }
    };
    fn eqlKeys(a: []const u8, b: []const u8) bool {
        return switch (native_os) {
            else => std.array_hash_map.eqlString(a, b),
            .windows => std.os.windows.eqlIgnoreCaseWtf8(a, b),
        };
    }

    pub fn validateKeyForPut(key: []const u8) bool {
        switch (native_os) {
            else => return key.len > 0 and mem.findAny(u8, key, &.{ 0, '=' }) == null,
            .windows => {
                if (!unicode.wtf8ValidateSlice(key)) return false;
                return key.len > 0 and key[0] != 0 and mem.findAnyPos(u8, key, 1, &.{ 0, '=' }) == null;
            },
        }
    }

    pub fn validateKeyForFetch(key: []const u8) bool {
        if (native_os == .windows and !unicode.wtf8ValidateSlice(key)) return false;
        return true;
    }

    /// Create a Map backed by a specific allocator.
    /// That allocator will be used for both backing allocations
    /// and string deduplication.
    pub fn init(allocator: Allocator) Map {
        return .{ .array_hash_map = .empty, .allocator = allocator };
    }

    /// Free the backing storage of the map, as well as all
    /// of the stored keys and values.
    pub fn deinit(self: *Map) void {
        const gpa = self.allocator;
        for (self.keys()) |key| gpa.free(key);
        for (self.values()) |value| gpa.free(value);
        self.array_hash_map.deinit(gpa);
        self.* = undefined;
    }

    pub fn keys(map: *const Map) [][]const u8 {
        return map.array_hash_map.keys();
    }

    pub fn values(map: *const Map) [][]const u8 {
        return map.array_hash_map.values();
    }

    pub fn putPosixBlock(map: *Map, view: PosixBlock.View) Allocator.Error!void {
        for (view.slice) |entry| {
            var entry_i: usize = 0;
            while (entry[entry_i] != 0 and entry[entry_i] != '=') : (entry_i += 1) {}
            const key = entry[0..entry_i];

            var end_i: usize = entry_i;
            while (entry[end_i] != 0) : (end_i += 1) {}
            const value = entry[entry_i + 1 .. end_i];

            try map.put(key, value);
        }
    }

    pub fn putWindowsBlock(map: *Map, view: WindowsBlock.View) Allocator.Error!void {
        var i: usize = 0;
        while (view.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 (view.ptr[key_start] == '=') i += 1;

            while (view.ptr[i] != 0 and view.ptr[i] != '=') : (i += 1) {}
            const key_w = view.ptr[key_start..i];
            const key = try unicode.wtf16LeToWtf8Alloc(map.allocator, key_w);
            errdefer map.allocator.free(key);

            if (view.ptr[i] == '=') i += 1;

            const value_start = i;
            while (view.ptr[i] != 0) : (i += 1) {}
            const value_w = view.ptr[value_start..i];
            const value = try unicode.wtf16LeToWtf8Alloc(map.allocator, value_w);
            errdefer map.allocator.free(value);

            i += 1; // skip over null byte

            try map.putMove(key, value);
        }
    }

    /// Same as `put` but the key and value become owned by the Map rather
    /// than being copied.
    /// If `putMove` fails, the ownership of key and value does not transfer.
    ///
    /// Asserts that `key` is valid:
    /// - It cannot contain a NUL (`'\x00') byte.
    /// - It must have a length > 0.
    /// - It cannot contain `=`, except on Windows where only the first code point is allowed to be `=`.
    /// - On Windows, it must be valid [WTF-8](https://wtf-8.codeberg.page/).
    pub fn putMove(self: *Map, key: []u8, value: []u8) Allocator.Error!void {
        assert(validateKeyForPut(key));
        const gpa = self.allocator;
        const get_or_put = try self.array_hash_map.getOrPut(gpa, key);
        if (get_or_put.found_existing) {
            gpa.free(get_or_put.key_ptr.*);
            gpa.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 Map.
    ///
    /// Asserts that `key` is valid:
    /// - It cannot contain a NUL (`'\x00') byte.
    /// - It must have a length > 0.
    /// - It cannot contain `=`, except on Windows where only the first code point is allowed to be `=`.
    /// - On Windows, it must be valid [WTF-8](https://wtf-8.codeberg.page/).
    pub fn put(self: *Map, key: []const u8, value: []const u8) Allocator.Error!void {
        assert(validateKeyForPut(key));
        const gpa = self.allocator;
        const value_copy = try gpa.dupe(u8, value);
        errdefer gpa.free(value_copy);
        const get_or_put = try self.array_hash_map.getOrPut(gpa, key);
        errdefer {
            if (!get_or_put.found_existing) assert(self.array_hash_map.pop() != null);
        }
        if (get_or_put.found_existing) {
            gpa.free(get_or_put.value_ptr.*);
        } else {
            get_or_put.key_ptr.* = try gpa.dupe(u8, key);
        }
        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, asserts that `key` is valid [WTF-8](https://wtf-8.codeberg.page/).
    pub fn getPtr(self: Map, key: []const u8) ?*[]const u8 {
        assert(validateKeyForFetch(key));
        return self.array_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, asserts that `key` is valid [WTF-8](https://wtf-8.codeberg.page/).
    pub fn get(self: Map, key: []const u8) ?[]const u8 {
        assert(validateKeyForFetch(key));
        return self.array_hash_map.get(key);
    }

    /// On Windows, asserts that `key` is valid [WTF-8](https://wtf-8.codeberg.page/).
    pub fn contains(m: *const Map, key: []const u8) bool {
        assert(validateKeyForFetch(key));
        return m.array_hash_map.contains(key);
    }

    /// If there is an entry with a matching key, it is deleted from the hash
    /// map. The entry is removed from the underlying array by swapping it with
    /// the last element.
    ///
    /// Returns true if an entry was removed, false otherwise.
    ///
    /// This invalidates the value returned by get() for this key.
    /// On Windows, asserts that `key` is valid [WTF-8](https://wtf-8.codeberg.page/).
    pub fn swapRemove(self: *Map, key: []const u8) bool {
        assert(validateKeyForFetch(key));
        const kv = self.array_hash_map.fetchSwapRemove(key) orelse return false;
        const gpa = self.allocator;
        gpa.free(kv.key);
        gpa.free(kv.value);
        return true;
    }

    /// If there is an entry with a matching key, it is deleted from the map.
    /// The entry is removed from the underlying array by shifting all elements
    /// forward, thereby maintaining the current ordering.
    ///
    /// Returns true if an entry was removed, false otherwise.
    ///
    /// This invalidates the value returned by get() for this key.
    /// On Windows, asserts that `key` is valid [WTF-8](https://wtf-8.codeberg.page/).
    pub fn orderedRemove(self: *Map, key: []const u8) bool {
        assert(validateKeyForFetch(key));
        const kv = self.array_hash_map.fetchOrderedRemove(key) orelse return false;
        const gpa = self.allocator;
        gpa.free(kv.key);
        gpa.free(kv.value);
        return true;
    }

    /// Returns the number of KV pairs stored in the map.
    pub fn count(self: Map) Size {
        return self.array_hash_map.count();
    }

    /// Returns an iterator over entries in the map.
    pub fn iterator(self: *const Map) ArrayHashMap.Iterator {
        return self.array_hash_map.iterator();
    }

    /// Returns a full copy of `em` allocated with `gpa`, which is not necessarily
    /// the same allocator used to allocate `em`.
    pub fn clone(m: *const Map, gpa: Allocator) Allocator.Error!Map {
        var new: Map = .init(gpa);
        errdefer new.deinit();
        try new.array_hash_map.ensureUnusedCapacity(gpa, m.array_hash_map.count());
        for (m.array_hash_map.keys(), m.array_hash_map.values()) |key, value| {
            try new.put(key, value);
        }
        return new;
    }

    /// Adds all the key-value pairs from `other` into this `m`.
    pub fn putAll(m: *Map, other: *const Map) Allocator.Error!void {
        const gpa = m.allocator;
        try m.array_hash_map.ensureUnusedCapacity(gpa, other.array_hash_map.count());
        const start = m.count();
        errdefer while (m.array_hash_map.count() > start) {
            const kv = m.array_hash_map.pop().?;
            gpa.free(kv.key);
            gpa.free(kv.value);
        };
        for (other.array_hash_map.keys(), other.array_hash_map.values()) |key, value| {
            try m.put(key, value);
        }
    }

    /// Set the length to zero, freeing all key and value memory, not freeing
    /// the allocation for the entries.
    pub fn clearRetainingCapacity(m: *Map) void {
        const gpa = m.allocator;
        for (m.array_hash_map.keys(), m.array_hash_map.values()) |k, v| {
            gpa.free(k);
            gpa.free(v);
        }
        m.array_hash_map.clearRetainingCapacity();
    }

    /// Creates a null-delimited environment variable block in the format
    /// expected by POSIX, from a hash map plus options.
    pub fn createPosixBlock(
        map: *const Map,
        gpa: Allocator,
        options: CreatePosixBlockOptions,
    ) Allocator.Error!PosixBlock {
        const ZigProgressAction = enum { nothing, edit, delete, add };
        const zig_progress_action: ZigProgressAction = action: {
            const fd = options.zig_progress_fd orelse break :action .nothing;
            const exists = map.contains("ZIG_PROGRESS");
            if (fd >= 0) {
                break :action if (exists) .edit else .add;
            } else {
                if (exists) break :action .delete;
            }
            break :action .nothing;
        };

        const envp = try gpa.allocSentinel(?[*:0]u8, len: {
            var len: usize = map.count();
            switch (zig_progress_action) {
                .add => len += 1,
                .delete => len -= 1,
                .nothing, .edit => {},
            }
            break :len len;
        }, null);
        var envp_len: usize = 0;
        errdefer {
            envp[envp_len] = null;
            PosixBlock.deinit(.{ .slice = envp[0..envp_len :null] }, gpa);
        }

        if (zig_progress_action == .add) {
            envp[envp_len] = try std.fmt.allocPrintSentinel(gpa, "ZIG_PROGRESS={d}", .{options.zig_progress_fd.?}, 0);
            envp_len += 1;
        }

        for (map.keys(), map.values()) |key, value| {
            if (mem.eql(u8, key, "ZIG_PROGRESS")) switch (zig_progress_action) {
                .add => unreachable,
                .delete => continue,
                .edit => {
                    envp[envp_len] = try std.fmt.allocPrintSentinel(gpa, "{s}={d}", .{
                        key, options.zig_progress_fd.?,
                    }, 0);
                    envp_len += 1;
                    continue;
                },
                .nothing => {},
            };

            envp[envp_len] = try std.fmt.allocPrintSentinel(gpa, "{s}={s}", .{ key, value }, 0);
            envp_len += 1;
        }

        assert(envp_len == envp.len);
        return .{ .slice = envp };
    }

    /// Caller owns result.
    pub fn createWindowsBlock(
        map: *const Map,
        gpa: Allocator,
        options: CreateWindowsBlockOptions,
    ) error{ OutOfMemory, InvalidWtf8 }!WindowsBlock {
        // count bytes needed
        const max_chars_needed = max_chars_needed: {
            var max_chars_needed: usize = "\x00".len;
            if (options.zig_progress_handle) |handle| if (handle != std.os.windows.INVALID_HANDLE_VALUE) {
                max_chars_needed += std.fmt.count("ZIG_PROGRESS={d}\x00", .{@intFromPtr(handle)});
            };
            for (map.keys(), map.values()) |key, value| {
                if (options.zig_progress_handle != null and eqlKeys(key, "ZIG_PROGRESS")) continue;
                max_chars_needed += key.len + "=".len + value.len + "\x00".len;
            }
            break :max_chars_needed @max("\x00\x00".len, max_chars_needed);
        };
        const block = try gpa.alloc(u16, max_chars_needed);
        errdefer gpa.free(block);

        var i: usize = 0;
        if (options.zig_progress_handle) |handle| if (handle != std.os.windows.INVALID_HANDLE_VALUE) {
            @memcpy(
                block[i..][0.."ZIG_PROGRESS=".len],
                &[_]u16{ 'Z', 'I', 'G', '_', 'P', 'R', 'O', 'G', 'R', 'E', 'S', 'S', '=' },
            );
            i += "ZIG_PROGRESS=".len;
            var value_buf: [std.fmt.count("{d}", .{std.math.maxInt(usize)})]u8 = undefined;
            const value = std.mem.print(&value_buf, "{d}", .{@intFromPtr(handle)}) catch unreachable;
            for (block[i..][0..value.len], value) |*r, v| r.* = v;
            i += value.len;
            block[i] = 0;
            i += 1;
        };
        for (map.keys(), map.values()) |key, value| {
            if (options.zig_progress_handle != null and eqlKeys(key, "ZIG_PROGRESS")) continue;
            i += try unicode.wtf8ToWtf16Le(block[i..], key);
            block[i] = '=';
            i += 1;
            i += try unicode.wtf8ToWtf16Le(block[i..], value);
            block[i] = 0;
            i += 1;
        }
        // An empty environment is a special case that requires a redundant
        // NUL terminator. CreateProcess will read the second code unit even
        // though theoretically the first should be enough to recognize that the
        // environment is empty (see https://nullprogram.com/blog/2023/08/23/)
        for (0..2) |_| {
            block[i] = 0;
            i += 1;
            if (i >= 2) break;
        } else unreachable;
        const reallocated = try gpa.realloc(block, i);
        return .{ .slice = reallocated[0 .. i - 1 :0] };
    }
}