huffman

const huffman = struct {
    const max_leafs = 286;
    const max_nodes = max_leafs * 2;

    const Node = packed struct(u32) {
        depth: u16,
        freq: u16,

        pub const Index = u16;

        /// `freq` is more significant than `depth`
        pub fn smaller(a: Node, b: Node) bool {
            return @as(u32, @bitCast(a)) < @as(u32, @bitCast(b));
        }
    };

    fn heapSiftDown(nodes: []Node, heap: []Node.Index, start: usize) void {
        var i = start;
        while (true) {
            var min = i;
            const l = i * 2 + 1;
            const r = l + 1;
            min = if (l < heap.len and nodes[heap[l]].smaller(nodes[heap[min]])) l else min;
            min = if (r < heap.len and nodes[heap[r]].smaller(nodes[heap[min]])) r else min;
            if (i == min) break;
            mem.swap(Node.Index, &heap[i], &heap[min]);
            i = min;
        }
    }

    fn heapRemoveRoot(nodes: []Node, heap: []Node.Index) void {
        heap[0] = heap[heap.len - 1];
        heapSiftDown(nodes, heap[0 .. heap.len - 1], 0);
    }

    /// Returns the total bits to encode `freqs` followed by the index of the last non-zero bits.
    /// For `freqs[i]` == 0, `out_codes[i]` will be undefined.
    /// It is asserted `out_bits` is zero-filled.
    /// It is asserted `out_bits.len` is at least a length of
    /// one if ncomplete trees are allowed and two otherwise.
    pub fn build(
        freqs: []const u16,
        out_codes: []u16,
        out_bits: []u4,
        max_bits: u4,
        incomplete_allowed: bool,
    ) struct { u32, u16 } {
        assert(out_codes.len - 1 >= @intFromBool(!incomplete_allowed));
        // freqs and out_codes are in the loop to assert they are all the same length
        for (freqs, out_codes, out_bits) |_, _, n| assert(n == 0);
        assert(out_codes.len <= @as(u16, 1) << max_bits);

        // Indexes 0..freqs are leafs, indexes max_leafs.. are internal nodes.
        var tree_nodes: [max_nodes]Node = undefined;
        var tree_parent_nodes: [max_nodes]Node.Index = undefined;
        var nodes_end: u16 = max_leafs;
        // Dual-purpose buffer. Nodes are ordered by least frequency or when equal, least depth.
        // The start is a min heap of level-zero nodes.
        // The end is a sorted buffer of nodes with the greatest first.
        var node_buf: [max_nodes]Node.Index = undefined;
        var heap_end: u16 = 0;
        var sorted_start: u16 = node_buf.len;

        for (0.., freqs) |n, freq| {
            tree_nodes[n] = .{ .freq = freq, .depth = 0 };
            node_buf[heap_end] = @intCast(n);
            heap_end += @intFromBool(freq != 0);
        }

        // There must be at least one code at minimum,
        node_buf[heap_end] = 0;
        heap_end += @intFromBool(heap_end == 0);
        // and at least two if incomplete must be avoided.
        if (heap_end == 1 and incomplete_allowed) {
            @branchHint(.unlikely); // LLVM 21 optimizes this branch as the more likely without

            // Codes must have at least one-bit, so this is a special case.
            out_bits[node_buf[0]] = 1;
            out_codes[node_buf[0]] = 0;
            return .{ freqs[node_buf[0]], node_buf[0] };
        }
        const last_nonzero = @max(node_buf[heap_end - 1], 1); // For heap_end > 1, last is not be 0
        node_buf[heap_end] = @intFromBool(node_buf[0] == 0);
        heap_end += @intFromBool(heap_end == 1);

        // Heapify the array of frequencies
        const heapify_final = heap_end - 1;
        const heapify_start = (heapify_final - 1) / 2; // Parent of final node
        var heapify_i = heapify_start;
        while (true) {
            heapSiftDown(&tree_nodes, node_buf[0..heap_end], heapify_i);
            if (heapify_i == 0) break;
            heapify_i -= 1;
        }

        // Build optimal tree. `max_bits` is not enforced yet.
        while (heap_end > 1) {
            const a = node_buf[0];
            heapRemoveRoot(&tree_nodes, node_buf[0..heap_end]);
            heap_end -= 1;
            const b = node_buf[0];

            sorted_start -= 2;
            node_buf[sorted_start..][0..2].* = .{ b, a };

            tree_nodes[nodes_end] = .{
                .freq = tree_nodes[a].freq + tree_nodes[b].freq,
                .depth = @max(tree_nodes[a].depth, tree_nodes[b].depth) + 1,
            };
            defer nodes_end += 1;
            tree_parent_nodes[a] = nodes_end;
            tree_parent_nodes[b] = nodes_end;

            node_buf[0] = nodes_end;
            heapSiftDown(&tree_nodes, node_buf[0..heap_end], 0);
        }
        sorted_start -= 1;
        node_buf[sorted_start] = node_buf[0];

        var bit_counts: [16]u16 = @splat(0);
        buildBits(out_bits, &bit_counts, &tree_parent_nodes, node_buf[sorted_start..], max_bits);
        return .{ buildValues(freqs, out_codes, out_bits, bit_counts), last_nonzero };
    }

    fn buildBits(
        out_bits: []u4,
        bit_counts: *[16]u16,
        parent_nodes: *[max_nodes]Node.Index,
        sorted: []Node.Index,
        max_bits: u4,
    ) void {
        var internal_node_bits: [max_nodes - max_leafs]u4 = undefined;
        var overflowed: u16 = 0;

        internal_node_bits[sorted[0] - max_leafs] = 0; // root
        for (sorted[1..]) |i| {
            const parent_bits = internal_node_bits[parent_nodes[i] - max_leafs];
            overflowed += @intFromBool(parent_bits == max_bits);
            const bits = parent_bits + @intFromBool(parent_bits != max_bits);
            bit_counts[bits] += @intFromBool(i < max_leafs);
            (if (i >= max_leafs) &internal_node_bits[i - max_leafs] else &out_bits[i]).* = bits;
        }

        if (overflowed == 0) {
            @branchHint(.likely);
            return;
        }

        outer: while (true) {
            var deepest: u4 = max_bits - 1;
            while (bit_counts[deepest] == 0) deepest -= 1;
            while (overflowed != 0) {
                // Insert an internal node under the leaf and move an overflow as its sibling
                bit_counts[deepest] -= 1;
                bit_counts[deepest + 1] += 2;
                // Only overflow moved. Its sibling's depth is one less, however is still >= depth.
                bit_counts[max_bits] -= 1;
                overflowed -= 2;

                if (overflowed == 0) break :outer;
                deepest += 1;
                if (deepest == max_bits) continue :outer;
            }
        }

        // Reassign bit lengths
        assert(bit_counts[0] == 0);
        var i: usize = 0;
        for (1.., bit_counts[1..]) |bits, all| {
            var remaining = all;
            while (remaining != 0) {
                defer i += 1;
                if (sorted[i] >= max_leafs) continue;
                out_bits[sorted[i]] = @intCast(bits);
                remaining -= 1;
            }
        }
        assert(for (sorted[i..]) |n| { // all leafs consumed
            if (n < max_leafs) break false;
        } else true);
    }

    fn buildValues(freqs: []const u16, out_codes: []u16, bits: []u4, bit_counts: [16]u16) u32 {
        var code: u16 = 0;
        var base: [16]u16 = undefined;
        assert(bit_counts[0] == 0);
        for (bit_counts[1..], base[1..]) |c, *b| {
            b.* = code;
            code +%= c;
            code <<= 1;
        }
        var freq_sums: [16]u16 = @splat(0);
        for (out_codes, bits, freqs) |*c, b, f| {
            c.* = @bitReverse(base[b]) >> -%b;
            base[b] += 1; // For `b == 0` this is fine since v is specified to be undefined.
            freq_sums[b] += f;
        }
        return @reduce(.Add, @as(@Vector(16, u32), freq_sums) * std.simd.iota(u32, 16));
    }

    test build {
        var codes: [8]u16 = undefined;
        var bits: [8]u4 = undefined;

        const regular_freqs: [8]u16 = .{ 1, 1, 0, 8, 8, 0, 2, 4 };
        // The optimal tree for the above frequencies is
        // 4             1   1
        //                \ /
        // 3           2   #
        //              \ /
        // 2   8   8 4   #
        //      \ /   \ /
        // 1     #     #
        //        \   /
        // 0        #
        bits = @splat(0);
        var n, var lnz = build(&regular_freqs, &codes, &bits, 15, true);
        codes[2] = 0;
        codes[5] = 0;
        try std.testing.expectEqualSlices(u4, &.{ 4, 4, 0, 2, 2, 0, 3, 2 }, &bits);
        try std.testing.expectEqualSlices(u16, &.{
            0b0111, 0b1111, 0, 0b00, 0b10, 0, 0b011, 0b01,
        }, &codes);
        try std.testing.expectEqual(54, n);
        try std.testing.expectEqual(7, lnz);
        // When constrained to 3 bits, it becomes
        // 3        1   1 2   4
        //           \ /   \ /
        // 2   8   8  #     #
        //      \ /    \   /
        // 1     #       #
        //        \     /
        // 0         #
        bits = @splat(0);
        n, lnz = build(&regular_freqs, &codes, &bits, 3, true);
        codes[2] = 0;
        codes[5] = 0;
        try std.testing.expectEqualSlices(u4, &.{ 3, 3, 0, 2, 2, 0, 3, 3 }, &bits);
        try std.testing.expectEqualSlices(u16, &.{
            0b001, 0b101, 0, 0b00, 0b10, 0, 0b011, 0b111,
        }, &codes);
        try std.testing.expectEqual(56, n);
        try std.testing.expectEqual(7, lnz);

        // Empty tree. At least one code should be present
        bits = @splat(0);
        n, lnz = build(&.{ 0, 0 }, codes[0..2], bits[0..2], 15, true);
        try std.testing.expectEqualSlices(u4, &.{ 1, 0 }, bits[0..2]);
        try std.testing.expectEqual(0b0, codes[0]);
        try std.testing.expectEqual(0, n);
        try std.testing.expectEqual(0, lnz);

        // Check all incompletable frequencies are completed
        for ([_][2]u16{ .{ 0, 0 }, .{ 0, 1 }, .{ 1, 0 } }) |incomplete| {
            // Empty tree. Both codes should be present to prevent incomplete trees
            bits = @splat(0);
            n, lnz = build(&incomplete, codes[0..2], bits[0..2], 15, false);
            try std.testing.expectEqualSlices(u4, &.{ 1, 1 }, bits[0..2]);
            try std.testing.expectEqualSlices(u16, &.{ 0b0, 0b1 }, codes[0..2]);
            try std.testing.expectEqual(incomplete[0] + incomplete[1], n);
            try std.testing.expectEqual(1, lnz);
        }

        try std.testing.fuzz({}, checkFuzzedBuildFreqs, .{});
    }

    fn checkFuzzedBuildFreqs(_: void, smith: *std.testing.Smith) !void {
        @disableInstrumentation();
        var freqs_limit: u16 = 65535;
        var freqs_buf: [max_leafs]u16 = undefined;
        var nfreqs: u15 = 0;

        const incomplete_allowed = smith.value(bool);
        while (nfreqs < @as(u8, @intFromBool(!incomplete_allowed)) + 1 or
            nfreqs != freqs_buf.len and freqs_limit != 0 and
                smith.eosWeightedSimple(15, 1))
        {
            const f = smith.valueWeighted(u16, &.{
                .rangeAtMost(u16, 0, @min(31, freqs_limit), @max(freqs_limit, 1)),
                .rangeAtMost(u16, 0, freqs_limit, 1),
            });
            freqs_buf[nfreqs] = f;
            freqs_limit -= f;
            nfreqs += 1;
        }

        var codes_buf: [max_leafs]u16 = undefined;
        var bits_buf: [max_leafs]u4 = @splat(0);
        const max_bits = smith.valueRangeAtMost(u4, math.log2_int_ceil(u15, nfreqs), 15);
        const total_bits, const last_nonzero = build(
            freqs_buf[0..nfreqs],
            codes_buf[0..nfreqs],
            bits_buf[0..nfreqs],
            max_bits,
            incomplete_allowed,
        );

        var has_bitlen_one: bool = false;
        var expected_total_bits: u32 = 0;
        var expected_last_nonzero: ?u16 = null;
        var weighted_sum: u32 = 0;
        for (freqs_buf[0..nfreqs], bits_buf[0..nfreqs], 0..) |f, nb, i| {
            has_bitlen_one = has_bitlen_one or nb == 1;
            weighted_sum += @shlExact(@as(u16, 1), 15 - nb) & ((1 << 15) - 1);
            expected_total_bits += @as(u32, f) * nb;
            if (nb != 0) expected_last_nonzero = @intCast(i);
        }

        errdefer std.log.err(
            \\ incomplete_allowed: {}
            \\ max_bits: {}
            \\ freqs: {any}
            \\ bits: {any}
            \\ # freqs: {}
            \\ weighted sum: {}
            \\ has_bitlen_one: {}
            \\ expected/actual total bits: {}/{}
            \\ expected/actual last nonzero: {?}/{}
        ++ "\n", .{
            incomplete_allowed,
            max_bits,
            freqs_buf[0..nfreqs],
            bits_buf[0..nfreqs],
            nfreqs,
            weighted_sum,
            has_bitlen_one,
            expected_total_bits,
            total_bits,
            expected_last_nonzero,
            last_nonzero,
        });

        try std.testing.expectEqual(expected_total_bits, total_bits);
        try std.testing.expectEqual(expected_last_nonzero, last_nonzero);
        if (weighted_sum > 1 << 15)
            return error.OversubscribedHuffmanTree;
        if (weighted_sum < 1 << 15 and
            !(incomplete_allowed and has_bitlen_one and weighted_sum == 1 << 14))
            return error.IncompleteHuffmanTree;
    }
}