//! BSON — Binary JSON. The foundation of the MongoDB wire protocol and the
//! storage engine. Documents parse into an arena-backed value tree; the tree
//! serializes back to canonical BSON bytes.

const std = @import("std");

pub const ObjectId = [12]u8;

pub const Binary = struct {
    subtype: u8,
    data: []const u8,
};

pub const Regex = struct {
    pattern: []const u8,
    options: []const u8,
};

/// Values we round-trip but never interpret: db_pointer (0x0C),
/// code_with_scope (0x0F), undefined (0x06). `data` is the raw payload that
/// follows the element type byte, re-emitted verbatim on serialize.
pub const Opaque = struct {
    kind: u8,
    data: []const u8,
};

pub const Pair = struct {
    key: []const u8,
    value: Value,
};

pub const Value = union(enum) {
    double: f64,
    string: []const u8,
    doc: []const Pair,
    array: []const Value,
    binary: Binary,
    object_id: ObjectId,
    bool: bool,
    datetime: i64,
    null,
    regex: Regex,
    code: []const u8,
    symbol: []const u8,
    int32: i32,
    timestamp: u64,
    int64: i64,
    decimal128: [16]u8,
    min_key,
    max_key,
    opaque_val: Opaque,

    pub fn type_tag(self: Value) u8 {
        return switch (self) {
            .double => 0x01,
            .string => 0x02,
            .doc => 0x03,
            .array => 0x04,
            .binary => 0x05,
            .object_id => 0x07,
            .bool => 0x08,
            .datetime => 0x09,
            .null => 0x0A,
            .regex => 0x0B,
            .code => 0x0D,
            .symbol => 0x0E,
            .int32 => 0x10,
            .timestamp => 0x11,
            .int64 => 0x12,
            .decimal128 => 0x13,
            .min_key => 0xFF, // not serializable; rank only
            .max_key => 0x7F, // not serializable; rank only
            .opaque_val => |o| o.kind,
        };
    }

    /// The name mongod uses for this type in a TypeMismatch message ("is the
    /// wrong type 'int', expected type 'object'"). Its own names, not Zig's:
    /// a driver that matches on the text is matching on these.
    pub fn type_name(self: Value) []const u8 {
        return switch (self) {
            .double => "double",
            .string => "string",
            .doc => "object",
            .array => "array",
            .binary => "binData",
            .object_id => "objectId",
            .bool => "bool",
            .datetime => "date",
            .null => "null",
            .regex => "regex",
            .code => "javascript",
            .symbol => "symbol",
            .int32 => "int",
            .timestamp => "timestamp",
            .int64 => "long",
            .decimal128 => "decimal",
            .min_key => "minKey",
            .max_key => "maxKey",
            // An unparsed value keeps only its tag byte, and the tags this
            // union does not name are the ones nothing here inspects.
            .opaque_val => "unknown",
        };
    }

    pub fn is_number(self: Value) bool {
        return switch (self) {
            .double, .int32, .int64 => true,
            else => false,
        };
    }

    /// Numeric value widened to f128 — exact for i64 and f64.
    pub fn as_f128(self: Value) f128 {
        return switch (self) {
            .double => |d| @as(f128, @floatCast(d)),
            .int32 => |i| @as(f128, @floatFromInt(i)),
            .int64 => |i| @as(f128, @floatFromInt(i)),
            else => unreachable,
        };
    }
};

/// A parsed document owns everything it references via its arena. Not
/// copyable — pass by pointer.
pub const Document = struct {
    arena: std.heap.ArenaAllocator,
    pairs: []const Pair,

    pub fn parse(allocator: std.mem.Allocator, bytes: []const u8) !Document {
        var arena = std.heap.ArenaAllocator.init(allocator);
        errdefer arena.deinit();
        var idx: usize = 0;
        const pairs = try parse_doc_into(arena.allocator(), bytes, &idx, false);
        return .{ .arena = arena, .pairs = pairs };
    }

    pub fn deinit(self: *Document) void {
        self.arena.deinit();
    }

    pub fn get(self: *const Document, key: []const u8) ?Value {
        return get_pair(self.pairs, key);
    }

    pub fn alloc(allocator: std.mem.Allocator, pairs: []const Pair) !Document {
        var arena = std.heap.ArenaAllocator.init(allocator);
        errdefer arena.deinit();
        const copied = try arena.allocator().dupe(Pair, pairs);
        return .{ .arena = arena, .pairs = copied };
    }

    /// Serialize the full document (length-prefixed) into `out`.
    pub fn to_bytes(
        self: *const Document,
        gpa: std.mem.Allocator,
        out: *std.ArrayListUnmanaged(u8),
    ) !void {
        try write_doc(self.pairs, gpa, out);
    }
};

pub fn get_pair(pairs: []const Pair, key: []const u8) ?Value {
    const i = get_pair_index(pairs, key) orelse return null;
    return pairs[i].value;
}

/// Position of `key` in `pairs` — for callers that need to replace or remove
/// the pair in place rather than just read it.
pub fn get_pair_index(pairs: []const Pair, key: []const u8) ?usize {
    for (pairs, 0..) |p, i| {
        if (std.mem.eql(u8, p.key, key)) return i;
    }
    return null;
}

// ---------------------------------------------------------------------------
// Parsing
// ---------------------------------------------------------------------------

const Parser = struct {
    allocator: std.mem.Allocator,
    bytes: []const u8,
    /// When set, strings and keys point into `bytes` instead of being
    /// copied, so the parsed spine is only valid while `bytes` is. Used for
    /// the borrowed spine of slab-resident documents; the full `parse`
    /// keeps the owned, self-contained behavior.
    borrow: bool,
};

