tiny.sql.search.vbyte

1pub const vbyte = @import("vbyte.zig");

1const std = @import("std");23pub const upstream_revision = "8afe8693964b33abb06b3b6b48fad4d33693a241";45pub const upstream_source = "https://github.com/fast-pack/streamvbyte";67pub const Error = error{8    InvalidSearchIndex,9    OutputTooSmall,10};1112pub const Decoder = struct {13    input: []const u8 = &.{},14    value_count: usize = 0,15    value_index: usize = 0,16    data_offset: usize = 0,1718    pub fn init(input: []const u8, value_count: usize) Error!Decoder {19        const control_count = controlByteCount(value_count);20        if (input.len < control_count) return error.InvalidSearchIndex;21        _ = std.math.mul(usize, value_count, 4) catch return error.InvalidSearchIndex;22        var data_count: usize = 0;23        const full_control_count = value_count / 4;24        for (input[0..full_control_count]) |control| data_count += controlDataByteCount(control);25        const remaining = value_count % 4;26        if (remaining != 0) {27            const control = input[full_control_count];28            if (control >> @intCast(remaining * 2) != 0) return error.InvalidSearchIndex;29            for (0..remaining) |index| data_count += @as(usize, (control >> @intCast(index * 2)) & 0x3) + 1;30        }31        const encoded_count = std.math.add(usize, control_count, data_count) catch return error.InvalidSearchIndex;32        if (encoded_count != input.len) return error.InvalidSearchIndex;33        return initValidated(input, value_count);34    }3536    pub fn initValidated(input: []const u8, value_count: usize) Decoder {37        return .{38            .input = input,39            .value_count = value_count,40            .data_offset = controlByteCount(value_count),41        };42    }4344    pub fn next(self: *Decoder) Error!?u32 {45        if (self.value_index >= self.value_count) return null;46        const control = self.input[self.value_index / 4];47        const shift: u3 = @intCast((self.value_index % 4) * 2);48        const code = (control >> shift) & 0x3;49        const value = try readValue(self.input, &self.data_offset, code);50        self.value_index += 1;51        return value;52    }5354    pub fn nextBlock(self: *Decoder, output: *[4]u32, count: usize) Error!void {55        if (count == 0 or count > 4 or self.value_index % 4 != 0 or count > self.value_count - self.value_index) return error.InvalidSearchIndex;56        const control = self.input[self.value_index / 4];57        if (control == 0) {58            const end = std.math.add(usize, self.data_offset, count) catch return error.InvalidSearchIndex;59            if (end > self.input.len) return error.InvalidSearchIndex;60            for (0..count) |index| output[index] = self.input[self.data_offset + index];61            self.data_offset = end;62            self.value_index += count;63            return;64        }65        for (0..count) |index| output[index] = (try self.next()) orelse return error.InvalidSearchIndex;66    }67};6869pub const Encoder = struct {70    output: []u8,71    value_count: usize,72    value_index: usize = 0,73    data_offset: usize,7475    pub fn init(output: []u8, value_count: usize) Error!Encoder {76        const control_count = controlByteCount(value_count);77        if (output.len < control_count) return error.OutputTooSmall;78        @memset(output[0..control_count], 0);79        return .{80            .output = output,81            .value_count = value_count,82            .data_offset = control_count,83        };84    }8586    pub fn append(self: *Encoder, value: u32) Error!void {87        if (self.value_index >= self.value_count) return error.InvalidSearchIndex;88        const byte_count = valueByteCount(value);89        const end = std.math.add(usize, self.data_offset, byte_count) catch return error.OutputTooSmall;90        if (end > self.output.len) return error.OutputTooSmall;91        const shift: u3 = @intCast((self.value_index % 4) * 2);92        self.output[self.value_index / 4] |= @as(u8, @intCast(byte_count - 1)) << shift;93        writeValue(self.output[self.data_offset..end], value);94        self.data_offset = end;95        self.value_index += 1;96    }9798    pub fn finish(self: *const Encoder) Error!void {99        if (self.value_index != self.value_count or self.data_offset != self.output.len) return error.InvalidSearchIndex;100    }101};102103pub fn controlByteCount(value_count: usize) usize {104    return value_count / 4 + @intFromBool(value_count % 4 != 0);105}106107pub fn encodedSize(values: []const u32) ?usize {108    var size = controlByteCount(values.len);109    for (values) |value| size = std.math.add(usize, size, valueByteCount(value)) catch return null;110    return size;111}112113pub fn encode(target: []u8, values: []const u32) Error![]const u8 {114    const encoded_size = encodedSize(values) orelse return error.InvalidSearchIndex;115    if (target.len < encoded_size) return error.OutputTooSmall;116    var encoder = try Encoder.init(target[0..encoded_size], values.len);117    for (values) |value| try encoder.append(value);118    try encoder.finish();119    return target[0..encoded_size];120}121122pub fn decode(target: []u32, input: []const u8) Error!void {123    var decoder = try Decoder.init(input, target.len);124    for (target) |*value| value.* = (try decoder.next()) orelse return error.InvalidSearchIndex;125    if (try decoder.next() != null) return error.InvalidSearchIndex;126}127128pub fn valueByteCount(value: u32) usize {129    return 1 + @as(usize, @intFromBool(value > 0xff)) + @as(usize, @intFromBool(value > 0xffff)) + @as(usize, @intFromBool(value > 0xff_ffff));130}131132fn controlDataByteCount(control: u8) usize {133    return 4 + @as(usize, control & 0x3) + @as(usize, (control >> 2) & 0x3) + @as(usize, (control >> 4) & 0x3) + @as(usize, control >> 6);134}135136fn writeValue(target: []u8, value: u32) void {137    std.debug.assert(target.len == valueByteCount(value));138    switch (target.len) {139        1 => target[0] = @truncate(value),140        2 => std.mem.writeInt(u16, target[0..2], @truncate(value), .little),141        3 => {142            target[0] = @truncate(value);143            target[1] = @truncate(value >> 8);144            target[2] = @truncate(value >> 16);145        },146        4 => std.mem.writeInt(u32, target[0..4], value, .little),147        else => unreachable,148    }149}150151fn readValue(input: []const u8, offset: *usize, code: u8) Error!u32 {152    const byte_count = @as(usize, code) + 1;153    const end = std.math.add(usize, offset.*, byte_count) catch return error.InvalidSearchIndex;154    if (end > input.len) return error.InvalidSearchIndex;155    const value: u32 = switch (byte_count) {156        1 => input[offset.*],157        2 => std.mem.readInt(u16, input[offset.*..][0..2], .little),158        3 => @as(u32, input[offset.*]) | (@as(u32, input[offset.* + 1]) << 8) | (@as(u32, input[offset.* + 2]) << 16),159        4 => std.mem.readInt(u32, input[offset.*..][0..4], .little),160        else => unreachable,161    };162    offset.* = end;163    return value;164}165166test "Stream VByte 8afe869 format fixture is exact" {167    const values = [_]u32{ 0, 100, 200, 300, 400, 500, 600, 700 };168    const expected = [_]u8{ 0x40, 0x55, 0x00, 0x64, 0xc8, 0x2c, 0x01, 0x90, 0x01, 0xf4, 0x01, 0x58, 0x02, 0xbc, 0x02 };169    var encoded: [expected.len]u8 = undefined;170    _ = try encode(&encoded, &values);171    try std.testing.expectEqualSlices(u8, &expected, &encoded);172173    var decoded: [values.len]u32 = undefined;174    try decode(&decoded, &encoded);175    try std.testing.expectEqualSlices(u32, &values, &decoded);176}177178test "Stream VByte preserves byte-width boundaries" {179    var empty_encoded: [0]u8 = .{};180    var empty_decoded: [0]u32 = .{};181    try std.testing.expectEqual(@as(usize, 0), (try encode(&empty_encoded, &empty_decoded)).len);182    try decode(&empty_decoded, &empty_encoded);183184    const values = [_]u32{ 0, 0xff, 0x100, 0xffff, 0x1_0000, 0xff_ffff, 0x100_0000, std.math.maxInt(u32) };185    var encoded: [22]u8 = undefined;186    _ = try encode(&encoded, &values);187    var decoded: [values.len]u32 = undefined;188    try decode(&decoded, &encoded);189    try std.testing.expectEqualSlices(u32, &values, &decoded);190}191192test "Stream VByte rejects malformed streams" {193    const values = [_]u32{ 1, 0x100, 3 };194    var encoded: [5]u8 = undefined;195    _ = try encode(&encoded, &values);196    var decoded: [values.len]u32 = undefined;197    try std.testing.expectError(error.InvalidSearchIndex, decode(&decoded, encoded[0 .. encoded.len - 1]));198199    var trailing: [encoded.len + 1]u8 = undefined;200    @memcpy(trailing[0..encoded.len], &encoded);201    trailing[encoded.len] = 0;202    try std.testing.expectError(error.InvalidSearchIndex, decode(&decoded, &trailing));203204    var noncanonical = encoded;205    noncanonical[0] |= 0xc0;206    try std.testing.expectError(error.InvalidSearchIndex, decode(&decoded, &noncanonical));207208    var too_small: [encoded.len - 1]u8 = undefined;209    try std.testing.expectError(error.OutputTooSmall, encode(&too_small, &values));210}