struct Decode [src]

Fields

properties: Properties
literal_probs: Vec2d
pos_slot_decoder: [4]BitTree(6)
align_decoder: BitTree(4)
pos_decoders: [115]u16
is_match: [192]u16
is_rep: [12]u16
is_rep_g0: [12]u16
is_rep_g1: [12]u16
is_rep_g2: [12]u16
is_rep_0long: [192]u16
state: usize
rep: [4]usize
len_decoder: LenDecoder
rep_len_decoder: LenDecoder
Members

BitTree (Type Function)
CircularBuffer (struct)
deinit (Function)
init (Function)
LenDecoder (struct)
Options (struct)
Params (struct)
process (Function)
Properties (struct)
resetState (Function)
UnpackedSize (union)
Vec2d (struct)
Source

  pub const Decode = struct {
    properties: Properties,
    literal_probs: Vec2d,
    pos_slot_decoder: [4]BitTree(6),
    align_decoder: BitTree(4),
    pos_decoders: [115]u16,
    is_match: [192]u16,
    is_rep: [12]u16,
    is_rep_g0: [12]u16,
    is_rep_g1: [12]u16,
    is_rep_g2: [12]u16,
    is_rep_0long: [192]u16,
    state: usize,
    rep: [4]usize,
    len_decoder: LenDecoder,
    rep_len_decoder: LenDecoder,

    pub fn init(gpa: Allocator, properties: Properties) !Decode {
        return .{
            .properties = properties,
            .literal_probs = try Vec2d.init(gpa, 0x400, @as(usize, 1) << (properties.lc + properties.lp), 0x300),
            .pos_slot_decoder = @splat(.{}),
            .align_decoder = .{},
            .pos_decoders = @splat(0x400),
            .is_match = @splat(0x400),
            .is_rep = @splat(0x400),
            .is_rep_g0 = @splat(0x400),
            .is_rep_g1 = @splat(0x400),
            .is_rep_g2 = @splat(0x400),
            .is_rep_0long = @splat(0x400),
            .state = 0,
            .rep = @splat(0),
            .len_decoder = .{},
            .rep_len_decoder = .{},
        };
    }

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

    pub fn resetState(self: *Decode, gpa: Allocator, new_props: Properties) !void {
        new_props.validate();
        if (self.properties.lc + self.properties.lp == new_props.lc + new_props.lp) {
            self.literal_probs.fill(0x400);
        } else {
            self.literal_probs.deinit(gpa);
            self.literal_probs = try Vec2d.init(gpa, 0x400, @as(usize, 1) << (new_props.lc + new_props.lp), 0x300);
        }

        self.properties = new_props;
        for (&self.pos_slot_decoder) |*t| t.reset();
        self.align_decoder.reset();
        self.pos_decoders = @splat(0x400);
        self.is_match = @splat(0x400);
        self.is_rep = @splat(0x400);
        self.is_rep_g0 = @splat(0x400);
        self.is_rep_g1 = @splat(0x400);
        self.is_rep_g2 = @splat(0x400);
        self.is_rep_0long = @splat(0x400);
        self.state = 0;
        self.rep = @splat(0);
        self.len_decoder.reset();
        self.rep_len_decoder.reset();
    }

    pub fn process(
        self: *Decode,
        reader: *Reader,
        allocating: *Writer.Allocating,
        /// `CircularBuffer` or `std.compress.lzma2.AccumBuffer`.
        buffer: anytype,
        decoder: *RangeDecoder,
        n_read: *u64,
    ) !ProcessingStatus {
        const gpa = allocating.allocator;
        const writer = &allocating.writer;
        const pos_state = buffer.len & ((@as(usize, 1) << self.properties.pb) - 1);

        if (!try decoder.decodeBit(reader, &self.is_match[(self.state << 4) + pos_state], n_read)) {
            const byte: u8 = try self.decodeLiteral(reader, buffer, decoder, n_read);

            try buffer.appendLiteral(gpa, byte, writer);

            self.state = if (self.state < 4)
                0
            else if (self.state < 10)
                self.state - 3
            else
                self.state - 6;
            return .more;
        }

        var len: usize = undefined;
        if (try decoder.decodeBit(reader, &self.is_rep[self.state], n_read)) {
            if (!try decoder.decodeBit(reader, &self.is_rep_g0[self.state], n_read)) {
                if (!try decoder.decodeBit(reader, &self.is_rep_0long[(self.state << 4) + pos_state], n_read)) {
                    self.state = if (self.state < 7) 9 else 11;
                    const dist = self.rep[0] + 1;
                    try buffer.appendLz(gpa, 1, dist, writer);
                    return .more;
                }
            } else {
                const idx: usize = if (!try decoder.decodeBit(reader, &self.is_rep_g1[self.state], n_read))
                    1
                else if (!try decoder.decodeBit(reader, &self.is_rep_g2[self.state], n_read))
                    2
                else
                    3;
                const dist = self.rep[idx];
                var i = idx;
                while (i > 0) : (i -= 1) {
                    self.rep[i] = self.rep[i - 1];
                }
                self.rep[0] = dist;
            }

            len = try self.rep_len_decoder.decode(reader, decoder, pos_state, n_read);

            self.state = if (self.state < 7) 8 else 11;
        } else {
            self.rep[3] = self.rep[2];
            self.rep[2] = self.rep[1];
            self.rep[1] = self.rep[0];

            len = try self.len_decoder.decode(reader, decoder, pos_state, n_read);

            self.state = if (self.state < 7) 7 else 10;

            const rep_0 = try self.decodeDistance(reader, decoder, len, n_read);

            self.rep[0] = rep_0;
            if (self.rep[0] == 0xFFFF_FFFF) {
                if (decoder.isFinished()) {
                    return .finished;
                }
                return error.CorruptInput;
            }
        }

        len += 2;

        const dist = self.rep[0] + 1;
        try buffer.appendLz(gpa, len, dist, writer);

        return .more;
    }

    fn decodeLiteral(
        self: *Decode,
        reader: *Reader,
        /// `CircularBuffer` or `std.compress.lzma2.AccumBuffer`.
        buffer: anytype,
        decoder: *RangeDecoder,
        n_read: *u64,
    ) !u8 {
        const def_prev_byte = 0;
        const prev_byte = @as(usize, buffer.lastOr(def_prev_byte));

        var result: usize = 1;
        const lit_state = ((buffer.len & ((@as(usize, 1) << self.properties.lp) - 1)) << self.properties.lc) +
            (prev_byte >> (8 - self.properties.lc));
        const probs = try self.literal_probs.get(lit_state);

        if (self.state >= 7) {
            var match_byte = @as(usize, try buffer.lastN(self.rep[0] + 1));

            while (result < 0x100) {
                const match_bit = (match_byte >> 7) & 1;
                match_byte <<= 1;
                const bit = @intFromBool(try decoder.decodeBit(
                    reader,
                    &probs[((@as(usize, 1) + match_bit) << 8) + result],
                    n_read,
                ));
                result = (result << 1) ^ bit;
                if (match_bit != bit) {
                    break;
                }
            }
        }

        while (result < 0x100) {
            result = (result << 1) ^ @intFromBool(try decoder.decodeBit(reader, &probs[result], n_read));
        }

        return @truncate(result - 0x100);
    }

    fn decodeDistance(
        self: *Decode,
        reader: *Reader,
        decoder: *RangeDecoder,
        length: usize,
        n_read: *u64,
    ) !usize {
        const len_state = if (length > 3) 3 else length;

        const pos_slot: usize = try self.pos_slot_decoder[len_state].parse(reader, decoder, n_read);
        if (pos_slot < 4) return pos_slot;

        const num_direct_bits = @as(u5, @intCast((pos_slot >> 1) - 1));
        var result = (2 ^ (pos_slot & 1)) << num_direct_bits;

        if (pos_slot < 14) {
            result += try decoder.parseReverseBitTree(
                reader,
                num_direct_bits,
                &self.pos_decoders,
                result - pos_slot,
                n_read,
            );
        } else {
            result += @as(usize, try decoder.get(reader, num_direct_bits - 4, n_read)) << 4;
            result += try self.align_decoder.parseReverse(reader, decoder, n_read);
        }

        return result;
    }

    /// A circular buffer for LZ sequences
    pub const CircularBuffer = struct {
        /// Circular buffer
        buf: ArrayList(u8),
        /// Length of the buffer
        dict_size: usize,
        /// Buffer memory limit
        mem_limit: usize,
        /// Current position
        cursor: usize,
        /// Total number of bytes sent through the buffer
        len: usize,

        pub fn init(dict_size: usize, mem_limit: usize) CircularBuffer {
            return .{
                .buf = .{},
                .dict_size = dict_size,
                .mem_limit = mem_limit,
                .cursor = 0,
                .len = 0,
            };
        }

        pub fn get(self: CircularBuffer, index: usize) u8 {
            return if (0 <= index and index < self.buf.items.len) self.buf.items[index] else 0;
        }

        pub fn set(self: *CircularBuffer, gpa: Allocator, index: usize, value: u8) !void {
            if (index >= self.mem_limit) {
                return error.CorruptInput;
            }
            try self.buf.ensureTotalCapacity(gpa, index + 1);
            while (self.buf.items.len < index) {
                self.buf.appendAssumeCapacity(0);
            }
            self.buf.appendAssumeCapacity(value);
        }

        /// Retrieve the last byte or return a default
        pub fn lastOr(self: CircularBuffer, lit: u8) u8 {
            return if (self.len == 0)
                lit
            else
                self.get((self.dict_size + self.cursor - 1) % self.dict_size);
        }

        /// Retrieve the n-th last byte
        pub fn lastN(self: CircularBuffer, dist: usize) !u8 {
            if (dist > self.dict_size or dist > self.len) {
                return error.CorruptInput;
            }

            const offset = (self.dict_size + self.cursor - dist) % self.dict_size;
            return self.get(offset);
        }

        /// Append a literal
        pub fn appendLiteral(
            self: *CircularBuffer,
            gpa: Allocator,
            lit: u8,
            writer: *Writer,
        ) !void {
            try self.set(gpa, self.cursor, lit);
            self.cursor += 1;
            self.len += 1;

            // Flush the circular buffer to the output
            if (self.cursor == self.dict_size) {
                try writer.writeAll(self.buf.items);
                self.cursor = 0;
            }
        }

        /// Fetch an LZ sequence (length, distance) from inside the buffer
        pub fn appendLz(
            self: *CircularBuffer,
            gpa: Allocator,
            len: usize,
            dist: usize,
            writer: *Writer,
        ) !void {
            if (dist > self.dict_size or dist > self.len) {
                return error.CorruptInput;
            }

            var offset = (self.dict_size + self.cursor - dist) % self.dict_size;
            var i: usize = 0;
            while (i < len) : (i += 1) {
                const x = self.get(offset);
                try self.appendLiteral(gpa, x, writer);
                offset += 1;
                if (offset == self.dict_size) {
                    offset = 0;
                }
            }
        }

        pub fn finish(self: *CircularBuffer, writer: *Writer) !void {
            if (self.cursor > 0) {
                try writer.writeAll(self.buf.items[0..self.cursor]);
                self.cursor = 0;
            }
        }

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

    pub fn BitTree(comptime num_bits: usize) type {
        return struct {
            probs: [1 << num_bits]u16 = @splat(0x400),

            pub fn parse(self: *@This(), reader: *Reader, decoder: *RangeDecoder, n_read: *u64) !u32 {
                return decoder.parseBitTree(reader, num_bits, &self.probs, n_read);
            }

            pub fn parseReverse(
                self: *@This(),
                reader: *Reader,
                decoder: *RangeDecoder,
                n_read: *u64,
            ) !u32 {
                return decoder.parseReverseBitTree(reader, num_bits, &self.probs, 0, n_read);
            }

            pub fn reset(self: *@This()) void {
                @memset(&self.probs, 0x400);
            }
        };
    }

    pub const LenDecoder = struct {
        choice: u16 = 0x400,
        choice2: u16 = 0x400,
        low_coder: [16]BitTree(3) = @splat(.{}),
        mid_coder: [16]BitTree(3) = @splat(.{}),
        high_coder: BitTree(8) = .{},

        pub fn decode(
            self: *LenDecoder,
            reader: *Reader,
            decoder: *RangeDecoder,
            pos_state: usize,
            n_read: *u64,
        ) !usize {
            if (!try decoder.decodeBit(reader, &self.choice, n_read)) {
                return @as(usize, try self.low_coder[pos_state].parse(reader, decoder, n_read));
            } else if (!try decoder.decodeBit(reader, &self.choice2, n_read)) {
                return @as(usize, try self.mid_coder[pos_state].parse(reader, decoder, n_read)) + 8;
            } else {
                return @as(usize, try self.high_coder.parse(reader, decoder, n_read)) + 16;
            }
        }

        pub fn reset(self: *LenDecoder) void {
            self.choice = 0x400;
            self.choice2 = 0x400;
            for (&self.low_coder) |*t| t.reset();
            for (&self.mid_coder) |*t| t.reset();
            self.high_coder.reset();
        }
    };

    pub const Vec2d = struct {
        data: []u16,
        cols: usize,

        pub fn init(gpa: Allocator, value: u16, w: usize, h: usize) !Vec2d {
            const len = try math.mul(usize, w, h);
            const data = try gpa.alloc(u16, len);
            @memset(data, value);
            return .{
                .data = data,
                .cols = h,
            };
        }

        pub fn deinit(v: *Vec2d, gpa: Allocator) void {
            gpa.free(v.data);
            v.* = undefined;
        }

        pub fn fill(v: *Vec2d, value: u16) void {
            @memset(v.data, value);
        }

        fn get(v: Vec2d, row: usize) ![]u16 {
            const start_row = try math.mul(usize, row, v.cols);
            const end_row = try math.add(usize, start_row, v.cols);
            return v.data[start_row..end_row];
        }
    };

    pub const Options = struct {
        unpacked_size: UnpackedSize = .read_from_header,
        mem_limit: ?usize = null,
        allow_incomplete: bool = false,
    };

    pub const UnpackedSize = union(enum) {
        read_from_header,
        read_header_but_use_provided: ?u64,
        use_provided: ?u64,
    };

    const ProcessingStatus = enum {
        more,
        finished,
    };

    pub const Properties = struct {
        lc: u4,
        lp: u3,
        pb: u3,

        fn validate(self: Properties) void {
            assert(self.lc <= 8);
            assert(self.lp <= 4);
            assert(self.pb <= 4);
        }
    };

    pub const Params = struct {
        properties: Properties,
        dict_size: u32,
        unpacked_size: ?u64,

        pub fn readHeader(reader: *Reader, options: Options) !Params {
            var props = try reader.takeByte();
            if (props >= 225) return error.CorruptInput;

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

            const dict_size_provided = try reader.takeInt(u32, .little);
            const dict_size = @max(0x1000, dict_size_provided);

            const unpacked_size = switch (options.unpacked_size) {
                .read_from_header => blk: {
                    const unpacked_size_provided = try reader.takeInt(u64, .little);
                    const marker_mandatory = unpacked_size_provided == 0xFFFF_FFFF_FFFF_FFFF;
                    break :blk if (marker_mandatory) null else unpacked_size_provided;
                },
                .read_header_but_use_provided => |x| blk: {
                    _ = try reader.takeInt(u64, .little);
                    break :blk x;
                },
                .use_provided => |x| x,
            };

            return .{
                .properties = .{ .lc = lc, .lp = lp, .pb = pb },
                .dict_size = dict_size,
                .unpacked_size = unpacked_size,
            };
        }
    };
}  