Decode

pub const Decode = struct {
    lzma_decode: lzma.Decode,

    pub fn init(gpa: Allocator) !Decode {
        return .{ .lzma_decode = try lzma.Decode.init(gpa, .{ .lc = 0, .lp = 0, .pb = 0 }) };
    }

    pub fn deinit(self: *Decode, gpa: Allocator) void {
        self.lzma_decode.deinit(gpa);
        self.* = undefined;
    }

    /// Returns how many compressed bytes were consumed.
    pub fn decompress(d: *Decode, reader: *Reader, allocating: *Writer.Allocating) !u64 {
        const gpa = allocating.allocator;

        var accum = AccumBuffer.init(std.math.maxInt(usize));
        defer accum.deinit(gpa);

        var n_read: u64 = 0;

        while (true) {
            const status = try reader.takeByte();
            n_read += 1;

            switch (status) {
                0 => break,
                1 => n_read += try parseUncompressed(reader, allocating, &accum, true),
                2 => n_read += try parseUncompressed(reader, allocating, &accum, false),
                else => n_read += try d.parseLzma(reader, allocating, &accum, status),
            }
        }

        try accum.finish(&allocating.writer);
        return n_read;
    }

    fn parseLzma(
        d: *Decode,
        reader: *Reader,
        allocating: *Writer.Allocating,
        accum: *AccumBuffer,
        status: u8,
    ) !u64 {
        if (status & 0x80 == 0) return error.CorruptInput;

        const Reset = struct {
            dict: bool,
            state: bool,
            props: bool,
        };

        const reset: Reset = switch ((status >> 5) & 0x3) {
            0 => .{
                .dict = false,
                .state = false,
                .props = false,
            },
            1 => .{
                .dict = false,
                .state = true,
                .props = false,
            },
            2 => .{
                .dict = false,
                .state = true,
                .props = true,
            },
            3 => .{
                .dict = true,
                .state = true,
                .props = true,
            },
            else => unreachable,
        };

        var n_read: u64 = 0;

        const unpacked_size = blk: {
            var tmp: u64 = status & 0x1F;
            tmp <<= 16;
            tmp |= try reader.takeInt(u16, .big);
            n_read += 2;
            break :blk tmp + 1;
        };

        const packed_size = blk: {
            const tmp: u17 = try reader.takeInt(u16, .big);
            n_read += 2;
            break :blk tmp + 1;
        };

        if (reset.dict) try accum.reset(&allocating.writer);

        const ld = &d.lzma_decode;

        if (reset.state) {
            var new_props = ld.properties;

            if (reset.props) {
                var props = try reader.takeByte();
                n_read += 1;
                if (props >= 225) {
                    return error.CorruptInput;
                }

                const lc = @as(u4, @intCast(props % 9));
                props /= 9;
                const lp = @as(u3, @intCast(props % 5));
                props /= 5;
                const pb = @as(u3, @intCast(props));

                if (lc + lp > 4) {
                    return error.CorruptInput;
                }

                new_props = .{ .lc = lc, .lp = lp, .pb = pb };
            }

            try ld.resetState(allocating.allocator, new_props);
        }

        const expected_unpacked_size = accum.len + unpacked_size;
        const start_count = n_read;
        var range_decoder = try lzma.RangeDecoder.initCounting(reader, &n_read);

        while (true) {
            if (accum.len >= expected_unpacked_size) break;
            switch (try ld.process(reader, allocating, accum, &range_decoder, &n_read)) {
                .more => continue,
                .finished => break,
            }
        }
        if (accum.len != expected_unpacked_size) return error.DecompressedSizeMismatch;
        if (n_read - start_count != packed_size) return error.CompressedSizeMismatch;

        return n_read;
    }

    fn parseUncompressed(
        reader: *Reader,
        allocating: *Writer.Allocating,
        accum: *AccumBuffer,
        reset_dict: bool,
    ) !usize {
        const unpacked_size = @as(u17, try reader.takeInt(u16, .big)) + 1;

        if (reset_dict) try accum.reset(&allocating.writer);

        const gpa = allocating.allocator;

        for (0..unpacked_size) |_| {
            try accum.appendByte(gpa, try reader.takeByte());
        }
        return 2 + unpacked_size;
    }
}