turbonss/src/compress.zig

//
// varint64 []const u8 variants
//
// Thanks to https://github.com/gsquire/zig-snappy/blob/master/snappy.zig and golang's
// varint implementation.

const std = @import("std");

// Represents a variable length integer that we read from a byte stream along
// with how many bytes were read to decode it.
pub const Varint = struct {
    value: u64,
    bytesRead: usize,
};

const maxVarintLen64 = 10;

// https://golang.org/pkg/encoding/binary/#Uvarint
pub fn uvarint(buf: []const u8) error{Overflow}!Varint {
    var x: u64 = 0;
    var s: u6 = 0;

    for (buf) |b, i| {
        if (i == maxVarintLen64) {
            // Catch byte reads past maxVarintLen64.
            // See issue https://golang.org/issues/41185
            return error.Overflow;
        }

        if (b < 0x80) {
            if (i == maxVarintLen64 - 1 and b > 1) {
                return error.Overflow;
            }
            return Varint{ .value = x | (@as(u64, b) << s), .bytesRead = i + 1 };
        }
        x |= (@as(u64, b & 0x7f) << s);
        s = try std.math.add(u6, s, 7);
    }

    return Varint{
        .value = 0,
        .bytesRead = 0,
    };
}

// https://golang.org/pkg/encoding/binary/#PutUvarint
pub fn putUvarint(buf: []u8, x: u64) usize {
    var i: usize = 0;
    var mutX = x;

    while (mutX >= 0x80) {
        buf[i] = @truncate(u8, mutX) | 0x80;
        mutX >>= 7;
        i += 1;
    }
    buf[i] = @truncate(u8, mutX);

    return i + 1;
}

const testing = std.testing;

const tests = [_]u64{
    0,
    1,
    2,
    10,
    20,
    63,
    64,
    65,
    127,
    128,
    129,
    255,
    256,
    257,
    1 << 63 - 1,
};

test "uvarint" {
    for (tests) |x| {
        var buf: [maxVarintLen64]u8 = undefined;
        const n = putUvarint(buf[0..], x);
        const got = try uvarint(buf[0..n]);

        try testing.expectEqual(x, got.value);
        try testing.expectEqual(n, got.bytesRead);
    }
}

test "overflow" {
    for ([_][]const u8{
        &[_]u8{ 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x2 },
        &[_]u8{ 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x1, 0, 0 },
        &[_]u8{ 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF },
    }) |t| {
        try testing.expectError(error.Overflow, uvarint(t));
    }
}