lib/sql/src/key.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const row = @import("row.zig");
3
4 pub const Error = error{
5 InvalidKey,
6 OutputTooSmall,
7 };
8
9 pub const rowid_size: usize = 8;
10
11 pub const DecodedIndex = struct {
12 values: []const row.Value,
13 rowid: i64,
14 };
15
16 const DecodedValue = struct {
17 value: row.Value,
18 key_len: usize,
19 scratch_len: usize = 0,
20 };
21
22 const DecodedBytes = struct {
23 bytes: []const u8,
24 key_len: usize,
25 };
26
27 const sign_bit: u64 = 0x8000_0000_0000_0000;
28 const nil_tag: u8 = 0x00;
29 const integer_tag: u8 = 0x10;
30 const text_tag: u8 = 0x20;
31 const blob_tag: u8 = 0x30;
32 const rowid_tag: u8 = 0x40;
33
34 pub fn encodeRowId(target: []u8, rowid: i64) Error![]const u8 {
35 if (target.len < rowid_size) return error.OutputTooSmall;
36 const sortable = @as(u64, @bitCast(rowid)) ^ sign_bit;
37 std.mem.writeInt(u64, target[0..rowid_size], sortable, .big);
38 return target[0..rowid_size];
39 }
40
41 pub fn decodeRowId(bytes: []const u8) Error!i64 {
42 if (bytes.len != rowid_size) return error.InvalidKey;
43 const sortable = std.mem.readInt(u64, bytes[0..rowid_size], .big);
44 return @bitCast(sortable ^ sign_bit);
45 }
46
47 pub fn encodeIndexPrefix(target: []u8, values: []const row.Value, columns: []const row.Column) Error![]const u8 {
48 var cursor: usize = 0;
49 for (values, 0..) |value, index| {
50 const column = if (index < columns.len) columns[index] else row.Column{};
51 cursor += try encodeValue(target[cursor..], value, column);
52 }
53 return target[0..cursor];
54 }
55
56 pub fn encodeIndex(target: []u8, values: []const row.Value, columns: []const row.Column, rowid: i64) Error![]const u8 {
57 const prefix = try encodeIndexPrefix(target, values, columns);
58 var cursor = prefix.len;
59 cursor += try encodeRowIdSegment(target[cursor..], rowid);
60 return target[0..cursor];
61 }
62
63 pub fn encodePrefixEnd(target: []u8, prefix: []const u8) Error![]const u8 {
64 if (target.len < prefix.len + 1) return error.OutputTooSmall;
65 @memcpy(target[0..prefix.len], prefix);
66 target[prefix.len] = 0xff;
67 return target[0 .. prefix.len + 1];
68 }
69
70 pub fn decodeIndexRowId(bytes: []const u8) Error!i64 {
71 const offset = try indexRowIdOffset(bytes);
72 return try decodeRowId(bytes[offset + 1 ..]);
73 }
74
75 pub fn decodeIndex(target: []row.Value, scratch: []u8, bytes: []const u8) Error!DecodedIndex {
76 const rowid_offset = try indexRowIdOffset(bytes);
77
78 var cursor: usize = 0;
79 var scratch_cursor: usize = 0;
80 var count: usize = 0;
81 while (cursor < rowid_offset) {
82 if (count >= target.len) return error.OutputTooSmall;
83 const decoded = try decodeValue(bytes[cursor..rowid_offset], scratch[scratch_cursor..]);
84 target[count] = decoded.value;
85 cursor += decoded.key_len;
86 scratch_cursor += decoded.scratch_len;
87 count += 1;
88 }
89
90 return .{
91 .values = target[0..count],
92 .rowid = try decodeRowId(bytes[rowid_offset + 1 ..]),
93 };
94 }
95
96 fn indexRowIdOffset(bytes: []const u8) Error!usize {
97 if (bytes.len < rowid_size + 1) return error.InvalidKey;
98 const offset = bytes.len - rowid_size - 1;
99 if (bytes[offset] != rowid_tag) return error.InvalidKey;
100 return offset;
101 }
102
103 fn encodeValue(target: []u8, value: row.Value, column: row.Column) Error!usize {
104 return switch (value) {
105 .nil => writeTag(target, nil_tag),
106 .integer => |integer| writeInteger(target, integer),
107 .text => |text| writeText(target, text, column.collation),
108 .blob => |blob| writeBlob(target, blob),
109 };
110 }
111
112 fn decodeValue(source: []const u8, scratch: []u8) Error!DecodedValue {
113 if (source.len == 0) return error.InvalidKey;
114 return switch (source[0]) {
115 nil_tag => .{
116 .value = .nil,
117 .key_len = 1,
118 },
119 integer_tag => integer: {
120 if (source.len < rowid_size + 1) return error.InvalidKey;
121 break :integer .{
122 .value = .{ .integer = try decodeRowId(source[1..][0..rowid_size]) },
123 .key_len = rowid_size + 1,
124 };
125 },
126 text_tag => text: {
127 const decoded = try decodeEscaped(source[1..], scratch);
128 break :text .{
129 .value = .{ .text = decoded.bytes },
130 .key_len = decoded.key_len + 1,
131 .scratch_len = decoded.bytes.len,
132 };
133 },
134 blob_tag => blob: {
135 const decoded = try decodeEscaped(source[1..], scratch);
136 break :blob .{
137 .value = .{ .blob = decoded.bytes },
138 .key_len = decoded.key_len + 1,
139 .scratch_len = decoded.bytes.len,
140 };
141 },
142 else => error.InvalidKey,
143 };
144 }
145
146 fn decodeEscaped(source: []const u8, scratch: []u8) Error!DecodedBytes {
147 var cursor: usize = 0;
148 var scratch_cursor: usize = 0;
149 while (cursor < source.len) {
150 const byte = source[cursor];
151 if (byte == 0) {
152 if (cursor + 1 >= source.len) return error.InvalidKey;
153 const next = source[cursor + 1];
154 if (next == 0) {
155 return .{
156 .bytes = scratch[0..scratch_cursor],
157 .key_len = cursor + 2,
158 };
159 }
160 if (next != 0xff) return error.InvalidKey;
161 if (scratch_cursor >= scratch.len) return error.OutputTooSmall;
162 scratch[scratch_cursor] = 0;
163 scratch_cursor += 1;
164 cursor += 2;
165 continue;
166 }
167 if (scratch_cursor >= scratch.len) return error.OutputTooSmall;
168 scratch[scratch_cursor] = byte;
169 scratch_cursor += 1;
170 cursor += 1;
171 }
172 return error.InvalidKey;
173 }
174
175 fn writeTag(target: []u8, tag: u8) Error!usize {
176 if (target.len < 1) return error.OutputTooSmall;
177 target[0] = tag;
178 return 1;
179 }
180
181 fn writeInteger(target: []u8, integer: i64) Error!usize {
182 if (target.len < rowid_size + 1) return error.OutputTooSmall;
183 target[0] = integer_tag;
184 _ = try encodeRowId(target[1..], integer);
185 return rowid_size + 1;
186 }
187
188 fn encodeRowIdSegment(target: []u8, rowid: i64) Error!usize {
189 if (target.len < rowid_size + 1) return error.OutputTooSmall;
190 target[0] = rowid_tag;
191 _ = try encodeRowId(target[1..], rowid);
192 return rowid_size + 1;
193 }
194
195 fn writeText(target: []u8, text: []const u8, collation: row.Collation) Error!usize {
196 if (target.len < 1) return error.OutputTooSmall;
197 target[0] = text_tag;
198 return 1 + try writeEscapedText(target[1..], text, collation);
199 }
200
201 fn writeBlob(target: []u8, blob: []const u8) Error!usize {
202 if (target.len < 1) return error.OutputTooSmall;
203 target[0] = blob_tag;
204 return 1 + try writeEscapedBytes(target[1..], blob);
205 }
206
207 fn writeEscapedText(target: []u8, text: []const u8, collation: row.Collation) Error!usize {
208 const bytes = switch (collation) {
209 .binary => text,
210 .nocase => text,
211 .rtrim => trimRightSpaces(text),
212 };
213 var cursor: usize = 0;
214 for (bytes) |byte| {
215 const encoded = switch (collation) {
216 .binary, .rtrim => byte,
217 .nocase => asciiLower(byte),
218 };
219 cursor += try writeEscapedByte(target[cursor..], encoded);
220 }
221 cursor += try writeTerminator(target[cursor..]);
222 return cursor;
223 }
224
225 fn writeEscapedBytes(target: []u8, bytes: []const u8) Error!usize {
226 var cursor: usize = 0;
227 for (bytes) |byte| cursor += try writeEscapedByte(target[cursor..], byte);
228 cursor += try writeTerminator(target[cursor..]);
229 return cursor;
230 }
231
232 fn writeEscapedByte(target: []u8, byte: u8) Error!usize {
233 if (byte == 0) {
234 if (target.len < 2) return error.OutputTooSmall;
235 target[0] = 0;
236 target[1] = 0xff;
237 return 2;
238 }
239 if (target.len < 1) return error.OutputTooSmall;
240 target[0] = byte;
241 return 1;
242 }
243
244 fn writeTerminator(target: []u8) Error!usize {
245 if (target.len < 2) return error.OutputTooSmall;
246 target[0] = 0;
247 target[1] = 0;
248 return 2;
249 }
250
251 fn asciiLower(byte: u8) u8 {
252 if (byte >= 'A' and byte <= 'Z') return byte + ('a' - 'A');
253 return byte;
254 }
255
256 fn trimRightSpaces(bytes: []const u8) []const u8 {
257 var end = bytes.len;
258 while (end > 0 and bytes[end - 1] == ' ') end -= 1;
259 return bytes[0..end];
260 }
261
262 test "rowid keys preserve signed integer ordering" {
263 const values = [_]i64{
264 std.math.minInt(i64),
265 -1_000_000,
266 -1,
267 0,
268 1,
269 1_000_000,
270 std.math.maxInt(i64),
271 };
272
273 var previous: [rowid_size]u8 = undefined;
274 for (values, 0..) |value, index| {
275 var encoded: [rowid_size]u8 = undefined;
276 _ = try encodeRowId(&encoded, value);
277 try std.testing.expectEqual(value, try decodeRowId(&encoded));
278 if (index > 0) try std.testing.expect(std.mem.order(u8, &previous, &encoded) == .lt);
279 previous = encoded;
280 }
281 }
282
283 test "rowid keys reject malformed buffers" {
284 var short: [rowid_size - 1]u8 = undefined;
285 try std.testing.expectError(error.OutputTooSmall, encodeRowId(&short, 1));
286 try std.testing.expectError(error.InvalidKey, decodeRowId(&short));
287 }
288
289 test "index keys sort by storage class value collation and rowid" {
290 try expectIndexOrder(
291 &.{row.Value.nil},
292 9,
293 &.{.{ .integer = std.math.minInt(i64) }},
294 0,
295 &.{},
296 .lt,
297 );
298 try expectIndexOrder(
299 &.{.{ .integer = -5 }},
300 9,
301 &.{.{ .integer = 7 }},
302 0,
303 &.{},
304 .lt,
305 );
306 try expectIndexOrder(
307 &.{.{ .text = "Alpha" }},
308 9,
309 &.{.{ .text = "alpha" }},
310 0,
311 &.{.{ .collation = .nocase }},
312 .gt,
313 );
314 try expectIndexOrder(
315 &.{.{ .text = "a " }},
316 1,
317 &.{.{ .text = "a" }},
318 2,
319 &.{.{ .collation = .rtrim }},
320 .lt,
321 );
322 try expectIndexOrder(
323 &.{.{ .text = "a\x00b" }},
324 9,
325 &.{.{ .text = "a\x00c" }},
326 0,
327 &.{},
328 .lt,
329 );
330 try expectIndexOrder(
331 &.{.{ .text = "z" }},
332 9,
333 &.{.{ .blob = "a" }},
334 0,
335 &.{},
336 .lt,
337 );
338 }
339
340 test "index prefix end bounds exact value prefix" {
341 var prefix_buffer: [64]u8 = undefined;
342 const prefix = try encodeIndexPrefix(&prefix_buffer, &.{.{ .integer = 5 }}, &.{});
343 var end_buffer: [64]u8 = undefined;
344 const end = try encodePrefixEnd(&end_buffer, prefix);
345
346 var lower_buffer: [64]u8 = undefined;
347 const lower = try encodeIndex(&lower_buffer, &.{.{ .integer = 4 }}, &.{}, 0);
348 var first_buffer: [64]u8 = undefined;
349 const first = try encodeIndex(&first_buffer, &.{.{ .integer = 5 }}, &.{}, -1);
350 var second_buffer: [64]u8 = undefined;
351 const second = try encodeIndex(&second_buffer, &.{.{ .integer = 5 }}, &.{}, 2);
352 var upper_buffer: [64]u8 = undefined;
353 const upper = try encodeIndex(&upper_buffer, &.{.{ .integer = 6 }}, &.{}, 0);
354
355 try std.testing.expect(std.mem.order(u8, lower, prefix) == .lt);
356 try std.testing.expect(std.mem.order(u8, first, prefix) == .gt);
357 try std.testing.expect(std.mem.order(u8, first, end) == .lt);
358 try std.testing.expect(std.mem.order(u8, second, end) == .lt);
359 try std.testing.expect(std.mem.order(u8, upper, end) == .gt);
360 try std.testing.expectEqual(@as(i64, -1), try decodeIndexRowId(first));
361 try std.testing.expectEqual(@as(i64, 2), try decodeIndexRowId(second));
362 }
363
364 test "index key decoder recovers stored values and rowid suffix" {
365 var key_buffer: [128]u8 = undefined;
366 const encoded = try encodeIndex(
367 &key_buffer,
368 &.{
369 .{ .integer = -42 },
370 .{ .text = "a\x00b" },
371 .{ .blob = "\x00z" },
372 row.Value.nil,
373 },
374 &.{},
375 77,
376 );
377
378 var values: [4]row.Value = undefined;
379 var scratch: [64]u8 = undefined;
380 const decoded = try decodeIndex(&values, &scratch, encoded);
381 try std.testing.expectEqual(@as(usize, 4), decoded.values.len);
382 try std.testing.expectEqual(@as(i64, 77), decoded.rowid);
383 try std.testing.expectEqual(@as(i64, -42), decoded.values[0].integer);
384 try std.testing.expectEqualStrings("a\x00b", decoded.values[1].text);
385 try std.testing.expectEqualSlices(u8, "\x00z", decoded.values[2].blob);
386 try std.testing.expectEqual(row.Value.nil, decoded.values[3]);
387 }
388
389 test "index key decoder exposes collation stored bytes" {
390 var key_buffer: [128]u8 = undefined;
391 const encoded = try encodeIndex(
392 &key_buffer,
393 &.{.{ .text = "Alpha " }},
394 &.{.{ .collation = .rtrim }},
395 1,
396 );
397
398 var values: [1]row.Value = undefined;
399 var scratch: [32]u8 = undefined;
400 const decoded = try decodeIndex(&values, &scratch, encoded);
401 try std.testing.expectEqualStrings("Alpha", decoded.values[0].text);
402 }
403
404 test "index rowid decoder rejects keys without rowid suffix" {
405 var prefix_buffer: [64]u8 = undefined;
406 const prefix = try encodeIndexPrefix(&prefix_buffer, &.{.{ .text = "x" }}, &.{});
407 try std.testing.expectError(error.InvalidKey, decodeIndexRowId(prefix));
408 var values: [1]row.Value = undefined;
409 var scratch: [16]u8 = undefined;
410 try std.testing.expectError(error.InvalidKey, decodeIndex(&values, &scratch, prefix));
411 }
412
413 fn expectIndexOrder(
414 left_values: []const row.Value,
415 left_rowid: i64,
416 right_values: []const row.Value,
417 right_rowid: i64,
418 columns: []const row.Column,
419 expected: std.math.Order,
420 ) !void {
421 var left_buffer: [128]u8 = undefined;
422 var right_buffer: [128]u8 = undefined;
423 const left = try encodeIndex(&left_buffer, left_values, columns, left_rowid);
424 const right = try encodeIndex(&right_buffer, right_values, columns, right_rowid);
425 try std.testing.expectEqual(expected, std.mem.order(u8, left, right));
426 }