const ParseError = error{ InvalidBson, OutOfMemory };

fn parse_doc_into(
    allocator: std.mem.Allocator,
    bytes: []const u8,
    idx: *usize,
    borrow: bool,
) ParseError![]const Pair {
    const p = Parser{ .allocator = allocator, .bytes = bytes, .borrow = borrow };
    return parse_doc_inner(p, idx);
}

fn ensure_available(bytes: []const u8, idx: usize, n: usize) error{InvalidBson}!void {
    if (bytes.len -| idx < n) return error.InvalidBson;
}

/// Validate the length-prefixed frame starting at `start` — documents and
/// arrays share it — and return the offset just past its terminating byte.
fn doc_extent(p: Parser, start: usize) ParseError!usize {
    try ensure_available(p.bytes, start, 4);
    const total: u32 = std.mem.readInt(u32, p.bytes[start..][0..4], .little);
    if (total < 5) return error.InvalidBson;
    if (p.bytes.len - start < total) return error.InvalidBson;
    const end = start + total;
    if (p.bytes[end - 1] != 0x00) return error.InvalidBson;
    return end;
}

fn parse_doc_inner(p: Parser, idx: *usize) ParseError![]const Pair {
    const start = idx.*;
    const end = try doc_extent(p, start);

    const gpa = p.allocator;
    var pairs: std.ArrayListUnmanaged(Pair) = .empty;
    errdefer pairs.deinit(gpa);

    idx.* = start + 4;
    while (idx.* < end - 1) {
        const value = try parse_element(p, idx);
        try pairs.append(gpa, value);
    }
    if (idx.* != end - 1) return error.InvalidBson;
    idx.* = end;
    return pairs.toOwnedSlice(gpa);
}

fn parse_element(p: Parser, idx: *usize) ParseError!Pair {
    try ensure_available(p.bytes, idx.*, 1);
    const tag = p.bytes[idx.*];
    idx.* += 1;
    const key = try parse_cstring(p, idx);
    const value = try parse_value(p, tag, idx);
    return .{ .key = key, .value = value };
}

fn parse_cstring(p: Parser, idx: *usize) ParseError![]const u8 {
    const start = idx.*;
    while (idx.* < p.bytes.len and p.bytes[idx.*] != 0) idx.* += 1;
    if (idx.* >= p.bytes.len) return error.InvalidBson;
    idx.* += 1;
    // Keys and strings are copied into the arena so owned documents are
    // self-contained and outlive the input buffer (wire messages and log
    // records are transient); a borrowed spine leaves them pointing at the
    // source bytes instead.
    if (p.borrow) return p.bytes[start .. idx.* - 1];
    return p.allocator.dupe(u8, p.bytes[start .. idx.* - 1]);
}

fn parse_string(p: Parser, idx: *usize) ParseError![]const u8 {
    try ensure_available(p.bytes, idx.*, 4);
    const len: u32 = std.mem.readInt(u32, p.bytes[idx.*..][0..4], .little);
    if (len == 0 or len > p.bytes.len - (idx.* + 4)) return error.InvalidBson;
    const str = p.bytes[idx.* + 4 .. idx.* + 4 + len];
    if (str[len - 1] != 0) return error.InvalidBson;
    idx.* += 4 + len;
    if (p.borrow) return str[0 .. len - 1];
    return p.allocator.dupe(u8, str[0 .. len - 1]);
}

