Self

fn HuffmanDecoder(
    comptime alphabet_size: u16,
    comptime max_code_bits: u4,
    comptime lookup_bits: u4,
) type {
    const lookup_shift = max_code_bits - lookup_bits;
    const lookup_mask = (1 << lookup_bits) - 1;

    return struct {
        // lookup table code -> symbol
        // for values with code_bits == 0, symbol is the index of the first node in linked
        // if the index of the first node is 0xfff, it is an invalid code
        lookup: [1 << lookup_bits]Symbol = undefined,
        linked: if (lookup_bits == max_code_bits) void else [alphabet_size]struct {
            // sym.value is the next index in linked where the current index ends the chain
            // the actual symbol is this nodes's index
            sym: Symbol,
            code: u16,
        } = undefined,

        const Self = @This();

        fn reverseIdx(idx: usize) u16 {
            return @bitReverse(@as(@Int(.unsigned, lookup_bits), @intCast(idx)));
        }

        /// Generates symbols and lookup tables from list of code lens for each symbol.
        pub fn generate(self: *Self, lens: []const u4) !void {
            try checkCompleteness(lens);

            var buckets: [1 + @as(usize, max_code_bits)][alphabet_size]Symbol = undefined;
            var bucket_len: [buckets.len]u16 = @splat(0);
            for (0.., lens) |symbol, bits| {
                buckets[bits][bucket_len[bits]] = .{
                    .value = @intCast(symbol),
                    .code_bits = bits,
                };
                bucket_len[bits] += 1;
            }

            var code: u16 = 0;
            var idx: u16 = 0;
            for (1..lookup_bits + 1) |bits| {
                const inc = @as(u16, 1) << @intCast(max_code_bits - bits);
                for (buckets[bits][0..bucket_len[bits]]) |lookup_sym| {
                    const next_code = code + inc;
                    const next_idx = next_code >> lookup_shift;
                    for (idx..next_idx) |i| {
                        self.lookup[reverseIdx(i)] = lookup_sym;
                    }
                    code = next_code;
                    idx = next_idx;
                }
            }
            for (lookup_bits + 1..buckets.len) |bits| {
                const inc = @as(u16, 1) << @intCast(max_code_bits - bits);
                for (buckets[bits][0..bucket_len[bits]]) |linked_sym| {
                    const next_code = code + inc;
                    const next_idx = next_code >> lookup_shift;

                    const ri = reverseIdx(idx);
                    const next: Symbol = .{
                        .value = self.lookup[ri].value,
                        .code_bits = linked_sym.code_bits,
                    };
                    self.linked[linked_sym.value] = .{
                        .sym = next,
                        .code = @bitReverse(@as(@Int(.unsigned, max_code_bits), @intCast(code))),
                    };
                    self.lookup[ri] = .{ .value = linked_sym.value, .code_bits = 0 };

                    code = next_code;
                    idx = next_idx;
                }
            }

            // Invalid codes
            for (idx..self.lookup.len) |i| {
                self.lookup[reverseIdx(i)] = .{ .value = 0xfff, .code_bits = 0 };
            }
        }

        /// Given the list of code lengths check that it represents a canonical
        /// Huffman code for n symbols.
        ///
        /// Reference: https://github.com/madler/zlib/blob/5c42a230b7b468dff011f444161c0145b5efae59/contrib/puff/puff.c#L340
        fn checkCompleteness(lens: []const u4) !void {
            if (alphabet_size == 286)
                if (lens[256] == 0) return error.MissingEndOfBlockCode;

            var count: [@as(usize, max_code_bits) + 1]u16 = @splat(0);
            var max: usize = 0;
            for (lens) |n| {
                if (n == 0) continue;
                if (n > max) max = n;
                count[n] += 1;
            }
            if (max == 0) // empty tree
                return;

            // check for an over-subscribed or incomplete set of lengths
            var left: usize = 1; // one possible code of zero length
            for (1..count.len) |len| {
                left <<= 1; // one more bit, double codes left
                if (count[len] > left)
                    return error.OversubscribedHuffmanTree;
                left -= count[len]; // deduct count from possible codes
            }
            if (left > 0) { // left > 0 means incomplete
                // incomplete code ok only for single length 1 code
                if (max_code_bits > 7 and max == count[0] + count[1]) return;
                return error.IncompleteHuffmanTree;
            }
        }

        /// Finds symbol for lookup table code.
        pub fn find(self: *Self, code: u16) !Symbol {
            // try to find in lookup table
            const idx = code & lookup_mask;
            const sym = self.lookup[idx];
            if (sym.code_bits != 0) return sym;
            // if not use linked list of symbols with same prefix
            return self.findLinked(code, sym.value);
        }

        fn findLinked(self: *Self, code: u16, start: u16) !Symbol {
            if (start == 0xfff) return error.InvalidCode;
            if (lookup_bits == max_code_bits) unreachable;
            var pos = start;
            while (true) {
                const node = self.linked[pos];
                const shift = -%node.sym.code_bits;
                // compare code_bits number of upper bits
                if ((code ^ node.code) << shift == 0)
                    return .{ .value = @intCast(pos), .code_bits = node.sym.code_bits };
                pos = node.sym.value;
            }
        }
    };
}