tiny.sql.key

1const std = @import("std");2const row = @import("row.zig");34pub const Error = error{5    InvalidKey,6    OutputTooSmall,7};89pub const rowid_size: usize = 8;1011pub const DecodedIndex = struct {12    values: []const row.Value,13    rowid: i64,14};1516const DecodedValue = struct {17    value: row.Value,18    key_len: usize,19    scratch_len: usize = 0,20};2122const DecodedBytes = struct {23    bytes: []const u8,24    key_len: usize,25};2627const sign_bit: u64 = 0x8000_0000_0000_0000;28const nil_tag: u8 = 0x00;29const integer_tag: u8 = 0x10;30const text_tag: u8 = 0x20;31const blob_tag: u8 = 0x30;32const rowid_tag: u8 = 0x40;3334pub 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}4041pub 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}4647pub 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}5556pub 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}6263pub 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}6970pub fn decodeIndexRowId(bytes: []const u8) Error!i64 {71    const offset = try indexRowIdOffset(bytes);72    return try decodeRowId(bytes[offset + 1 ..]);73}7475pub fn decodeIndex(target: []row.Value, scratch: []u8, bytes: []const u8) Error!DecodedIndex {76    const rowid_offset = try indexRowIdOffset(bytes);7778    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    }8990    return .{91        .values = target[0..count],92        .rowid = try decodeRowId(bytes[rowid_offset + 1 ..]),93    };94}9596fn 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}102103fn 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}111112fn 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}145146fn 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}174175fn writeTag(target: []u8, tag: u8) Error!usize {176    if (target.len < 1) return error.OutputTooSmall;177    target[0] = tag;178    return 1;179}180181fn 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}187188fn 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}194195fn 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}200201fn 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}206207fn 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}224225fn 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}231232fn 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}243244fn 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}250251fn asciiLower(byte: u8) u8 {252    if (byte >= 'A' and byte <= 'Z') return byte + ('a' - 'A');253    return byte;254}255256fn 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}261262test "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    };272273    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}282283test "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}288289test "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}339340test "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);345346    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);354355    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}363364test "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    );377378    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}388389test "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    );397398    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}403404test "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}412413fn 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}

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