fn parse_value(p: Parser, tag: u8, idx: *usize) ParseError!Value {
    return switch (tag) {
        0x01 => blk: {
            try ensure_available(p.bytes, idx.*, 8);
            const v: f64 = @bitCast(std.mem.readInt(u64, p.bytes[idx.*..][0..8], .little));
            idx.* += 8;
            break :blk .{ .double = v };
        },
        0x02 => .{ .string = try parse_string(p, idx) },
        0x03 => .{ .doc = try parse_doc_inner(p, idx) },
        0x04 => .{ .array = try parse_array(p, idx) },
        0x05 => blk: {
            try ensure_available(p.bytes, idx.*, 5);
            const len: u32 = std.mem.readInt(u32, p.bytes[idx.*..][0..4], .little);
            const subtype = p.bytes[idx.* + 4];
            try ensure_available(p.bytes, idx.* + 5, len);
            const data = p.bytes[idx.* + 5 .. idx.* + 5 + len];
            idx.* += 5 + len;
            break :blk .{ .binary = .{ .subtype = subtype, .data = try p.allocator.dupe(u8, data) } };
        },
        0x06 => .{ .opaque_val = .{ .kind = 0x06, .data = &.{} } },
        0x07 => blk: {
            try ensure_available(p.bytes, idx.*, 12);
            const oid: ObjectId = p.bytes[idx.*..][0..12].*;
            idx.* += 12;
            break :blk .{ .object_id = oid };
        },
        0x08 => blk: {
            try ensure_available(p.bytes, idx.*, 1);
            const v = p.bytes[idx.*];
            if (v > 1) return error.InvalidBson;
            idx.* += 1;
            break :blk .{ .bool = v == 1 };
        },
        0x09 => blk: {
            try ensure_available(p.bytes, idx.*, 8);
            const v: i64 = std.mem.readInt(i64, p.bytes[idx.*..][0..8], .little);
            idx.* += 8;
            break :blk .{ .datetime = v };
        },
        0x0A => .null,
        0x0B => .{ .regex = .{
            .pattern = try parse_cstring(p, idx),
            .options = try parse_cstring(p, idx),
        } },
        0x0C => blk: {
            const start = idx.*;
            _ = try parse_string(p, idx);
            try ensure_available(p.bytes, idx.*, 12);
            idx.* += 12;
            // Like all other types, the payload is copied into the arena so
            // documents stay valid after the input buffer is reused.
            break :blk .{ .opaque_val = .{ .kind = 0x0C, .data = try p.allocator.dupe(u8, p.bytes[start..idx.*]) } };
        },
        0x0D => .{ .code = try parse_string(p, idx) },
        0x0E => .{ .symbol = try parse_string(p, idx) },
        0x0F => blk: {
            try ensure_available(p.bytes, idx.*, 4);
            const total: u32 = std.mem.readInt(u32, p.bytes[idx.*..][0..4], .little);
            if (total < 4 or total > p.bytes.len - idx.*) return error.InvalidBson;
            const data = p.bytes[idx.* .. idx.* + total];
            idx.* += total;
            break :blk .{ .opaque_val = .{ .kind = 0x0F, .data = try p.allocator.dupe(u8, data) } };
        },
        0x10 => blk: {
            try ensure_available(p.bytes, idx.*, 4);
            const v: i32 = std.mem.readInt(i32, p.bytes[idx.*..][0..4], .little);
            idx.* += 4;
            break :blk .{ .int32 = v };
        },
        0x11 => blk: {
            try ensure_available(p.bytes, idx.*, 8);
            const v: u64 = std.mem.readInt(u64, p.bytes[idx.*..][0..8], .little);
            idx.* += 8;
            break :blk .{ .timestamp = v };
        },
        0x12 => blk: {
            try ensure_available(p.bytes, idx.*, 8);
            const v: i64 = std.mem.readInt(i64, p.bytes[idx.*..][0..8], .little);
            idx.* += 8;
            break :blk .{ .int64 = v };
        },
        0x13 => blk: {
            try ensure_available(p.bytes, idx.*, 16);
            const v: [16]u8 = p.bytes[idx.*..][0..16].*;
            idx.* += 16;
            break :blk .{ .decimal128 = v };
        },
        else => return error.InvalidBson,
    };
}

fn parse_array(p: Parser, idx: *usize) ParseError![]const Value {
    const start = idx.*;
    const end = try doc_extent(p, start);

    const gpa = p.allocator;
    var values: std.ArrayListUnmanaged(Value) = .empty;
    errdefer values.deinit(gpa);

    idx.* = start + 4;
    while (idx.* < end - 1) {
        try ensure_available(p.bytes, idx.*, 1);
        const tag = p.bytes[idx.*];
        idx.* += 1;
        _ = try parse_cstring(p, idx);
        try values.append(gpa, try parse_value(p, tag, idx));
    }
    if (idx.* != end - 1) return error.InvalidBson;
    idx.* = end;
    return values.toOwnedSlice(gpa);
}

// ---------------------------------------------------------------------------
// Borrowed spine and byte walking
// ---------------------------------------------------------------------------
//
// Slab-resident documents are stored as canonical BSON bytes; the tree
// machinery needs a `[]const Pair` spine over them. `spine` parses one with
// keys and leaf values pointing at the source bytes (valid while the bytes
// live — the collection slab outlives any query). The byte-walking
// primitives below let a matcher skip over fields it does not name instead
// of materializing the whole document.

/// A borrowed parse of `bytes` into an arena: keys, strings, binary and
/// regex payloads point into `bytes`; only the pair/value skeleton is
/// allocated. The result is valid while both `bytes` and `arena` live.
pub fn spine(
    allocator: std.mem.Allocator,
    bytes: []const u8,
) error{ InvalidBson, OutOfMemory }![]const Pair {
    var idx: usize = 0;
    return parse_doc_into(allocator, bytes, &idx, true);
}

/// The first element's key, or null when the remaining bytes are not an
/// element. Advances `idx` past the key.
pub fn element_key(bytes: []const u8, idx: *usize) ?[]const u8 {
    if (idx.* >= bytes.len) return null;
    const start = idx.*;
    while (idx.* < bytes.len and bytes[idx.*] != 0) idx.* += 1;
    if (idx.* >= bytes.len) return null;
    idx.* += 1;
    return bytes[start .. idx.* - 1];
}

/// Advance past one value of `tag` (the key is already consumed). The input
/// is canonical (validated when it was stored), so skips are cheap; the
/// length-prefixed types still bounds-check their prefix.
pub fn skip_value(bytes: []const u8, idx: *usize, tag: u8) error{InvalidBson}!void {
    switch (tag) {
        0x01, 0x09, 0x11, 0x12 => idx.* += 8,
        0x10 => idx.* += 4,
        0x13 => idx.* += 16,
        0x07 => idx.* += 12,
        0x08 => idx.* += 1,
        0x0A, 0x06 => {},
        0x02, 0x0D, 0x0E => {
            try ensure_available(bytes, idx.*, 4);
            idx.* += 4 + std.mem.readInt(u32, bytes[idx.*..][0..4], .little);
        },
        0x03, 0x04 => {
            try ensure_available(bytes, idx.*, 4);
            idx.* += std.mem.readInt(u32, bytes[idx.*..][0..4], .little);
        },
        0x05 => {
            try ensure_available(bytes, idx.*, 5);
            idx.* += 5 + std.mem.readInt(u32, bytes[idx.*..][0..4], .little);
        },
        0x0B => {
            _ = element_key(bytes, idx) orelse return error.InvalidBson;
            _ = element_key(bytes, idx) orelse return error.InvalidBson;
        },
        0x0C => {
            try ensure_available(bytes, idx.*, 4);
            idx.* += 4 + std.mem.readInt(u32, bytes[idx.*..][0..4], .little) + 12;
        },
        0x0F => {
            try ensure_available(bytes, idx.*, 4);
            idx.* += std.mem.readInt(u32, bytes[idx.*..][0..4], .little);
        },
        else => return error.InvalidBson,
    }
}

