lib/sql/src/lattice.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const simd = @import("simd");
3
4 const Aes128 = std.crypto.core.aes.Aes128;
5 const Blake3 = std.crypto.hash.Blake3;
6 const Sha256 = std.crypto.hash.sha2.Sha256;
7 const native_endian = @import("builtin").cpu.arch.endian();
8
9 pub const lane_count = 1024;
10 pub const encoded_size = lane_count * @sizeOf(u16);
11 pub const digest_size = Sha256.digest_length;
12
13 const entry_tag = "sql.lattice.entry";
14 const digest_tag = "sql.lattice";
15
16 const Lanes = simd.ScalableTag(u16);
17
18 /// Counter blocks the entry expansion builds and encrypts together.
19 const group_blocks = 8;
20 const block_size = 16;
21 const group_size = group_blocks * block_size;
22 const Words = @Vector(group_blocks, u64);
23
24 comptime {
25 std.debug.assert(lane_count % Lanes.lane_count == 0);
26 std.debug.assert(encoded_size % group_size == 0);
27 }
28
29 pub const State = struct {
30 lanes: [lane_count]u16,
31
32 pub const empty = State{ .lanes = @splat(0) };
33
34 pub fn add(self: *State, entry: *const State) void {
35 simd.transform1(Lanes, &self.lanes, &entry.lanes, WrappingAdd{});
36 }
37
38 pub fn subtract(self: *State, entry: *const State) void {
39 simd.transform1(Lanes, &self.lanes, &entry.lanes, WrappingSubtract{});
40 }
41
42 pub fn isEmpty(self: *const State) bool {
43 return simd.allEqual(Lanes, &self.lanes, 0);
44 }
45
46 pub fn eql(self: *const State, other: *const State) bool {
47 return simd.equal(Lanes, &self.lanes, &other.lanes);
48 }
49
50 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 }
59
60 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 }
66
67 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 }
80
81 /// 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 };
86
87 const WrappingAdd = struct {
88 pub fn call(_: WrappingAdd, comptime D: type, lanes: D.Vector, entry: D.Vector) D.Vector {
89 return lanes +% entry;
90 }
91 };
92
93 const WrappingSubtract = struct {
94 pub fn call(_: WrappingSubtract, comptime D: type, lanes: D.Vector, entry: D.Vector) D.Vector {
95 return lanes -% entry;
96 }
97 };
98
99 pub const EntryHasher = struct {
100 hasher: Blake3,
101
102 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 }
111
112 pub fn update(self: *EntryHasher, value_chunk: []const u8) void {
113 self.hasher.update(value_chunk);
114 }
115
116 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 };
125
126 pub 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 }
131
132 /// 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 group
135 /// of counter blocks is written in place and encrypted there.
136 fn 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 }
147
148 /// Writes the big-endian counter blocks `first` through
149 /// `first + group_blocks - 1`. Each vector lane adds its block's offset to
150 /// the low word, and a lane whose low word wraps carries one into its high
151 /// word, so no block formats a 128-bit integer a byte at a time.
152 fn 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 }
163
164 /// Shuffle mask that takes each block's high word and then its low word.
165 const 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 };
173
174 fn bigEndian(words: Words) Words {
175 return if (native_endian == .little) @byteSwap(words) else words;
176 }
177
178 /// Expands a digest with the standard library's counter mode, the
179 /// construction `expand` reproduces.
180 fn 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 }
185
186 test "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 }
206
207 test "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 }
221
222 fn 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 }
227
228 test "lattice add and subtract invert" {
229 const first = entryState("alpha", "one");
230 const second = entryState("beta", "two");
231
232 var state = State.empty;
233 state.add(&first);
234 state.add(&second);
235
236 var expected = State.empty;
237 expected.add(&first);
238
239 state.subtract(&second);
240 try std.testing.expect(state.eql(&expected));
241
242 state.subtract(&first);
243 try std.testing.expect(state.isEmpty());
244 }
245
246 test "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));
250
251 const key_only = entryState("a", "");
252 const value_only = entryState("", "a");
253 try std.testing.expect(!key_only.eql(&value_only));
254 }
255
256 test "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();
262
263 const oneshot = entryState("stream", "hello");
264 try std.testing.expect(streamed.eql(&oneshot));
265 }
266
267 test "lattice encode decode round trips" {
268 var state = State.empty;
269 const entry = entryState("round", "trip");
270 state.add(&entry);
271
272 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 }
277
278 test "lattice digest binds entry count" {
279 var state = State.empty;
280 const entry = entryState("count", "bound");
281 state.add(&entry);
282
283 const one = state.digest(1);
284 const two = state.digest(2);
285 try std.testing.expect(!std.mem.eql(u8, &one, &two));
286 }
287
288 test "lattice empty digest differs from populated digest" {
289 var state = State.empty;
290 const empty_digest = State.empty.digest(0);
291
292 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 }