1
Fork 0
turbonss/lib/compress.zig

324 lines
9.0 KiB
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");
const ArrayList = std.ArrayList;
const Allocator = std.mem.Allocator;
const assert = std.debug.assert;
const math = std.math;
// compresses a strictly incrementing sorted slice of integers using delta
// compression. Compression is in-place.
pub fn deltaCompress(comptime T: type, elems: []T) error{NotSorted}!void {
if (elems.len <= 1) {
return;
}
var prev: T = elems[0];
var i: usize = 1;
while (i < elems.len) : (i += 1) {
const cur = elems[i];
if (cur <= prev) {
return error.NotSorted;
}
elems[i] = cur - prev - 1;
prev = cur;
}
}
// decompresses a slice compressed by deltaCompress. In-place.
pub fn deltaDecompress(comptime T: type, elems: []T) error{Overflow}!void {
if (elems.len <= 1) {
return;
}
var i: usize = 1;
while (i < elems.len) : (i += 1) {
const x = try math.add(T, elems[i - 1], 1);
elems[i] = try math.add(T, elems[i], x);
}
}
// 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,
bytes_read: usize,
};
pub 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),
.bytes_read = i + 1,
};
}
x |= (@as(u64, b & 0x7f) << s);
s = try math.add(u6, s, 7);
}
return Varint{
.value = 0,
.bytes_read = 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;
}
// VarintSliceIterator iterates over varint-encoded slice.
// The first element is the length of the slice, in decoded numbers.
const varintSliceIterator = struct {
remaining: usize,
arr: []const u8,
idx: usize,
pub fn next(self: *varintSliceIterator) error{Overflow}!?u64 {
if (self.remaining == 0)
return null;
const value = try uvarint(self.arr[self.idx..]);
self.idx += value.bytes_read;
self.remaining -= 1;
return value.value;
}
pub fn nextMust(self: *varintSliceIterator) ?u64 {
return self.next() catch |err| switch (err) {
error.Overflow => unreachable,
};
}
// returns the number of remaining items. If called before the first
// next(), returns the length of the slice.
pub fn remaining(self: *const varintSliceIterator) usize {
return self.remaining;
}
};
pub fn VarintSliceIterator(arr: []const u8) error{Overflow}!varintSliceIterator {
const firstnumber = try uvarint(arr);
return varintSliceIterator{
.remaining = firstnumber.value,
.arr = arr,
.idx = firstnumber.bytes_read,
};
}
pub fn VarintSliceIteratorMust(arr: []const u8) varintSliceIterator {
return VarintSliceIterator(arr) catch |err| switch (err) {
error.Overflow => unreachable,
};
}
const deltaDecompressionIterator = struct {
vit: *varintSliceIterator,
prev: u64,
add_to_prev: u1,
pub fn next(self: *deltaDecompressionIterator) error{Overflow}!?u64 {
const current = try self.vit.next();
if (current == null) return null;
const prevExtra = try math.add(u64, self.prev, self.add_to_prev);
const result = try math.add(u64, current.?, prevExtra);
self.prev = result;
self.add_to_prev = 1;
return result;
}
// returns the number of remaining items. If called before the first
// next(), returns the length of the slice.
pub fn remaining(self: *const deltaDecompressionIterator) usize {
return self.vit.remaining;
}
};
pub fn DeltaDecompressionIterator(vit: *varintSliceIterator) deltaDecompressionIterator {
return deltaDecompressionIterator{
.vit = vit,
.prev = 0,
.add_to_prev = 0,
};
}
pub fn appendUvarint(arr: *ArrayList(u8), x: u64) Allocator.Error!void {
var buf: [maxVarintLen64]u8 = undefined;
const n = putUvarint(&buf, x);
try arr.appendSlice(buf[0..n]);
}
const testing = std.testing;
const uvarint_tests = [_]u64{
0,
1,
2,
10,
20,
63,
64,
65,
127,
128,
129,
255,
256,
257,
1 << 63 - 1,
};
test "putUvarint/uvarint" {
for (uvarint_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.bytes_read);
}
}
test "VarintSliceIterator" {
var buf = ArrayList(u8).init(testing.allocator);
defer buf.deinit();
try appendUvarint(&buf, uvarint_tests.len);
for (uvarint_tests) |x|
try appendUvarint(&buf, x);
var it = try VarintSliceIterator(buf.items);
var i: usize = 0;
while (try it.next()) |got| : (i += 1) {
try testing.expectEqual(uvarint_tests[i], got);
}
try testing.expectEqual(i, uvarint_tests.len);
}
test "delta compress/decompress" {
const tests = [_]struct { input: []const u8, want: []const u8 }{
.{ .input = &[_]u8{}, .want = &[_]u8{} },
.{ .input = &[_]u8{0}, .want = &[_]u8{0} },
.{ .input = &[_]u8{10}, .want = &[_]u8{10} },
.{ .input = &[_]u8{ 0, 1, 2 }, .want = &[_]u8{ 0, 0, 0 } },
.{ .input = &[_]u8{ 10, 20, 30, 255 }, .want = &[_]u8{ 10, 9, 9, 224 } },
.{ .input = &[_]u8{ 0, 254, 255 }, .want = &[_]u8{ 0, 253, 0 } },
};
for (tests) |t| {
var arr = try ArrayList(u8).initCapacity(
testing.allocator,
t.input.len,
);
defer arr.deinit();
try arr.appendSlice(t.input);
try deltaCompress(u8, arr.items);
try testing.expectEqualSlices(u8, arr.items, t.want);
try deltaDecompress(u8, arr.items);
try testing.expectEqualSlices(u8, arr.items, t.input);
}
}
test "delta compression with varint tests" {
var scratch: [uvarint_tests.len]u64 = undefined;
std.mem.copy(u64, scratch[0..], uvarint_tests[0..]);
try deltaCompress(u64, scratch[0..]);
try deltaDecompress(u64, scratch[0..]);
try testing.expectEqualSlices(u64, uvarint_tests[0..], scratch[0..]);
}
test "delta compression negative tests" {
for ([_][]const u8{
&[_]u8{ 0, 0 },
&[_]u8{ 0, 1, 1 },
&[_]u8{ 0, 1, 2, 1 },
}) |t| {
var arr = try ArrayList(u8).initCapacity(testing.allocator, t.len);
defer arr.deinit();
try arr.appendSlice(t);
try testing.expectError(error.NotSorted, deltaCompress(u8, arr.items));
}
}
test "delta decompress overflow" {
for ([_][]const u8{
&[_]u8{ 255, 0 },
&[_]u8{ 0, 128, 127 },
}) |t| {
var arr = try ArrayList(u8).initCapacity(testing.allocator, t.len);
defer arr.deinit();
try arr.appendSlice(t);
try testing.expectError(error.Overflow, deltaDecompress(u8, arr.items));
}
}
test "delta decompression with an iterator" {
var compressed: [uvarint_tests.len]u64 = undefined;
std.mem.copy(u64, compressed[0..], uvarint_tests[0..]);
try deltaCompress(u64, compressed[0..]);
var buf = ArrayList(u8).init(testing.allocator);
defer buf.deinit();
try appendUvarint(&buf, compressed.len);
for (compressed) |x|
try appendUvarint(&buf, x);
var it = DeltaDecompressionIterator(&try VarintSliceIterator(buf.items));
var i: usize = 0;
try testing.expectEqual(it.remaining(), uvarint_tests.len);
while (try it.next()) |got| : (i += 1) {
try testing.expectEqual(uvarint_tests[i], got);
}
try testing.expectEqual(i, uvarint_tests.len);
}
test "appendUvarint" {
for (uvarint_tests) |x| {
var buf = ArrayList(u8).init(testing.allocator);
defer buf.deinit();
try appendUvarint(&buf, x);
const got = try uvarint(buf.items);
try testing.expectEqual(x, got.value);
}
}
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));
}
}