/// Read one value of `tag` into a Value whose leaves borrow `bytes`; nested
/// documents and arrays materialize their spines into `arena`. Advances
/// `idx` past the value.
pub fn read_value(
    allocator: std.mem.Allocator,
    bytes: []const u8,
    idx: *usize,
    tag: u8,
) error{ InvalidBson, OutOfMemory }!Value {
    switch (tag) {
        0x01 => {
            try ensure_available(bytes, idx.*, 8);
            const v: f64 = @bitCast(std.mem.readInt(u64, bytes[idx.*..][0..8], .little));
            idx.* += 8;
            return .{ .double = v };
        },
        0x02 => {
            try ensure_available(bytes, idx.*, 4);
            const len: u32 = std.mem.readInt(u32, bytes[idx.*..][0..4], .little);
            const str = bytes[idx.* + 4 .. idx.* + 4 + len];
            idx.* += 4 + len;
            return .{ .string = str[0 .. len - 1] };
        },
        0x03 => {
            const start = idx.*;
            const total: u32 = std.mem.readInt(u32, bytes[start..][0..4], .little);
            var sub: usize = 0;
            const pairs = try parse_doc_into(allocator, bytes[start..], &sub, true);
            idx.* = start + total;
            return .{ .doc = pairs };
        },
        0x04 => {
            const start = idx.*;
            const total: u32 = std.mem.readInt(u32, bytes[start..][0..4], .little);
            var sub: usize = start + 4;
            var values: std.ArrayListUnmanaged(Value) = .empty;
            errdefer values.deinit(allocator);
            while (sub < start + total - 1 and bytes[sub] != 0) {
                const atag = bytes[sub];
                sub += 1;
                _ = element_key(bytes, &sub) orelse return error.InvalidBson;
                try values.append(allocator, try read_value(allocator, bytes, &sub, atag));
            }
            idx.* = start + total;
            return .{ .array = try values.toOwnedSlice(allocator) };
        },
        0x05 => {
            try ensure_available(bytes, idx.*, 5);
            const len: u32 = std.mem.readInt(u32, bytes[idx.*..][0..4], .little);
            const subtype = bytes[idx.* + 4];
            const data = bytes[idx.* + 5 .. idx.* + 5 + len];
            idx.* += 5 + len;
            return .{ .binary = .{ .subtype = subtype, .data = data } };
        },
        0x07 => {
            try ensure_available(bytes, idx.*, 12);
            const oid: ObjectId = bytes[idx.*..][0..12].*;
            idx.* += 12;
            return .{ .object_id = oid };
        },
        0x08 => {
            try ensure_available(bytes, idx.*, 1);
            const v = bytes[idx.*] != 0;
            idx.* += 1;
            return .{ .bool = v };
        },
        0x09 => {
            try ensure_available(bytes, idx.*, 8);
            const v: i64 = std.mem.readInt(i64, bytes[idx.*..][0..8], .little);
            idx.* += 8;
            return .{ .datetime = v };
        },
        0x0A => return .null,
        0x0B => {
            const pattern = element_key(bytes, idx) orelse return error.InvalidBson;
            const options = element_key(bytes, idx) orelse return error.InvalidBson;
            return .{ .regex = .{ .pattern = pattern, .options = options } };
        },
        0x0C => {
            try ensure_available(bytes, idx.*, 4);
            const len: u32 = std.mem.readInt(u32, bytes[idx.*..][0..4], .little);
            const start = idx.*;
            idx.* += 4 + len + 12;
            return .{ .opaque_val = .{ .kind = 0x0C, .data = bytes[start..idx.*] } };
        },
        0x0D => {
            try ensure_available(bytes, idx.*, 4);
            const len: u32 = std.mem.readInt(u32, bytes[idx.*..][0..4], .little);
            const str = bytes[idx.* + 4 .. idx.* + 4 + len];
            idx.* += 4 + len;
            return .{ .code = str[0 .. len - 1] };
        },
        0x0E => {
            try ensure_available(bytes, idx.*, 4);
            const len: u32 = std.mem.readInt(u32, bytes[idx.*..][0..4], .little);
            const str = bytes[idx.* + 4 .. idx.* + 4 + len];
            idx.* += 4 + len;
            return .{ .symbol = str[0 .. len - 1] };
        },
        0x0F => {
            try ensure_available(bytes, idx.*, 4);
            const total: u32 = std.mem.readInt(u32, bytes[idx.*..][0..4], .little);
            const start = idx.*;
            idx.* += total;
            return .{ .opaque_val = .{ .kind = 0x0F, .data = bytes[start..idx.*] } };
        },
        0x10 => {
            try ensure_available(bytes, idx.*, 4);
            const v: i32 = std.mem.readInt(i32, bytes[idx.*..][0..4], .little);
            idx.* += 4;
            return .{ .int32 = v };
        },
        0x11 => {
            try ensure_available(bytes, idx.*, 8);
            const v: u64 = std.mem.readInt(u64, bytes[idx.*..][0..8], .little);
            idx.* += 8;
            return .{ .timestamp = v };
        },
        0x12 => {
            try ensure_available(bytes, idx.*, 8);
            const v: i64 = std.mem.readInt(i64, bytes[idx.*..][0..8], .little);
            idx.* += 8;
            return .{ .int64 = v };
        },
        0x13 => {
            try ensure_available(bytes, idx.*, 16);
            const v: [16]u8 = bytes[idx.*..][0..16].*;
            idx.* += 16;
            return .{ .decimal128 = v };
        },
        else => return error.InvalidBson,
    }
}

