lib/sql/src/search/vbyte.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2
3 pub const upstream_revision = "8afe8693964b33abb06b3b6b48fad4d33693a241";
4
5 pub const upstream_source = "https://github.com/fast-pack/streamvbyte";
6
7 pub const Error = error{
8 InvalidSearchIndex,
9 OutputTooSmall,
10 };
11
12 pub const Decoder = struct {
13 input: []const u8 = &.{},
14 value_count: usize = 0,
15 value_index: usize = 0,
16 data_offset: usize = 0,
17
18 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 }
35
36 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 }
43
44 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 }
53
54 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 };
68
69 pub const Encoder = struct {
70 output: []u8,
71 value_count: usize,
72 value_index: usize = 0,
73 data_offset: usize,
74
75 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 }
85
86 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 }
97
98 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 };
102
103 pub fn controlByteCount(value_count: usize) usize {
104 return value_count / 4 + @intFromBool(value_count % 4 != 0);
105 }
106
107 pub 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 }
112
113 pub 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 }
121
122 pub 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 }
127
128 pub 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 }
131
132 fn 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 }
135
136 fn 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 }
150
151 fn 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 }
165
166 test "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);
172
173 var decoded: [values.len]u32 = undefined;
174 try decode(&decoded, &encoded);
175 try std.testing.expectEqualSlices(u32, &values, &decoded);
176 }
177
178 test "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);
183
184 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 }
191
192 test "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]));
198
199 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));
203
204 var noncanonical = encoded;
205 noncanonical[0] |= 0xc0;
206 try std.testing.expectError(error.InvalidSearchIndex, decode(&decoded, &noncanonical));
207
208 var too_small: [encoded.len - 1]u8 = undefined;
209 try std.testing.expectError(error.OutputTooSmall, encode(&too_small, &values));
210 }