Simple LZ4 block enc/dec in 100 LOC

lz4_block.zig

const std = @import("std");

fn writeExtended(w: *std.Io.Writer, value: u32) !void {
    const num_max_bytes = value / 0xff;
    try w.splatByteAll(0xff, num_max_bytes);
    try w.writeByte(@intCast(value - (num_max_bytes * 0xff)));
}

fn readExtended(r: *std.Io.Reader) !u32 {
    var value: u32 = 0;
    while (true) {
        const n = try r.takeByte();
        value +|= n; // saturate to avoid panic on bad inputs
        if (n < 0xff) return value;
    }
}

fn writeToken(w: *std.Io.Writer, literals: []const u8, offset: u16, matches: u32) !void {
    const matches_sub_4 = matches -| 4;
    const header = (@as(u8, @min(literals.len, 0xf)) << 4) | @min(matches_sub_4, 0xf);
    try w.writeByte(@intCast(header));
    
    if (literals.len >= 0xf) try writeExtended(w, @intCast(literals.len - 0xf));
    try w.writeAll(literals);

    if (matches > 0) {
        std.debug.assert(matches >= 4);
        try w.writeInt(u16, offset, .little);
        if (matches_sub_4 >= 0xf) try writeExtended(w, matches_sub_4 - 0xf);
    }
}

fn readToken(r: *std.Io.Reader, literals_buf: []u8) !struct{ literals: u32, offset: u16, matches: u32 } {
    const header = try r.takeByte();
    
    var literals: u32 = header >> 4;
    if (literals >= 0xf) literals += try readExtended(r);

    if (literals > literals_buf.len) return error.OutOfSpace;
    try r.readSliceAll(literals_buf[0..literals]);

    // read first byte manually to differentiate last-literals token from truncated stream.
    const offset_byte_0 = r.takeByte() catch |err| switch (err) {
        error.EndOfStream => return .{ .literals = literals, .offset = 0, .matches = 0 },
        else => |e| return e,
    };
    const offset = (@as(u16, try r.takeByte()) << 8) | offset_byte_0;

    var matches_sub_4: u32 = header & 0xf;
    if (matches_sub_4 >= 0xf) matches_sub_4 += try readExtended(r);
    return .{ .literals = literals, .offset = offset, .matches = matches_sub_4 + 4 };
}

fn compressBlock(w: *std.Io.Writer, src: []const u8) !void {
    var table: [4096]u32 = @splat(0); // hash table of previous pos stored as +1 to diff from 0
    var literals_start_pos: u32 = 0; 

    var pos: u32 = 0;
    while (pos + 4 <= src.len -| (12 - 4)) : (pos += 1) { // spec: last match must start 12b before end of block
        const word: u32 = @bitCast(src[pos..][0..4].*);
        const hash = (word *% 0x9e3779b9) >> @intCast(32 - @ctz(table.len)); // fib hash

        const match_pos = table[hash] -% 1; // if empty, match_pos = ~0
        table[hash] = pos +% 1; // simple eviction strat (overwrite what's there)

        const offset = pos -% match_pos; // convert match_pos to pos-relative offset
        if (match_pos == ~@as(u32, 0) or offset > 0xffff) continue; // invalid match or spec: offset > u16

        const match_cutoff = src.len - 5; // spec: last 5b must be literals
        const matches: u32 = @intCast(std.mem.indexOfDiff(u8, src[pos..match_cutoff], src[match_pos..]) orelse 0);
        if (matches < 4) continue; // spec: at least 4 bytes at the match_pos must match current pos

        try writeToken(w, src[literals_start_pos..pos], @intCast(offset), matches);
        pos += matches - 1; // minus one to cancel out while-loop's `pos += 1`
        literals_start_pos = pos + 1;
    }

    try writeToken(w, src[literals_start_pos..], 0, 0); // write remaining literals.
}

fn decompressBlock(r: *std.Io.Reader, dst: []u8) ![]const u8 {
    var pos: u32 = 0;
    while (true) {
        const token = try readToken(r, dst[pos..]);        
        pos += token.literals; // copied that many literals into dst[pos..]

        if (token.matches == 0) return dst[0..pos]; // was last literals token
        if (pos + token.matches > dst.len) return error.OutOfSpace; // copy would overflow dst
        if (token.offset == 0 or token.offset > pos) return error.Invalid; // bad offset
        
        for (0..token.matches) |i| dst[pos + i] = dst[(pos - token.offset) + i]; // not memmove, as matches can be > offset
        pos += token.matches;
    }
}