/// The value stored under `key` in a document's bytes, or null when absent.
/// Nested documents and arrays materialize their spines into `arena`.
pub fn get_at(
    arena: std.mem.Allocator,
    bytes: []const u8,
    key: []const u8,
) error{ InvalidBson, OutOfMemory }!?Value {
    var idx: usize = 4;
    while (idx + 1 < bytes.len and bytes[idx] != 0) {
        const tag = bytes[idx];
        idx += 1;
        const k = element_key(bytes, &idx) orelse return error.InvalidBson;
        if (std.mem.eql(u8, k, key)) return try read_value(arena, bytes, &idx, tag);
        try skip_value(bytes, &idx, tag);
    }
    return null;
}

// ---------------------------------------------------------------------------
// Serialization
// ---------------------------------------------------------------------------

pub const SerializeError = error{
    BsonTooLarge,
    BsonNulInKey,
    BsonNotSerializable,
    OutOfMemory,
};

pub fn write_value(
    v: Value,
    gpa: std.mem.Allocator,
    out: *std.ArrayListUnmanaged(u8),
) SerializeError!void {
    switch (v) {
        .double => |d| {
            var buf: [8]u8 = undefined;
            std.mem.writeInt(u64, &buf, @bitCast(d), .little);
            try out.appendSlice(gpa, &buf);
        },
        .string => |s| try write_string(s, gpa, out),
        .doc => |pairs| try write_doc(pairs, gpa, out),
        .array => |items| try write_array(items, gpa, out),
        .binary => |b| {
            if (b.data.len > std.math.maxInt(u32)) return error.BsonTooLarge;
            var buf: [5]u8 = undefined;
            std.mem.writeInt(u32, buf[0..4], @intCast(b.data.len), .little);
            buf[4] = b.subtype;
            try out.appendSlice(gpa, &buf);
            try out.appendSlice(gpa, b.data);
        },
        .object_id => |oid| try out.appendSlice(gpa, &oid),
        .bool => |b| try out.append(gpa, @intFromBool(b)),
        .datetime => |t| {
            var buf: [8]u8 = undefined;
            std.mem.writeInt(i64, &buf, t, .little);
            try out.appendSlice(gpa, &buf);
        },
        .null => {},
        .regex => |r| {
            try write_cstring(r.pattern, gpa, out);
            try write_cstring(r.options, gpa, out);
        },
        .code => |c| try write_string(c, gpa, out),
        .symbol => |s| try write_string(s, gpa, out),
        .int32 => |i| {
            var buf: [4]u8 = undefined;
            std.mem.writeInt(i32, &buf, i, .little);
            try out.appendSlice(gpa, &buf);
        },
        .timestamp => |t| {
            var buf: [8]u8 = undefined;
            std.mem.writeInt(u64, &buf, t, .little);
            try out.appendSlice(gpa, &buf);
        },
        .int64 => |i| {
            var buf: [8]u8 = undefined;
            std.mem.writeInt(i64, &buf, i, .little);
            try out.appendSlice(gpa, &buf);
        },
        .decimal128 => |d| try out.appendSlice(gpa, &d),
        .opaque_val => |o| try out.appendSlice(gpa, o.data),
        .min_key, .max_key => return error.BsonNotSerializable,
    }
}

pub fn write_cstring(
    s: []const u8,
    gpa: std.mem.Allocator,
    out: *std.ArrayListUnmanaged(u8),
) SerializeError!void {
    if (std.mem.indexOfScalar(u8, s, 0) != null) return error.BsonNulInKey;
    try out.appendSlice(gpa, s);
    try out.append(gpa, 0);
}

pub fn write_string(
    s: []const u8,
    gpa: std.mem.Allocator,
    out: *std.ArrayListUnmanaged(u8),
) SerializeError!void {
    if (s.len + 1 > std.math.maxInt(u32)) return error.BsonTooLarge;
    var buf: [4]u8 = undefined;
    std.mem.writeInt(u32, &buf, @intCast(s.len + 1), .little);
    try out.appendSlice(gpa, &buf);
    try out.appendSlice(gpa, s);
    try out.append(gpa, 0);
}

pub fn write_element(
    pair: Pair,
    gpa: std.mem.Allocator,
    out: *std.ArrayListUnmanaged(u8),
) SerializeError!void {
    try out.append(gpa, pair.value.type_tag());
    try write_cstring(pair.key, gpa, out);
    try write_value(pair.value, gpa, out);
}

/// Reserve the 4-byte length prefix of a document or array frame. Returns the
/// offset to hand back to `end_frame`.
fn begin_frame(gpa: std.mem.Allocator, out: *std.ArrayListUnmanaged(u8)) SerializeError!usize {
    const len_pos = out.items.len;
    try out.appendSlice(gpa, &[4]u8{ 0, 0, 0, 0 });
    return len_pos;
}

/// Terminate the frame opened at `len_pos` and patch in its total length.
fn end_frame(
    gpa: std.mem.Allocator,
    out: *std.ArrayListUnmanaged(u8),
    len_pos: usize,
) SerializeError!void {
    try out.append(gpa, 0);
    const total = out.items.len - len_pos;
    if (total > std.math.maxInt(u32)) return error.BsonTooLarge;
    std.mem.writeInt(u32, out.items[len_pos..][0..4], @intCast(total), .little);
}

