diff --git a/src/bson.zig b/src/bson.zig index 8b7c80c..cd8158d 100644 --- a/src/bson.zig +++ b/src/bson.zig @@ -557,6 +557,126 @@ fn rank(v: Value) u8 { }; } +// --------------------------------------------------------------------------- +// Order-preserving key encoding +// --------------------------------------------------------------------------- + +/// Encode `v` so that `std.mem.order(u8, encode(a), encode(b))` equals +/// `compare(a, b)` for every pair of values. That equivalence is the whole +/// point: it turns index keys into plain byte strings, so an index can +/// binary-search, range-scan and be stored as bytes without carrying +/// `Value` trees around and chasing pointers into per-document arenas. +/// +/// Layout is `[rank + 1]` then a self-delimiting payload. The `+ 1` keeps +/// `0x00` out of the tag space so it can terminate variable-length payloads +/// everywhere. +/// +/// Every payload is self-delimiting, which is what lets several columns be +/// concatenated into one key and still support prefix search. +pub fn encode_key(v: Value, gpa: std.mem.Allocator, out: *std.ArrayListUnmanaged(u8)) !void { + try out.append(gpa, rank(v) + 1); + switch (v) { + // Ranks that compare `.eq` to themselves carry no payload, so equal + // values produce identical bytes. + .min_key, .null, .max_key => {}, + + .double, .int32, .int64 => try encode_number(v.as_f128(), gpa, out), + + // One rank for all three, so "x" as a string and as a symbol encode + // identically — exactly as compare treats them. + .string, .symbol, .code => try encode_escaped(as_str(v), gpa, out), + + // compare_docs: key, then value, then length. The terminator gives + // the length tie-break for free, since it sorts below any real key. + .doc => |pairs| { + for (pairs) |p| { + try encode_escaped(p.key, gpa, out); + try encode_key(p.value, gpa, out); + } + try out.appendSlice(gpa, &.{ 0x00, 0x00 }); + }, + .array => |items| { + for (items) |item| try encode_key(item, gpa, out); + try out.appendSlice(gpa, &.{ 0x00, 0x00 }); + }, + + // compare_binary orders by length first, then bytes, then subtype. + .binary => |b| { + var len_be: [4]u8 = undefined; + std.mem.writeInt(u32, &len_be, @intCast(b.data.len), .big); + try out.appendSlice(gpa, &len_be); + try out.appendSlice(gpa, b.data); + try out.append(gpa, b.subtype); + }, + + // Fixed width, compared as raw bytes. + .object_id => |oid| try out.appendSlice(gpa, &oid), + .decimal128 => |d| try out.appendSlice(gpa, &d), + .bool => |b| try out.append(gpa, @intFromBool(b)), + + // Flip the sign bit so the two's-complement order becomes unsigned + // byte order. + .datetime => |ms| { + var buf: [8]u8 = undefined; + std.mem.writeInt(u64, &buf, @as(u64, @bitCast(ms)) ^ (1 << 63), .big); + try out.appendSlice(gpa, &buf); + }, + .timestamp => |ts| { + var buf: [8]u8 = undefined; + std.mem.writeInt(u64, &buf, ts, .big); + try out.appendSlice(gpa, &buf); + }, + + .regex => |r| { + try encode_escaped(r.pattern, gpa, out); + try encode_escaped(r.options, gpa, out); + }, + + // compare ignores `kind` and orders the data lexicographically, so + // this escapes rather than length-prefixes. + .opaque_val => |o| try encode_escaped(o.data, gpa, out), + } +} + +/// Byte string, terminated so it stays self-delimiting, with the terminator +/// ordering below any content. +/// +/// A BSON string may contain NUL, so a bare `00` terminator would be +/// ambiguous. Escaping `00` as `00 FF` fixes both problems at once: a real +/// NUL encodes above the `00 00` terminator, and any byte >= 01 is above it +/// too, so "shorter is less" falls out to match `std.mem.order`. +fn encode_escaped(bytes: []const u8, gpa: std.mem.Allocator, out: *std.ArrayListUnmanaged(u8)) !void { + for (bytes) |b| { + try out.append(gpa, b); + if (b == 0x00) try out.append(gpa, 0xFF); + } + try out.appendSlice(gpa, &.{ 0x00, 0x00 }); +} + +/// Map an f128 onto 16 bytes whose unsigned order is the numeric order. +/// +/// `compare` widens int32, int64 and double to f128 before comparing, and +/// that widening is exact for all three. Encoding the widened value is what +/// makes `int32 1`, `int64 1` and `double 1.0` produce identical bytes -- +/// the cross-type equality that index lookups on a numeric field depend on. +fn encode_number(x: f128, gpa: std.mem.Allocator, out: *std.ArrayListUnmanaged(u8)) !void { + // compare_f128 makes every NaN equal and greater than everything else. + // All-ones is above any finite or infinite encoding below. + if (std.math.isNan(x)) { + try out.appendNTimes(gpa, 0xFF, 16); + return; + } + // -0.0 and +0.0 compare equal, so they must encode identically. + const normalized: f128 = if (x == 0) 0 else x; + const bits: u128 = @bitCast(normalized); + // Negatives: invert everything, so more-negative sorts lower. + // Positives: set the sign bit, lifting them above every negative. + const mono: u128 = if (bits >> 127 != 0) ~bits else bits | (@as(u128, 1) << 127); + var buf: [16]u8 = undefined; + std.mem.writeInt(u128, &buf, mono, .big); + try out.appendSlice(gpa, &buf); +} + pub fn compare(a: Value, b: Value) std.math.Order { const ra = rank(a); const rb = rank(b); @@ -752,6 +872,126 @@ test "compare: NaN is greatest number" { try testing.expectEqual(std.math.Order.gt, compare(.{ .double = std.math.nan(f64) }, .{ .int64 = 1 << 62 })); } +test "encode_key order matches bson.compare on every pair" { + // This equivalence is what the whole encoded-key index rests on: if it + // holds, byte comparison can replace value comparison everywhere. + // Check it exhaustively over a corpus spanning every rank and the + // boundaries inside them. + const gpa = testing.allocator; + + const nested = [_]Pair{.{ .key = "k", .value = .{ .int32 = 1 } }}; + const nested2 = [_]Pair{.{ .key = "k", .value = .{ .int32 = 2 } }}; + const nested_l = [_]Pair{.{ .key = "l", .value = .{ .int32 = 1 } }}; + const two_pairs = [_]Pair{ + .{ .key = "k", .value = .{ .int32 = 1 } }, + .{ .key = "z", .value = .{ .int32 = 1 } }, + }; + const arr1 = [_]Value{.{ .int32 = 1 }}; + const arr2 = [_]Value{ .{ .int32 = 1 }, .{ .int32 = 2 } }; + const arr_str = [_]Value{.{ .string = "a" }}; + + const corpus = [_]Value{ + .min_key, + .null, + // Numbers: cross-type equality, sign, zero, extremes, NaN. + .{ .int32 = -2147483648 }, .{ .int32 = -1 }, .{ .int32 = 0 }, + .{ .int32 = 1 }, .{ .int32 = 2 }, .{ .int32 = 2147483647 }, + .{ .int64 = -9223372036854775807 }, .{ .int64 = -1 }, .{ .int64 = 0 }, + .{ .int64 = 1 }, .{ .int64 = 9223372036854775807 }, + .{ .double = -std.math.inf(f64) }, .{ .double = -1.5 }, .{ .double = -0.0 }, + .{ .double = 0.0 }, .{ .double = 0.5 }, .{ .double = 1.0 }, + .{ .double = 1.5 }, .{ .double = std.math.inf(f64) }, + .{ .double = std.math.nan(f64) }, + // Strings, including embedded NUL and prefix relationships. + .{ .string = "" }, .{ .string = "\x00" }, .{ .string = "\x00b" }, + .{ .string = "a" }, .{ .string = "a\x00" }, .{ .string = "a\x00b" }, + .{ .string = "ab" }, .{ .string = "b" }, .{ .string = "\xff" }, + // Same rank as string, so these must interleave with them. + .{ .symbol = "a" }, .{ .code = "ab" }, + .{ .doc = &.{} }, .{ .doc = &nested }, .{ .doc = &nested2 }, + .{ .doc = &nested_l }, .{ .doc = &two_pairs }, + .{ .array = &.{} }, .{ .array = &arr1 }, .{ .array = &arr2 }, + .{ .array = &arr_str }, + // Binary orders by length first, then bytes, then subtype. + .{ .binary = .{ .subtype = 0, .data = "" } }, + .{ .binary = .{ .subtype = 0, .data = "a" } }, + .{ .binary = .{ .subtype = 1, .data = "a" } }, + .{ .binary = .{ .subtype = 0, .data = "b" } }, + .{ .binary = .{ .subtype = 0, .data = "aa" } }, + .{ .object_id = [_]u8{0} ** 12 }, + .{ .object_id = [_]u8{0} ** 11 ++ [_]u8{1} }, + .{ .object_id = [_]u8{255} ** 12 }, + .{ .bool = false }, .{ .bool = true }, + .{ .datetime = std.math.minInt(i64) }, .{ .datetime = -1 }, + .{ .datetime = 0 }, .{ .datetime = 1 }, + .{ .datetime = std.math.maxInt(i64) }, + .{ .timestamp = 0 }, .{ .timestamp = 1 }, + .{ .timestamp = std.math.maxInt(u64) }, + .{ .regex = .{ .pattern = "a", .options = "" } }, + .{ .regex = .{ .pattern = "a", .options = "i" } }, + .{ .regex = .{ .pattern = "b", .options = "" } }, + .{ .opaque_val = .{ .kind = 1, .data = "a" } }, + // compare ignores kind, so this must encode equal to the one above. + .{ .opaque_val = .{ .kind = 2, .data = "a" } }, + .{ .opaque_val = .{ .kind = 1, .data = "ab" } }, + .{ .decimal128 = [_]u8{0} ** 16 }, + .{ .decimal128 = [_]u8{255} ** 16 }, + .max_key, + }; + + var encoded: [corpus.len]std.ArrayListUnmanaged(u8) = undefined; + for (&encoded) |*e| e.* = .empty; + defer for (&encoded) |*e| e.deinit(gpa); + for (corpus, 0..) |v, i| try encode_key(v, gpa, &encoded[i]); + + for (corpus, 0..) |a, i| { + for (corpus, 0..) |b, j| { + const want = compare(a, b); + const got = std.mem.order(u8, encoded[i].items, encoded[j].items); + testing.expectEqual(want, got) catch |err| { + std.debug.print( + "pair ({d},{d}) {s} vs {s}: compare={s} encoded={s}\n", + .{ i, j, @tagName(a), @tagName(b), @tagName(want), @tagName(got) }, + ); + return err; + }; + } + } +} + +test "encode_key concatenates into unambiguous compound keys" { + // Multi-column index keys are just concatenated encodings, so a column + // must never be able to run into the next one. Compare two-column keys + // against the component-wise order they are supposed to reproduce. + const gpa = testing.allocator; + const parts = [_]Value{ + .{ .string = "" }, .{ .string = "a" }, .{ .string = "a\x00" }, + .{ .string = "ab" }, .{ .int32 = 1 }, .{ .int32 = 2 }, + .null, .{ .array = &.{} }, .{ .bool = true }, + }; + + for (parts) |a1| { + for (parts) |a2| { + for (parts) |b1| { + for (parts) |b2| { + var ka: std.ArrayListUnmanaged(u8) = .empty; + defer ka.deinit(gpa); + var kb: std.ArrayListUnmanaged(u8) = .empty; + defer kb.deinit(gpa); + try encode_key(a1, gpa, &ka); + try encode_key(a2, gpa, &ka); + try encode_key(b1, gpa, &kb); + try encode_key(b2, gpa, &kb); + + const first = compare(a1, b1); + const want = if (first != .eq) first else compare(a2, b2); + try testing.expectEqual(want, std.mem.order(u8, ka.items, kb.items)); + } + } + } + } +} + test "object id generation" { var threaded: std.Io.Threaded = .init_single_threaded; defer threaded.deinit();