tiny.sql.lattice

1const std = @import("std");2const simd = @import("simd");34const Aes128 = std.crypto.core.aes.Aes128;5const Blake3 = std.crypto.hash.Blake3;6const Sha256 = std.crypto.hash.sha2.Sha256;7const native_endian = @import("builtin").cpu.arch.endian();89pub const lane_count = 1024;10pub const encoded_size = lane_count * @sizeOf(u16);11pub const digest_size = Sha256.digest_length;1213const entry_tag = "sql.lattice.entry";14const digest_tag = "sql.lattice";1516const Lanes = simd.ScalableTag(u16);1718/// Counter blocks the entry expansion builds and encrypts together.19const group_blocks = 8;20const block_size = 16;21const group_size = group_blocks * block_size;22const Words = @Vector(group_blocks, u64);2324comptime {25    std.debug.assert(lane_count % Lanes.lane_count == 0);26    std.debug.assert(encoded_size % group_size == 0);27}2829pub const State = struct {30    lanes: [lane_count]u16,3132    pub const empty = State{ .lanes = @splat(0) };3334    pub fn add(self: *State, entry: *const State) void {35        simd.transform1(Lanes, &self.lanes, &entry.lanes, WrappingAdd{});36    }3738    pub fn subtract(self: *State, entry: *const State) void {39        simd.transform1(Lanes, &self.lanes, &entry.lanes, WrappingSubtract{});40    }4142    pub fn isEmpty(self: *const State) bool {43        return simd.allEqual(Lanes, &self.lanes, 0);44    }4546    pub fn eql(self: *const State, other: *const State) bool {47        return simd.equal(Lanes, &self.lanes, &other.lanes);48    }4950    pub fn encode(self: *const State, out: *[encoded_size]u8) void {51        if (native_endian == .little) {52            @memcpy(out, std.mem.asBytes(&self.lanes));53        } else {54            for (&self.lanes, 0..) |*lane, index| {55                std.mem.writeInt(u16, out[index * 2 ..][0..2], lane.*, .little);56            }57        }58    }5960    pub fn decode(bytes: *const [encoded_size]u8) State {61        var state: State = undefined;62        @memcpy(std.mem.asBytes(&state.lanes), bytes);63        state.lanesFromLittle();64        return state;65    }6667    pub fn digest(self: *const State, entries: u64) [digest_size]u8 {68        var hasher = Sha256.init(.{});69        hasher.update(digest_tag);70        var count: [8]u8 = undefined;71        std.mem.writeInt(u64, &count, entries, .big);72        hasher.update(&count);73        var encoded: [encoded_size]u8 = undefined;74        self.encode(&encoded);75        hasher.update(&encoded);76        var out: [digest_size]u8 = undefined;77        hasher.final(&out);78        return out;79    }8081    /// Puts lanes read from little-endian bytes in native byte order.82    fn lanesFromLittle(self: *State) void {83        if (native_endian == .big) std.mem.byteSwapAllElements(u16, &self.lanes);84    }85};8687const WrappingAdd = struct {88    pub fn call(_: WrappingAdd, comptime D: type, lanes: D.Vector, entry: D.Vector) D.Vector {89        return lanes +% entry;90    }91};9293const WrappingSubtract = struct {94    pub fn call(_: WrappingSubtract, comptime D: type, lanes: D.Vector, entry: D.Vector) D.Vector {95        return lanes -% entry;96    }97};9899pub const EntryHasher = struct {100    hasher: Blake3,101102    pub fn init(entry_key: []const u8) EntryHasher {103        var hasher = Blake3.init(.{});104        hasher.update(entry_tag);105        var length: [8]u8 = undefined;106        std.mem.writeInt(u64, &length, entry_key.len, .big);107        hasher.update(&length);108        hasher.update(entry_key);109        return .{ .hasher = hasher };110    }111112    pub fn update(self: *EntryHasher, value_chunk: []const u8) void {113        self.hasher.update(value_chunk);114    }115116    pub fn finish(self: *const EntryHasher) State {117        var digest: [32]u8 = undefined;118        self.hasher.final(&digest);119        var state: State = undefined;120        expand(&digest, std.mem.asBytes(&state.lanes));121        state.lanesFromLittle();122        return state;123    }124};125126pub fn entryState(entry_key: []const u8, value: []const u8) State {127    var hasher = EntryHasher.init(entry_key);128    hasher.update(value);129    return hasher.finish();130}131132/// Writes the keystream that expands an entry digest into lattice lanes:133/// AES-128 in counter mode over zeros, keyed by the digest's first half,134/// counting from its second half read as a big-endian integer. Each group135/// of counter blocks is written in place and encrypted there.136fn expand(digest: *const [32]u8, out: *[encoded_size]u8) void {137    const aes = Aes128.initEnc(digest[0..16].*);138    var counter = std.mem.readInt(u128, digest[16..32], .big);139    var offset: usize = 0;140    while (offset < encoded_size) : (offset += group_size) {141        const group = out[offset..][0..group_size];142        writeCounters(counter, group);143        aes.encryptWide(group_blocks, group, group);144        counter +%= group_blocks;145    }146}147148/// Writes the big-endian counter blocks `first` through149/// `first + group_blocks - 1`. Each vector lane adds its block's offset to150/// the low word, and a lane whose low word wraps carries one into its high151/// word, so no block formats a 128-bit integer a byte at a time.152fn writeCounters(first: u128, out: *[group_size]u8) void {153    const low_first: u64 = @truncate(first);154    const high_first: u64 = @truncate(first >> 64);155    const low = @as(Words, @splat(low_first)) +% std.simd.iota(u64, group_blocks);156    const wrapped = low < @as(Words, @splat(low_first));157    const carry = @select(u64, wrapped, @as(Words, @splat(1)), @as(Words, @splat(0)));158    const high = @as(Words, @splat(high_first)) +% carry;159    const words: [2 * group_blocks]u64 =160        @shuffle(u64, bigEndian(high), bigEndian(low), high_then_low);161    @memcpy(out, std.mem.asBytes(&words));162}163164/// Shuffle mask that takes each block's high word and then its low word.165const high_then_low = mask: {166    var lanes: [2 * group_blocks]i32 = undefined;167    for (0..group_blocks) |block| {168        lanes[2 * block] = @intCast(block);169        lanes[2 * block + 1] = ~@as(i32, @intCast(block));170    }171    break :mask lanes;172};173174fn bigEndian(words: Words) Words {175    return if (native_endian == .little) @byteSwap(words) else words;176}177178/// Expands a digest with the standard library's counter mode, the179/// construction `expand` reproduces.180fn expandReference(digest: *const [32]u8, out: *[encoded_size]u8) void {181    const aes = Aes128.initEnc(digest[0..16].*);182    const zeros: [encoded_size]u8 = @splat(0);183    std.crypto.core.modes.ctr(@TypeOf(aes), aes, out, &zeros, digest[16..32].*, .big);184}185186test "lattice expansion matches counter mode across counter carries" {187    var prng = std.Random.DefaultPrng.init(0x1a77_1ce5);188    const random = prng.random();189    const low_ends = [_]u64{ 0, 1, std.math.maxInt(u64) - group_blocks, std.math.maxInt(u64) };190    for (0..640) |round| {191        var digest: [32]u8 = undefined;192        random.bytes(&digest);193        if (round % 5 != 0) {194            const low_end = low_ends[round % low_ends.len];195            const offset = random.uintAtMost(u64, 2 * group_blocks);196            std.mem.writeInt(u64, digest[24..32], low_end -% offset, .big);197        }198        if (round % 8 == 0) @memset(digest[16..24], 0xff);199        var fresh: [encoded_size]u8 = undefined;200        var expected: [encoded_size]u8 = undefined;201        expand(&digest, &fresh);202        expandReference(&digest, &expected);203        try std.testing.expectEqualSlices(u8, &expected, &fresh);204    }205}206207test "lattice entry digests keep their persisted values" {208    const alpha = entryState("alpha", "one");209    const empty = entryState("", "");210    const beta = entryState("beta", "two");211    var pair = State.empty;212    pair.add(&alpha);213    pair.add(&beta);214    const alpha_hex = "0fa6d7690ea57bfcb0e5c0fff010ad1d11a101560d69c9471d71dc2f87dbf558";215    const empty_hex = "088cb4b17ab38962b95e7df2e7fe65dbaadab2ee20a7e48b21ce2c6210faaaad";216    const pair_hex = "53e5f002f002639e7c61c98d23876fbd8ddd8afcd3132bb61b527b973370a9e1";217    try expectDigest(alpha_hex, alpha.digest(1));218    try expectDigest(empty_hex, empty.digest(1));219    try expectDigest(pair_hex, pair.digest(2));220}221222fn expectDigest(comptime hex: []const u8, actual: [digest_size]u8) !void {223    var expected: [digest_size]u8 = undefined;224    _ = try std.fmt.hexToBytes(&expected, hex);225    try std.testing.expectEqualSlices(u8, &expected, &actual);226}227228test "lattice add and subtract invert" {229    const first = entryState("alpha", "one");230    const second = entryState("beta", "two");231232    var state = State.empty;233    state.add(&first);234    state.add(&second);235236    var expected = State.empty;237    expected.add(&first);238239    state.subtract(&second);240    try std.testing.expect(state.eql(&expected));241242    state.subtract(&first);243    try std.testing.expect(state.isEmpty());244}245246test "lattice entry framing separates key and value boundaries" {247    const joined_left = entryState("ab", "c");248    const joined_right = entryState("a", "bc");249    try std.testing.expect(!joined_left.eql(&joined_right));250251    const key_only = entryState("a", "");252    const value_only = entryState("", "a");253    try std.testing.expect(!key_only.eql(&value_only));254}255256test "lattice streaming value chunks match one shot" {257    var hasher = EntryHasher.init("stream");258    hasher.update("he");259    hasher.update("l");260    hasher.update("lo");261    const streamed = hasher.finish();262263    const oneshot = entryState("stream", "hello");264    try std.testing.expect(streamed.eql(&oneshot));265}266267test "lattice encode decode round trips" {268    var state = State.empty;269    const entry = entryState("round", "trip");270    state.add(&entry);271272    var encoded: [encoded_size]u8 = undefined;273    state.encode(&encoded);274    const decoded = State.decode(&encoded);275    try std.testing.expect(state.eql(&decoded));276}277278test "lattice digest binds entry count" {279    var state = State.empty;280    const entry = entryState("count", "bound");281    state.add(&entry);282283    const one = state.digest(1);284    const two = state.digest(2);285    try std.testing.expect(!std.mem.eql(u8, &one, &two));286}287288test "lattice empty digest differs from populated digest" {289    var state = State.empty;290    const empty_digest = State.empty.digest(0);291292    const entry = entryState("k", "v");293    state.add(&entry);294    const populated = state.digest(1);295    try std.testing.expect(!std.mem.eql(u8, &empty_digest, &populated));296}

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