/// Write a length-prefixed document. Length is patched in after the body.
pub fn write_doc(
    pairs: []const Pair,
    gpa: std.mem.Allocator,
    out: *std.ArrayListUnmanaged(u8),
) SerializeError!void {
    const len_pos = try begin_frame(gpa, out);
    for (pairs) |p| try write_element(p, gpa, out);
    try end_frame(gpa, out, len_pos);
}

fn write_array(
    items: []const Value,
    gpa: std.mem.Allocator,
    out: *std.ArrayListUnmanaged(u8),
) SerializeError!void {
    const len_pos = try begin_frame(gpa, out);
    var buf: [16]u8 = undefined;
    for (items, 0..) |item, i| {
        try out.append(gpa, item.type_tag());
        const key = std.fmt.bufPrint(&buf, "{d}", .{i}) catch unreachable;
        try write_cstring(key, gpa, out);
        try write_value(item, gpa, out);
    }
    try end_frame(gpa, out, len_pos);
}

/// Serialize a single value with its type byte (no key) — used for `_id`
/// map keys and index entries.
pub fn serialize_value(gpa: std.mem.Allocator, v: Value) ![]u8 {
    var out: std.ArrayListUnmanaged(u8) = .empty;
    errdefer out.deinit(gpa);
    try write_serialized_value(v, gpa, &out);
    return out.toOwnedSlice(gpa);
}

/// Append the serialized-key bytes of `v` (type tag + payload) to `out`. The
/// appending form of serialize_value, for callers reusing one scratch buffer
/// across many keys.
pub fn write_serialized_value(
    v: Value,
    gpa: std.mem.Allocator,
    out: *std.ArrayListUnmanaged(u8),
) !void {
    try out.append(gpa, v.type_tag());
    try write_value(v, gpa, out);
}

/// Deep-copy a value into `arena`, so the copy is self-contained.
pub fn copy_value(arena: std.mem.Allocator, v: Value) std.mem.Allocator.Error!Value {
    return switch (v) {
        .string => |s| .{ .string = try arena.dupe(u8, s) },
        .symbol => |s| .{ .symbol = try arena.dupe(u8, s) },
        .code => |c| .{ .code = try arena.dupe(u8, c) },
        .doc => |pairs| .{ .doc = try copy_pairs(arena, pairs) },
        .array => |items| .{ .array = try copy_values(arena, items) },
        .binary => |b| .{ .binary = .{ .subtype = b.subtype, .data = try arena.dupe(u8, b.data) } },
        .regex => |r| .{ .regex = .{
            .pattern = try arena.dupe(u8, r.pattern),
            .options = try arena.dupe(u8, r.options),
        } },
        .opaque_val => |o| .{ .opaque_val = .{ .kind = o.kind, .data = try arena.dupe(u8, o.data) } },
        else => v,
    };
}

pub fn copy_pairs(
    arena: std.mem.Allocator,
    pairs: []const Pair,
) std.mem.Allocator.Error![]const Pair {
    const out = try arena.alloc(Pair, pairs.len);
    for (pairs, 0..) |p, i| {
        out[i] = .{ .key = try arena.dupe(u8, p.key), .value = try copy_value(arena, p.value) };
    }
    return out;
}

fn copy_values(
    arena: std.mem.Allocator,
    items: []const Value,
) std.mem.Allocator.Error![]const Value {
    const out = try arena.alloc(Value, items.len);
    for (items, 0..) |item, i| out[i] = try copy_value(arena, item);
    return out;
}

// ---------------------------------------------------------------------------
// ObjectId generation
// ---------------------------------------------------------------------------

pub const ObjectIdGen = struct {
    random_prefix: [5]u8,
    // Atomic so concurrent connections (e.g. two hellos, or a hello racing
    // an insert) can share one generator without a data race. The counter
    // only needs uniqueness within a second + random prefix, so monotonic
    // fetchAdd is fine.
    counter: std.atomic.Value(u32),

    pub fn init(io: std.Io) ObjectIdGen {
        var self: ObjectIdGen = undefined;
        io.random(&self.random_prefix);
        self.counter = .init(0);
        return self;
    }

    pub fn new(self: *ObjectIdGen, io: std.Io) ObjectId {
        var oid: ObjectId = undefined;
        const now = std.Io.Timestamp.now(io, .real);
        const secs: u32 = @truncate(@as(u64, @intCast(now.toSeconds())));
        std.mem.writeInt(u32, oid[0..4], secs, .big);
        @memcpy(oid[4..9], &self.random_prefix);
        const n = self.counter.fetchAdd(1, .monotonic) +% 1;
        std.mem.writeInt(u24, oid[9..12], @truncate(n), .big);
        return oid;
    }
};

// ---------------------------------------------------------------------------
// Canonical BSON comparison order
// ---------------------------------------------------------------------------

fn rank(v: Value) u8 {
    return switch (v) {
        .min_key => 0,
        .null => 1,
        .double, .int32, .int64 => 2,
        .string, .symbol, .code => 3,
        .doc => 4,
        .array => 5,
        .binary => 6,
        .object_id => 7,
        .bool => 8,
        .datetime => 9,
        .timestamp => 10,
        .regex => 11,
        .opaque_val => 12,
        .decimal128 => 13,
        .max_key => 14,
    };
}

// ---------------------------------------------------------------------------
// 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| {
            std.debug.assert(rank(v) + 1 == encoded_datetime_tag);
            var buf: [8]u8 = undefined;
            std.mem.writeInt(u64, &buf, @as(u64, @bitCast(ms)) ^ (@as(u64, 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),
    }
}

/// The tag `encode_key` writes for a datetime.
pub const encoded_datetime_tag: u8 = 9 + 1;

/// The datetime in an encoded key's first column, or null when that column
/// holds anything else. Lets a TTL sweep read entry keys without decoding
/// them back into Values.
pub fn encoded_leading_datetime(key: []const u8) ?i64 {
    if (key.len < 1 + 8 or key[0] != encoded_datetime_tag) return null;
    const biased = std.mem.readInt(u64, key[1..9], .big);
    return @bitCast(biased ^ (@as(u64, 1) << 63));
}

/// 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);
    if (ra != rb) return std.math.order(ra, rb);
    return switch (ra) {
        0, 1 => .eq,
        2 => compare_f128(a.as_f128(), b.as_f128()),
        3 => std.mem.order(u8, as_str(a), as_str(b)),
        4 => compare_docs(a.doc, b.doc),
        5 => compare_arrays(a.array, b.array),
        6 => compare_binary(a.binary, b.binary),
        7 => std.mem.order(u8, &a.object_id, &b.object_id),
        8 => std.math.order(@intFromBool(a.bool), @intFromBool(b.bool)),
        9 => std.math.order(a.datetime, b.datetime),
        10 => std.math.order(a.timestamp, b.timestamp),
        11 => blk: {
            const p = std.mem.order(u8, as_regex(a).pattern, as_regex(b).pattern);
            break :blk if (p != .eq) p else std.mem.order(u8, as_regex(a).options, as_regex(b).options);
        },
        12 => std.mem.order(u8, a.opaque_val.data, b.opaque_val.data),
        13 => std.mem.order(u8, &a.decimal128, &b.decimal128),
        14 => .eq,
        else => unreachable,
    };
}

fn as_str(v: Value) []const u8 {
    return switch (v) {
        .string => |s| s,
        .symbol => |s| s,
        .code => |c| c,
        else => unreachable,
    };
}

fn as_regex(v: Value) Regex {
    return switch (v) {
        .regex => |r| r,
        else => unreachable,
    };
}

fn compare_f128(a: f128, b: f128) std.math.Order {
    if (a < b) return .lt;
    if (a > b) return .gt;
    // Distinguish -0.0 from +0.0 like MongoDB does (they are equal); NaN
    // sorts greater than every number (MongoDB treats NaN as largest).
    if (std.math.isNan(a)) {
        if (std.math.isNan(b)) return .eq;
        return .gt;
    }
    if (std.math.isNan(b)) return .lt;
    return .eq;
}

fn compare_docs(a: []const Pair, b: []const Pair) std.math.Order {
    const n = @min(a.len, b.len);
    for (a[0..n], b[0..n]) |pa, pb| {
        // Keys tie-break equal values, so {a: 1} and {b: 1} are distinct.
        const ko = std.mem.order(u8, pa.key, pb.key);
        if (ko != .eq) return ko;
        const o = compare(pa.value, pb.value);
        if (o != .eq) return o;
    }
    return std.math.order(a.len, b.len);
}

fn compare_arrays(a: []const Value, b: []const Value) std.math.Order {
    const n = @min(a.len, b.len);
    for (a[0..n], b[0..n]) |va, vb| {
        const o = compare(va, vb);
        if (o != .eq) return o;
    }
    return std.math.order(a.len, b.len);
}

fn compare_binary(a: Binary, b: Binary) std.math.Order {
    const l = std.math.order(a.data.len, b.data.len);
    if (l != .eq) return l;
    const d = std.mem.order(u8, a.data, b.data);
    if (d != .eq) return d;
    return std.math.order(a.subtype, b.subtype);
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

const testing = std.testing;

test "document round-trip" {
    var out: std.ArrayListUnmanaged(u8) = .empty;
    defer out.deinit(testing.allocator);

    try write_doc(&.{
        .{ .key = "_id", .value = .{ .int32 = 7 } },
        .{ .key = "name", .value = .{ .string = "héllo" } },
        .{ .key = "pi", .value = .{ .double = 3.25 } },
        .{ .key = "ok", .value = .{ .bool = true } },
        .{ .key = "nul", .value = .null },
        .{ .key = "big", .value = .{ .int64 = 1 << 40 } },
        .{ .key = "when", .value = .{ .datetime = 1_700_000_000_000 } },
        .{ .key = "re", .value = .{ .regex = .{ .pattern = "^a", .options = "i" } } },
        .{ .key = "bin", .value = .{ .binary = .{ .subtype = 0x80, .data = &[_]u8{ 1, 2, 3 } } } },
        .{ .key = "arr", .value = .{ .array = &.{ .{ .int32 = 1 }, .{ .string = "x" } } } },
        .{ .key = "sub", .value = .{ .doc = &.{.{ .key = "a", .value = .{ .int32 = 1 } }} } },
        .{ .key = "ts", .value = .{ .timestamp = 42 } },
        .{ .key = "oid", .value = .{ .object_id = [_]u8{ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12 } } },
    }, testing.allocator, &out);

    var doc = try Document.parse(testing.allocator, out.items);
    defer doc.deinit();

    try testing.expectEqual(@as(i64, 7), doc.get("_id").?.int32);
    try testing.expectEqualStrings("héllo", doc.get("name").?.string);
    try testing.expectEqual(@as(f64, 3.25), doc.get("pi").?.double);
    try testing.expect(doc.get("ok").?.bool);
    try testing.expectEqual(@as(i64, 1 << 40), doc.get("big").?.int64);
    try testing.expectEqual(@as(i64, 1_700_000_000_000), doc.get("when").?.datetime);
    try testing.expectEqualStrings("^a", doc.get("re").?.regex.pattern);
    try testing.expectEqualStrings("i", doc.get("re").?.regex.options);
    try testing.expectEqualSlices(u8, &[_]u8{ 1, 2, 3 }, doc.get("bin").?.binary.data);
    try testing.expectEqual(@as(usize, 2), doc.get("arr").?.array.len);
    try testing.expectEqual(@as(i64, 1), doc.get("sub").?.doc[0].value.int32);
    try testing.expectEqual(@as(u64, 42), doc.get("ts").?.timestamp);
    try testing.expectEqualSlices(u8, &[_]u8{ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12 }, &doc.get("oid").?.object_id);
}

test "array and nested doc round-trip" {
    var out: std.ArrayListUnmanaged(u8) = .empty;
    defer out.deinit(testing.allocator);
    try write_doc(&.{
        .{ .key = "a", .value = .{ .array = &.{
            .{ .int32 = 10 },
            .{ .doc = &.{.{ .key = "deep", .value = .{ .string = "v" } }} },
        } } },
    }, testing.allocator, &out);

    var doc = try Document.parse(testing.allocator, out.items);
    defer doc.deinit();
    const a = doc.get("a").?.array;
    try testing.expectEqual(@as(i64, 10), a[0].int32);
    try testing.expectEqualStrings("v", a[1].doc[0].value.string);
}

test "reject truncated document" {
    try testing.expectError(error.InvalidBson, Document.parse(testing.allocator, &[_]u8{ 6, 0, 0, 0 }));
    try testing.expectError(error.InvalidBson, Document.parse(testing.allocator, &[_]u8{ 4, 0, 0, 0, 0 }));
    try testing.expectError(error.InvalidBson, Document.parse(testing.allocator, &[_]u8{ 9, 0, 0, 0, 0x01, 'a', 0 }));
}

test "compare: canonical order" {
    const min = Value{ .min_key = {} };
    const nul = Value.null;
    const i32a = Value{ .int32 = 5 };
    const i64a = Value{ .int64 = 5 };
    const dbl = Value{ .double = 4.9 };
    const str = Value{ .string = "a" };
    const obj = Value{ .doc = &.{} };
    const arr = Value{ .array = &.{} };
    const oid = Value{ .object_id = [_]u8{0} ** 12 };
    const btrue = Value{ .bool = true };
    const bfalse = Value{ .bool = false };
    const max = Value{ .max_key = {} };

    try testing.expectEqual(std.math.Order.lt, compare(min, nul));
    try testing.expectEqual(std.math.Order.lt, compare(nul, i32a));
    try testing.expectEqual(std.math.Order.eq, compare(i32a, i64a)); // numeric equality across widths
    try testing.expectEqual(std.math.Order.gt, compare(i32a, dbl)); // 5 > 4.9
    try testing.expectEqual(std.math.Order.lt, compare(i32a, str));
    try testing.expectEqual(std.math.Order.lt, compare(str, obj));
    try testing.expectEqual(std.math.Order.lt, compare(obj, arr));
    try testing.expectEqual(std.math.Order.lt, compare(arr, oid));
    try testing.expectEqual(std.math.Order.lt, compare(oid, bfalse));
    try testing.expectEqual(std.math.Order.lt, compare(bfalse, btrue));
    try testing.expectEqual(std.math.Order.gt, compare(max, nul));
}

test "compare: strings and docs" {
    try testing.expectEqual(std.math.Order.lt, compare(.{ .string = "a" }, .{ .string = "b" }));
    try testing.expectEqual(std.math.Order.eq, compare(.{ .string = "x" }, .{ .string = "x" }));
    try testing.expectEqual(std.math.Order.lt, compare(
        .{ .doc = &.{.{ .key = "a", .value = .{ .int32 = 1 } }} },
        .{ .doc = &.{.{ .key = "a", .value = .{ .int32 = 2 } }} },
    ));
    try testing.expectEqual(std.math.Order.lt, compare(
        .{ .array = &.{.{ .int32 = 1 }} },
        .{ .array = &.{ .{ .int32 = 1 }, .{ .int32 = 1 } } },
    ));
}

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();
    const io = threaded.io();

    var gen = ObjectIdGen.init(io);
    const a = gen.new(io);
    const b = gen.new(io);
    try testing.expect(!std.mem.eql(u8, &a, &b));
    // timestamp bytes match wall clock roughly
    try testing.expect(a[0] >= 0x66); // 2024+ in big-endian seconds
    // The 5-byte random prefix is untouched by the counter (spec layout:
    // 4s timestamp | 5B random | 3B counter).
    try testing.expectEqualSlices(u8, a[0..9], b[0..9]);
    try testing.expect(!std.mem.eql(u8, a[9..12], b[9..12]));
}

test "serialize_value deterministic for _id keys" {
    const gpa = testing.allocator;
    const v1 = try serialize_value(gpa, .{ .doc = &.{
        .{ .key = "a", .value = .{ .int32 = 1 } },
        .{ .key = "b", .value = .{ .string = "s" } },
    } });
    defer gpa.free(v1);
    const v2 = try serialize_value(gpa, .{ .doc = &.{
        .{ .key = "a", .value = .{ .int32 = 1 } },
        .{ .key = "b", .value = .{ .string = "s" } },
    } });
    defer gpa.free(v2);
    try testing.expectEqualSlices(u8, v1, v2);
}