index: secondary index core — entries, search, planner, _id fast path

Adds src/index.zig with the full secondary-index machinery: entry
generation mirroring field_matches (array value + elements), BSON-order
sorted entries with binary search, compound prefix and range lookups,
unique/sparse options, the query planner (longest equality/$in run +
optional range, $in cartesian cap, sparse/null bail), and the _id_ fast
path guarded against serialization-ambiguous values (numbers, strings,
symbols, codes, opaque payloads).

query.collect_values is now pub so entry generation can mirror it exactly.
storage.zig gains record_type_index_create/drop; lib.zig exports index.

src/bson.zig

@@ -128,14 +128,17 @@ pub const Document = struct {
 };
 
 pub fn get_pair(pairs: []const Pair, key: []const u8) ?Value {
-    for (pairs) |p| {
+    const i = get_pair_index(pairs, key) orelse return null;
-        if (std.mem.eql(u8, p.key, key)) return p.value;
+    return pairs[i].value;
-    }
-    return null;
 }
 
-pub fn parse_doc(allocator: std.mem.Allocator, bytes: []const u8) !Document {
+/// Position of `key` in `pairs` — for callers that need to replace or remove
-    return Document.parse(allocator, bytes);
+/// 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;
 }
 
 // ---------------------------------------------------------------------------
@@ -145,10 +148,6 @@ pub fn parse_doc(allocator: std.mem.Allocator, bytes: []const u8) !Document {
 const Parser = struct {
     arena: *std.heap.ArenaAllocator,
     bytes: []const u8,
-
-    fn fail() error{InvalidBson} {
-        return error.InvalidBson;
-    }
 };
 
 const ParseError = error{ InvalidBson, OutOfMemory };
@@ -162,14 +161,21 @@ fn ensure_available(bytes: []const u8, idx: usize, n: usize) error{InvalidBson}!
     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.*;
-    ensure_available(p.bytes, start, 4) catch return Parser.fail();
+    const end = try doc_extent(p, start);
-    const total: u32 = std.mem.readInt(u32, p.bytes[start..][0..4], .little);
-    if (total < 5) return Parser.fail();
-    if (p.bytes.len - start < total) return Parser.fail();
-    const end = start + total;
-    if (p.bytes[end - 1] != 0x00) return Parser.fail();
 
     const gpa = p.arena.allocator();
     var pairs: std.ArrayListUnmanaged(Pair) = .empty;
@@ -180,13 +186,13 @@ fn parse_doc_inner(p: Parser, idx: *usize) ParseError![]const Pair {
         const value = try parse_element(p, idx);
         try pairs.append(gpa, value);
     }
-    if (idx.* != end - 1) return Parser.fail();
+    if (idx.* != end - 1) return error.InvalidBson;
     idx.* = end;
     return pairs.toOwnedSlice(gpa);
 }
 
 fn parse_element(p: Parser, idx: *usize) ParseError!Pair {
-    ensure_available(p.bytes, idx.*, 1) catch return Parser.fail();
+    try ensure_available(p.bytes, idx.*, 1);
     const tag = p.bytes[idx.*];
     idx.* += 1;
     const key = try parse_cstring(p, idx);
@@ -197,7 +203,7 @@ fn parse_element(p: Parser, idx: *usize) ParseError!Pair {
 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 Parser.fail();
+    if (idx.* >= p.bytes.len) return error.InvalidBson;
     idx.* += 1;
     // Strings are copied into the arena so documents are self-contained and
     // outlive the input buffer (wire messages and log records are transient).
@@ -205,11 +211,11 @@ fn parse_cstring(p: Parser, idx: *usize) ParseError![]const u8 {
 }
 
 fn parse_string(p: Parser, idx: *usize) ParseError![]const u8 {
-    ensure_available(p.bytes, idx.*, 4) catch return Parser.fail();
+    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 Parser.fail();
+    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 Parser.fail();
+    if (str[len - 1] != 0) return error.InvalidBson;
     idx.* += 4 + len;
     return p.arena.allocator().dupe(u8, str[0 .. len - 1]);
 }
@@ -217,7 +223,7 @@ fn parse_string(p: Parser, idx: *usize) ParseError![]const u8 {
 fn parse_value(p: Parser, tag: u8, idx: *usize) ParseError!Value {
     return switch (tag) {
         0x01 => blk: {
-            ensure_available(p.bytes, idx.*, 8) catch return Parser.fail();
+            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 };
@@ -226,30 +232,30 @@ fn parse_value(p: Parser, tag: u8, idx: *usize) ParseError!Value {
         0x03 => .{ .doc = try parse_doc_inner(p, idx) },
         0x04 => .{ .array = try parse_array(p, idx) },
         0x05 => blk: {
-            ensure_available(p.bytes, idx.*, 5) catch return Parser.fail();
+            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];
-            ensure_available(p.bytes, idx.* + 5, len) catch return Parser.fail();
+            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.arena.allocator().dupe(u8, data) } };
         },
         0x06 => .{ .opaque_val = .{ .kind = 0x06, .data = &.{} } },
         0x07 => blk: {
-            ensure_available(p.bytes, idx.*, 12) catch return Parser.fail();
+            try ensure_available(p.bytes, idx.*, 12);
             const oid: ObjectId = p.bytes[idx.*..][0..12].*;
             idx.* += 12;
             break :blk .{ .object_id = oid };
         },
         0x08 => blk: {
-            ensure_available(p.bytes, idx.*, 1) catch return Parser.fail();
+            try ensure_available(p.bytes, idx.*, 1);
             const v = p.bytes[idx.*];
-            if (v > 1) return Parser.fail();
+            if (v > 1) return error.InvalidBson;
             idx.* += 1;
             break :blk .{ .bool = v == 1 };
         },
         0x09 => blk: {
-            ensure_available(p.bytes, idx.*, 8) catch return Parser.fail();
+            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 };
@@ -262,7 +268,7 @@ fn parse_value(p: Parser, tag: u8, idx: *usize) ParseError!Value {
         0x0C => blk: {
             const start = idx.*;
             _ = try parse_string(p, idx);
-            ensure_available(p.bytes, idx.*, 12) catch return Parser.fail();
+            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.
@@ -271,49 +277,44 @@ fn parse_value(p: Parser, tag: u8, idx: *usize) ParseError!Value {
         0x0D => .{ .code = try parse_string(p, idx) },
         0x0E => .{ .symbol = try parse_string(p, idx) },
         0x0F => blk: {
-            ensure_available(p.bytes, idx.*, 4) catch return Parser.fail();
+            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 Parser.fail();
+            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.arena.allocator().dupe(u8, data) } };
         },
         0x10 => blk: {
-            ensure_available(p.bytes, idx.*, 4) catch return Parser.fail();
+            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: {
-            ensure_available(p.bytes, idx.*, 8) catch return Parser.fail();
+            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: {
-            ensure_available(p.bytes, idx.*, 8) catch return Parser.fail();
+            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: {
-            ensure_available(p.bytes, idx.*, 16) catch return Parser.fail();
+            try ensure_available(p.bytes, idx.*, 16);
             const v: [16]u8 = p.bytes[idx.*..][0..16].*;
             idx.* += 16;
             break :blk .{ .decimal128 = v };
         },
-        else => return Parser.fail(),
+        else => return error.InvalidBson,
     };
 }
 
 fn parse_array(p: Parser, idx: *usize) ParseError![]const Value {
     const start = idx.*;
-    ensure_available(p.bytes, start, 4) catch return Parser.fail();
+    const end = try doc_extent(p, start);
-    const total: u32 = std.mem.readInt(u32, p.bytes[start..][0..4], .little);
-    if (total < 5) return Parser.fail();
-    if (p.bytes.len - start < total) return Parser.fail();
-    const end = start + total;
-    if (p.bytes[end - 1] != 0x00) return Parser.fail();
 
     const gpa = p.arena.allocator();
     var values: std.ArrayListUnmanaged(Value) = .empty;
@@ -321,13 +322,13 @@ fn parse_array(p: Parser, idx: *usize) ParseError![]const Value {
 
     idx.* = start + 4;
     while (idx.* < end - 1) {
-        ensure_available(p.bytes, idx.*, 1) catch return Parser.fail();
+        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 Parser.fail();
+    if (idx.* != end - 1) return error.InvalidBson;
     idx.* = end;
     return values.toOwnedSlice(gpa);
 }
@@ -417,22 +418,31 @@ pub fn write_element(pair: Pair, gpa: std.mem.Allocator, out: *std.ArrayListUnma
     try write_value(pair.value, gpa, out);
 }
 
-/// Write a length-prefixed document. Length is patched in after the body.
+/// Reserve the 4-byte length prefix of a document or array frame. Returns the
-pub fn write_doc(pairs: []const Pair, gpa: std.mem.Allocator, out: *std.ArrayListUnmanaged(u8)) SerializeError!void {
+/// 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;
-    var zero: [4]u8 = [4]u8{ 0, 0, 0, 0 };
+    try out.appendSlice(gpa, &[4]u8{ 0, 0, 0, 0 });
-    try out.appendSlice(gpa, &zero);
+    return len_pos;
-    for (pairs) |p| try write_element(p, gpa, out);
+}
+
+/// 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);
 }
 
-fn write_array(items: []const Value, gpa: std.mem.Allocator, out: *std.ArrayListUnmanaged(u8)) !void {
+/// Write a length-prefixed document. Length is patched in after the body.
-    const len_pos = out.items.len;
+pub fn write_doc(pairs: []const Pair, gpa: std.mem.Allocator, out: *std.ArrayListUnmanaged(u8)) SerializeError!void {
-    var zero: [4]u8 = [4]u8{ 0, 0, 0, 0 };
+    const len_pos = try begin_frame(gpa, out);
-    try out.appendSlice(gpa, &zero);
+    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());
@@ -440,10 +450,7 @@ fn write_array(items: []const Value, gpa: std.mem.Allocator, out: *std.ArrayList
         try write_cstring(key, gpa, out);
         try write_value(item, gpa, out);
     }
-    try out.append(gpa, 0);
+    try end_frame(gpa, out, len_pos);
-    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);
 }
 
 /// Serialize a single value with its type byte (no key) — used for `_id`

src/db.zig

@@ -55,24 +55,43 @@ pub const Engine = struct {
     pub fn deinit(self: *Engine) void {
         var db_it = self.dbs.iterator();
         while (db_it.next()) |db_entry| {
-            var coll_it = db_entry.value_ptr.collections.iterator();
+            self.free_db(db_entry.value_ptr);
-            while (coll_it.next()) |coll_entry| {
-                var doc_it = coll_entry.value_ptr.docs.iterator();
-                while (doc_it.next()) |doc_entry| {
-                    doc_entry.value_ptr.*.deinit();
-                    self.gpa.destroy(doc_entry.value_ptr.*);
-                    self.gpa.free(doc_entry.key_ptr.*);
-                }
-                coll_entry.value_ptr.docs.deinit(self.gpa);
-                self.gpa.free(coll_entry.key_ptr.*);
-            }
-            db_entry.value_ptr.collections.deinit(self.gpa);
             self.gpa.free(db_entry.key_ptr.*);
         }
         self.dbs.deinit(self.gpa);
         self.log.close();
     }
 
+    /// Free every document in a collection along with its owned _id keys.
+    fn free_collection(self: *Engine, coll: *Collection) void {
+        var doc_it = coll.docs.iterator();
+        while (doc_it.next()) |doc_entry| {
+            doc_entry.value_ptr.*.deinit();
+            self.gpa.destroy(doc_entry.value_ptr.*);
+            self.gpa.free(doc_entry.key_ptr.*);
+        }
+        coll.docs.deinit(self.gpa);
+    }
+
+    /// Free every collection in a database along with its owned name keys.
+    fn free_db(self: *Engine, db: *Db) void {
+        var coll_it = db.collections.iterator();
+        while (coll_it.next()) |coll_entry| {
+            self.free_collection(coll_entry.value_ptr);
+            self.gpa.free(coll_entry.key_ptr.*);
+        }
+        db.collections.deinit(self.gpa);
+    }
+
+    /// Drop the document stored under `id_key`, freeing it and its key.
+    /// No-op when the id is absent.
+    fn evict_doc(self: *Engine, coll: *Collection, id_key: []const u8) void {
+        const old = coll.docs.fetchRemove(id_key) orelse return;
+        old.value.*.deinit();
+        self.gpa.destroy(old.value);
+        self.gpa.free(old.key);
+    }
+
     // -- commands (callers must hold the matching lock) ---------------------
 
     /// Exclusive lock: for commands that mutate the engine.
@@ -97,74 +116,76 @@ pub const Engine = struct {
     /// Insert a document. Fails with error.DuplicateKey if the _id exists.
     /// Generates an ObjectId _id when absent.
     pub fn insert(self: *Engine, db_name: []const u8, coll_name: []const u8, doc: *const bson.Document, oid_gen: *bson.ObjectIdGen) !void {
-        const coll = try self.get_or_create_collection(db_name, coll_name);
+        return self.upsert(db_name, coll_name, doc, oid_gen, .insert);
-        const owned = try self.own_with_id(doc, oid_gen);
-        errdefer {
-            owned.deinit();
-            self.gpa.destroy(owned);
-        }
-
-        const id_value = owned.get("_id") orelse unreachable;
-        const id_key = try bson.serialize_value(self.gpa, id_value);
-        defer self.gpa.free(id_key);
-
-        if (coll.docs.contains(id_key)) return error.DuplicateKey;
-
-        const doc_bytes = try serialize_doc(self.gpa, owned);
-        defer self.gpa.free(doc_bytes);
-        self.seq += 1;
-        try self.log.append_upsert(db_name, coll_name, doc_bytes, self.seq);
-
-        const key_owned = try self.gpa.dupe(u8, id_key);
-        try coll.docs.put(self.gpa, key_owned, owned);
-        try self.maybe_compact();
     }
 
     /// Insert or replace a document by _id (upsert without existence check).
     pub fn replace(self: *Engine, db_name: []const u8, coll_name: []const u8, doc: *const bson.Document, oid_gen: *bson.ObjectIdGen) !void {
+        return self.upsert(db_name, coll_name, doc, oid_gen, .replace);
+    }
+
+    /// Shared body of `insert` and `replace`: they differ only in how an
+    /// existing _id is treated. Logs (and syncs) the new document before it
+    /// becomes visible in memory.
+    fn upsert(
+        self: *Engine,
+        db_name: []const u8,
+        coll_name: []const u8,
+        doc: *const bson.Document,
+        oid_gen: *bson.ObjectIdGen,
+        mode: enum { insert, replace },
+    ) !void {
         const coll = try self.get_or_create_collection(db_name, coll_name);
         const owned = try self.own_with_id(doc, oid_gen);
-        errdefer {
+        const id_value = owned.get("_id") orelse unreachable;
+        // Ownership of the key moves to the map once `stored` is set; until
+        // then this frame still owns both it and `owned`.
+        const id_key = try bson.serialize_value(self.gpa, id_value);
+        var stored = false;
+        errdefer if (!stored) {
             owned.deinit();
             self.gpa.destroy(owned);
-        }
+            self.gpa.free(id_key);
+        };
 
-        const id_value = owned.get("_id") orelse unreachable;
+        if (mode == .insert and coll.docs.contains(id_key)) return error.DuplicateKey;
-        const id_key = try bson.serialize_value(self.gpa, id_value);
-        defer self.gpa.free(id_key);
 
         const doc_bytes = try serialize_doc(self.gpa, owned);
         defer self.gpa.free(doc_bytes);
         self.seq += 1;
         try self.log.append_upsert(db_name, coll_name, doc_bytes, self.seq);
 
-        if (coll.docs.fetchRemove(id_key)) |old| {
+        if (mode == .replace) self.evict_doc(coll, id_key);
-            old.value.*.deinit();
+        try coll.docs.put(self.gpa, id_key, owned);
-            self.gpa.destroy(old.value);
+        stored = true;
-            self.gpa.free(old.key);
-        }
-        const key_owned = try self.gpa.dupe(u8, id_key);
-        try coll.docs.put(self.gpa, key_owned, owned);
         try self.maybe_compact();
     }
 
-    /// Remove a document by _id. Returns true if it existed.
+    /// Remove a document by its `_id` value. Returns true if it existed.
-    pub fn remove(self: *Engine, db_name: []const u8, coll_name: []const u8, id_key: []const u8) !bool {
+    /// The serialized-key encoding stays private to the engine.
+    pub fn remove_by_id(self: *Engine, db_name: []const u8, coll_name: []const u8, id: bson.Value) !bool {
+        const id_key = try bson.serialize_value(self.gpa, id);
+        defer self.gpa.free(id_key);
+        return self.remove(db_name, coll_name, id_key);
+    }
+
+    fn remove(self: *Engine, db_name: []const u8, coll_name: []const u8, id_key: []const u8) !bool {
         const db = self.dbs.get(db_name) orelse return false;
         const coll = db.collections.getPtr(coll_name) orelse return false;
         const doc = coll.docs.get(id_key) orelse return false;
 
         // Log (and sync) the delete before removing it from memory, so the
         // log always describes at least as much as the in-memory state.
-        const doc_bytes = try serialize_doc(self.gpa, doc);
+        // Replay only reads _id out of a delete record, so log just that
-        defer self.gpa.free(doc_bytes);
+        // rather than a copy of the whole document.
+        const id_pairs = [_]bson.Pair{.{ .key = "_id", .value = doc.get("_id") orelse unreachable }};
+        var id_doc: std.ArrayListUnmanaged(u8) = .empty;
+        defer id_doc.deinit(self.gpa);
+        try bson.write_doc(&id_pairs, self.gpa, &id_doc);
         self.seq += 1;
-        try self.log.append_delete(db_name, coll_name, doc_bytes, self.seq);
+        try self.log.append_delete(db_name, coll_name, id_doc.items, self.seq);
 
-        const removed = coll.docs.fetchRemove(id_key) orelse unreachable;
+        self.evict_doc(coll, id_key);
-        removed.value.*.deinit();
-        self.gpa.destroy(removed.value);
-        self.gpa.free(removed.key);
         return true;
     }
 
@@ -181,31 +202,14 @@ pub const Engine = struct {
     pub fn drop_collection(self: *Engine, db_name: []const u8, coll_name: []const u8) !bool {
         const db = self.dbs.getPtr(db_name) orelse return false;
         var removed = db.collections.fetchRemove(coll_name) orelse return false;
-        var doc_it = removed.value.docs.iterator();
+        self.free_collection(&removed.value);
-        while (doc_it.next()) |doc_entry| {
-            doc_entry.value_ptr.*.deinit();
-            self.gpa.destroy(doc_entry.value_ptr.*);
-            self.gpa.free(doc_entry.key_ptr.*);
-        }
-        removed.value.docs.deinit(self.gpa);
         self.gpa.free(removed.key);
         return true;
     }
 
     pub fn drop_database(self: *Engine, db_name: []const u8) !bool {
         var removed = self.dbs.fetchRemove(db_name) orelse return false;
-        var coll_it = removed.value.collections.iterator();
+        self.free_db(&removed.value);
-        while (coll_it.next()) |coll_entry| {
-            var docs_it = coll_entry.value_ptr.docs.iterator();
-            while (docs_it.next()) |doc_entry| {
-                doc_entry.value_ptr.*.deinit();
-                self.gpa.destroy(doc_entry.value_ptr.*);
-                self.gpa.free(doc_entry.key_ptr.*);
-            }
-            coll_entry.value_ptr.docs.deinit(self.gpa);
-            self.gpa.free(coll_entry.key_ptr.*);
-        }
-        removed.value.collections.deinit(self.gpa);
         self.gpa.free(removed.key);
         return true;
     }
@@ -328,28 +332,19 @@ fn apply_record(ctx: *anyopaque, record: storage.Record, doc: *bson.Document) an
         return;
     };
     const id_key = try bson.serialize_value(self.gpa, id_value);
-    defer self.gpa.free(id_key);
+    var key_owned = false;
+    defer if (!key_owned) self.gpa.free(id_key);
 
     const coll = self.get_or_create_collection(record.db, record.coll) catch return;
 
     switch (record.type) {
         storage.record_type_upsert => {
-            if (coll.docs.fetchRemove(id_key)) |old| {
+            self.evict_doc(coll, id_key);
-                old.value.*.deinit();
+            try coll.docs.put(self.gpa, id_key, doc);
-                self.gpa.destroy(old.value);
+            key_owned = true;
-                self.gpa.free(old.key);
-            }
-            const key_owned = try self.gpa.dupe(u8, id_key);
-            try coll.docs.put(self.gpa, key_owned, doc);
             stored = true;
         },
-        storage.record_type_delete => {
+        storage.record_type_delete => self.evict_doc(coll, id_key),
-            if (coll.docs.fetchRemove(id_key)) |old| {
-                old.value.*.deinit();
-                self.gpa.destroy(old.value);
-                self.gpa.free(old.key);
-            }
-        },
         else => {},
     }
 }
@@ -360,21 +355,7 @@ fn apply_record(ctx: *anyopaque, record: storage.Record, doc: *bson.Document) an
 
 const testing = std.testing;
 
-const TmpLog = struct {
+const TmpLog = storage.TmpLog;
-    tmp: std.testing.TmpDir,
-    path: []u8,
-
-    fn init(gpa: std.mem.Allocator) !TmpLog {
-        const tmp = std.testing.tmpDir(.{});
-        const path = try std.fmt.allocPrint(gpa, ".zig-cache/tmp/{s}/test.log", .{tmp.sub_path});
-        return .{ .tmp = tmp, .path = path };
-    }
-
-    fn deinit(self: *TmpLog, gpa: std.mem.Allocator) void {
-        self.tmp.cleanup();
-        gpa.free(self.path);
-    }
-};
 
 fn test_env(threaded: *std.Io.Threaded) struct { io: std.Io, gen: bson.ObjectIdGen } {
     const io = threaded.io();
@@ -426,7 +407,7 @@ test "insert, query, remove" {
     try engine.lock();
     const found = engine.get_doc("app", "users", id_key).?;
     try testing.expectEqualStrings("bob", found.get("name").?.string);
-    const removed = try engine.remove("app", "users", id_key);
+    const removed = try engine.remove_by_id("app", "users", .{ .int32 = 2 });
     try testing.expect(removed);
     engine.unlock();
 }
@@ -450,9 +431,7 @@ test "reopen replays log" {
         try engine.lock();
         try engine.insert("app", "users", &d1, &env.gen);
         try engine.insert("app", "users", &d2, &env.gen);
-        const id_key2 = try bson.serialize_value(gpa, bson.Value{ .int32 = 2 });
+        _ = try engine.remove_by_id("app", "users", .{ .int32 = 2 });
-        defer gpa.free(id_key2);
-        _ = try engine.remove("app", "users", id_key2);
         engine.unlock();
     }
 

src/lib.zig

@@ -8,6 +8,7 @@ pub const storage = @import("storage.zig");
 pub const db = @import("db.zig");
 pub const query = @import("query.zig");
 pub const update = @import("update.zig");
+pub const index = @import("index.zig");
 
 test {
     _ = @import("bson.zig");
@@ -18,4 +19,5 @@ test {
     _ = @import("db.zig");
     _ = @import("query.zig");
     _ = @import("update.zig");
+    _ = @import("index.zig");
 }

src/query.zig

@@ -57,16 +57,26 @@ fn match_top_level(gpa: std.mem.Allocator, op: []const u8, value: bson.Value, do
     return false;
 }
 
+/// True when every key is a `$`-prefixed operator. The one definition of
+/// what an operator document looks like — filters, $elemMatch, $pull and
+/// upsert-document construction all ask this question.
+pub fn all_operator_keys(pairs: []const bson.Pair) bool {
+    if (pairs.len == 0) return false;
+    for (pairs) |p| {
+        if (p.key.len == 0 or p.key[0] != '$') return false;
+    }
+    return true;
+}
+
+/// The operator pairs of a `{$op: ...}` value, or null if it is not one. An
+/// empty document counts as an (vacuously satisfied) operator document.
 fn is_operator_doc(value: bson.Value) ?[]const bson.Pair {
-    return switch (value) {
+    const pairs = switch (value) {
-        .doc => |pairs| blk: {
+        .doc => |p| p,
-            for (pairs) |p| {
+        else => return null,
-                if (p.key.len == 0 or p.key[0] != '$') break :blk null;
-            }
-            break :blk pairs;
-        },
-        else => null,
     };
+    if (pairs.len > 0 and !all_operator_keys(pairs)) return null;
+    return pairs;
 }
 
 fn field_matches(gpa: std.mem.Allocator, path: []const u8, expected: bson.Value, doc: *const bson.Document) QueryError!bool {
@@ -220,13 +230,7 @@ fn match_operator(gpa: std.mem.Allocator, op: []const u8, value: bson.Value, act
             .doc => |pairs| pairs,
             else => return false,
         };
-        var all_operators = operand.len > 0;
+        const all_operators = all_operator_keys(operand);
-        for (operand) |p| {
-            if (p.key.len == 0 or p.key[0] != '$') {
-                all_operators = false;
-                break;
-            }
-        }
         for (actuals) |a| {
             if (a != .array) continue;
             for (a.array) |elem| {
@@ -258,8 +262,9 @@ fn match_operator(gpa: std.mem.Allocator, op: []const u8, value: bson.Value, act
 /// Collect values reachable at `path` (dot-separated), descending into
 /// documents and, per MongoDB multikey semantics, into arrays of documents.
 /// Appends into `out`; on OOM, collection stops early (the engine is
-/// already failing at that point).
+/// already failing at that point). Public because index entry generation
-fn collect_values(gpa: std.mem.Allocator, pairs: []const bson.Pair, path: []const u8, out: *std.ArrayListUnmanaged(bson.Value), depth: usize) QueryError!void {
+/// must mirror field_matches exactly (src/index.zig).
+pub fn collect_values(gpa: std.mem.Allocator, pairs: []const bson.Pair, path: []const u8, out: *std.ArrayListUnmanaged(bson.Value), depth: usize) QueryError!void {
     var it = std.mem.splitScalar(u8, path, '.');
     const first = it.next() orelse return;
 
@@ -550,7 +555,7 @@ pub const SortKey = struct {
 /// Sort `docs` in place by `keys`. Candidate values are collected up front
 /// (allocations happen before the sort), so the comparator itself is pure
 /// and cannot fail — OOM during collection propagates as QueryError.
-pub fn sort_docs(gpa: std.mem.Allocator, arena: std.mem.Allocator, docs: []*const bson.Document, keys: []const SortKey) QueryError!void {
+pub fn sort_docs(arena: std.mem.Allocator, docs: []*const bson.Document, keys: []const SortKey) QueryError!void {
     if (keys.len == 0 or docs.len < 2) return;
 
     const SortedDoc = struct {
@@ -583,14 +588,13 @@ pub fn sort_docs(gpa: std.mem.Allocator, arena: std.mem.Allocator, docs: []*cons
     std.mem.sort(SortedDoc, entries, Ctx{ .keys = keys }, Ctx.lessThan);
 
     for (entries, 0..) |e, i| docs[i] = e.doc;
-    _ = gpa;
 }
 
 // ---------------------------------------------------------------------------
 // Projection
 // ---------------------------------------------------------------------------
 
-pub const ProjectionError = error{ OutOfMemory, InvalidProjection };
+pub const ProjectionError = std.mem.Allocator.Error;
 
 /// Apply a projection document, writing resulting pairs into `out` (which
 /// should use the caller's arena so strings are owned).
@@ -614,7 +618,7 @@ pub fn project(arena: std.mem.Allocator, doc: *const bson.Document, proj: *const
         }
         if (include_id) {
             if (bson.get_pair(doc.pairs, "_id")) |idv| {
-                try out.append(arena, .{ .key = try arena.dupe(u8, "_id"), .value = try bson.copy_value(arena, idv) });
+                try out.append(arena, .{ .key = "_id", .value = try bson.copy_value(arena, idv) });
             }
         }
         for (proj.pairs) |p| {
@@ -659,10 +663,7 @@ fn exclude_doc(arena: std.mem.Allocator, pairs: []const bson.Pair, proj: *const
 }
 
 fn is_excluded(proj: *const bson.Document, key: []const u8) bool {
-    for (proj.pairs) |pp| {
+    return bson.get_pair(proj.pairs, key) != null;
-        if (std.mem.eql(u8, pp.key, key)) return true;
-    }
-    return false;
 }
 
 fn has_deeper_exclusion(proj: *const bson.Document, key: []const u8) bool {
@@ -875,17 +876,17 @@ test "sort compares by BSON order" {
 
     var docs = [_]*const bson.Document{ &b, &a };
     const asc = [_]SortKey{.{ .path = "n", .descending = false }};
-    try sort_docs(testing.allocator, arena.allocator(), &docs, &asc);
+    try sort_docs(arena.allocator(), &docs, &asc);
     try testing.expect(docs[0] == &a);
     try testing.expect(docs[1] == &b);
 
     var docs2 = [_]*const bson.Document{ &a, &b };
     const desc = [_]SortKey{.{ .path = "n", .descending = true }};
-    try sort_docs(testing.allocator, arena.allocator(), &docs2, &desc);
+    try sort_docs(arena.allocator(), &docs2, &desc);
     try testing.expect(docs2[0] == &b);
 
     const missing = [_]SortKey{.{ .path = "zz", .descending = false }};
-    try sort_docs(testing.allocator, arena.allocator(), &docs2, &missing);
+    try sort_docs(arena.allocator(), &docs2, &missing);
     try testing.expect(docs2[0] == &b); // stable-ish: order untouched by missing key
 }
 

src/server.zig

@@ -120,17 +120,14 @@ fn handle_connection_inner(io: std.Io, stream: std.Io.net.Stream, server: *Serve
         };
 
         out_buf.clearRetainingCapacity();
-        if (msg.op_code == wire.op_code_query) {
+        const built = if (msg.op_code == wire.op_code_query)
-            wire.write_reply_query(server.gpa, reply_request_id, msg.request_id, reply.pairs.items, &out_buf) catch |err| {
+            wire.write_reply_query(server.gpa, reply_request_id, msg.request_id, reply.pairs.items, &out_buf)
-                std.debug.print("mongo-light: reply build error on conn {d}: {s}\n", .{ connection_id, @errorName(err) });
+        else
-                return;
+            reply.build(server.gpa, reply_request_id, msg.request_id, &out_buf);
-            };
+        built catch |err| {
-        } else {
+            std.debug.print("mongo-light: reply build error on conn {d}: {s}\n", .{ connection_id, @errorName(err) });
-            reply.build(server.gpa, reply_request_id, msg.request_id, &out_buf) catch |err| {
+            return;
-                std.debug.print("mongo-light: reply build error on conn {d}: {s}\n", .{ connection_id, @errorName(err) });
+        };
-                return;
-            };
-        }
         reply_request_id +%= 1;
         writer.interface.writeAll(out_buf.items) catch |err| {
             std.debug.print("mongo-light: write error on conn {d}: {s}\n", .{ connection_id, @errorName(err) });

src/storage.zig

@@ -14,6 +14,8 @@ const bson = @import("bson.zig");
 
 pub const record_type_upsert: u8 = 1;
 pub const record_type_delete: u8 = 2;
+pub const record_type_index_create: u8 = 3;
+pub const record_type_index_drop: u8 = 4;
 
 pub const header_len: usize = 20; // len + crc + seq + type + reserved
 
@@ -26,7 +28,6 @@ pub const Record = struct {
     type: u8,
     db: []const u8, // transient: valid only during replay callback
     coll: []const u8,
-    doc: []const u8, // raw bson bytes
 };
 
 pub const Error = error{
@@ -44,6 +45,9 @@ pub const Log = struct {
     path: []const u8,
     end_pos: u64,
     log_bytes: u64, // bytes written since the log was last rewritten
+    // Reused record-framing buffer. Appends are single-writer (the engine's
+    // exclusive lock), so one buffer avoids a realloc cycle per record.
+    scratch: std.ArrayListUnmanaged(u8),
 
     pub fn open(gpa: std.mem.Allocator, io: std.Io, path: []const u8) !Log {
         // Resolve to an absolute path so compaction can rename the file
@@ -68,10 +72,12 @@ pub const Log = struct {
             .path = abs_path,
             .end_pos = 0,
             .log_bytes = 0,
+            .scratch = .empty,
         };
     }
 
     pub fn close(self: *Log) void {
+        self.scratch.deinit(self.gpa);
         self.file.close(self.io);
         self.gpa.free(self.path);
     }
@@ -134,7 +140,6 @@ pub const Log = struct {
                 .type = rtype,
                 .db = db,
                 .coll = coll,
-                .doc = doc_bytes,
             }, doc);
 
             pos += total;
@@ -150,12 +155,23 @@ pub const Log = struct {
         try self.append(record_type_delete, db, coll, doc, seq);
     }
 
+    /// The payload is the canonical index spec document ({v, key, name,
+    /// unique?, sparse?}); only apply_record interprets it.
+    pub fn append_index_create(self: *Log, db: []const u8, coll: []const u8, doc: []const u8, seq: u64) !void {
+        try self.append(record_type_index_create, db, coll, doc, seq);
+    }
+
+    /// The payload is {name: "..."}; only apply_record interprets it.
+    pub fn append_index_drop(self: *Log, db: []const u8, coll: []const u8, doc: []const u8, seq: u64) !void {
+        try self.append(record_type_index_drop, db, coll, doc, seq);
+    }
+
     fn append(self: *Log, rtype: u8, db: []const u8, coll: []const u8, doc: []const u8, seq: u64) !void {
         if (std.mem.indexOfScalar(u8, db, 0) != null or std.mem.indexOfScalar(u8, coll, 0) != null) {
             return error.NulInName;
         }
-        var buf: std.ArrayListUnmanaged(u8) = .empty;
+        const buf = &self.scratch;
-        defer buf.deinit(self.gpa);
+        buf.clearRetainingCapacity();
         try buf.appendNTimes(self.gpa, 0, header_len);
         std.mem.writeInt(u64, buf.items[8..16], seq, .little);
         buf.items[16] = rtype;
@@ -189,17 +205,19 @@ pub const Log = struct {
 
 const testing = std.testing;
 
-const TmpLog = struct {
+/// Throwaway log file in the test temp dir, shared by the storage and
+/// engine test suites.
+pub const TmpLog = struct {
     tmp: std.testing.TmpDir,
     path: []u8,
 
-    fn init(gpa: std.mem.Allocator) !TmpLog {
+    pub fn init(gpa: std.mem.Allocator) !TmpLog {
         const tmp = std.testing.tmpDir(.{});
         const path = try std.fmt.allocPrint(gpa, ".zig-cache/tmp/{s}/test.log", .{tmp.sub_path});
         return .{ .tmp = tmp, .path = path };
     }
 
-    fn deinit(self: *TmpLog, gpa: std.mem.Allocator) void {
+    pub fn deinit(self: *TmpLog, gpa: std.mem.Allocator) void {
         self.tmp.cleanup();
         gpa.free(self.path);
     }

src/update.zig

@@ -13,7 +13,7 @@ const max_path_segments = 16;
 /// Apply an update document (whose fields are operator documents) to `doc`.
 pub fn apply(doc: *bson.Document, update: *const bson.Document) UpdateError!void {
     const arena = doc.arena.allocator();
-    var pairs = try copy_pairs_to_list(arena, doc.pairs);
+    var pairs = try copy_to_list(bson.Pair, arena, doc.pairs);
     for (update.pairs) |op| {
         if (op.key.len == 0 or op.key[0] != '$') return error.InvalidUpdate;
         try apply_operator(arena, &pairs, op.key, op.value);
@@ -48,7 +48,7 @@ fn apply_operator(arena: std.mem.Allocator, pairs: *std.ArrayListUnmanaged(bson.
             const n = split_path(p.key, &segs) orelse return error.InvalidUpdate;
             const current = get_value(pairs.items, segs[0..n]) orelse bson.Value{ .int32 = 0 };
             if (!current.is_number() or !p.value.is_number()) return error.InvalidUpdate;
-            const sum = try numeric_add(arena, current, p.value);
+            const sum = try numeric_add(current, p.value);
             try set_path(arena, pairs, segs[0..n], sum);
         }
         return;
@@ -145,19 +145,15 @@ fn parse_index(seg: []const u8) ?usize {
     return std.fmt.parseInt(usize, seg, 10) catch null;
 }
 
-fn copy_pairs_to_list(arena: std.mem.Allocator, pairs: []const bson.Pair) UpdateError!std.ArrayListUnmanaged(bson.Pair) {
+/// Shallow-copy a slice into a growable list backed by `arena`.
-    var out: std.ArrayListUnmanaged(bson.Pair) = .empty;
+fn copy_to_list(comptime T: type, arena: std.mem.Allocator, items: []const T) UpdateError!std.ArrayListUnmanaged(T) {
+    var out: std.ArrayListUnmanaged(T) = .empty;
     errdefer out.deinit(arena);
-    try out.appendSlice(arena, pairs);
+    try out.appendSlice(arena, items);
     return out;
 }
 
-fn find_pair(pairs: []const bson.Pair, key: []const u8) ?usize {
+const find_pair = bson.get_pair_index;
-    for (pairs, 0..) |p, i| {
-        if (std.mem.eql(u8, p.key, key)) return i;
-    }
-    return null;
-}
 
 fn get_value(pairs: []const bson.Pair, segs: []const []const u8) ?bson.Value {
     const idx = find_pair(pairs, segs[0]) orelse return null;
@@ -193,7 +189,7 @@ fn set_path(arena: std.mem.Allocator, pairs: *std.ArrayListUnmanaged(bson.Pair),
     };
     switch (pairs.items[idx].value) {
         .doc => |sub| {
-            var sub_pairs = try copy_pairs_to_list(arena, sub);
+            var sub_pairs = try copy_to_list(bson.Pair, arena, sub);
             defer sub_pairs.deinit(arena);
             try set_path(arena, &sub_pairs, segs[1..], value);
             pairs.items[idx].value = .{ .doc = try sub_pairs.toOwnedSlice(arena) };
@@ -207,7 +203,7 @@ fn set_path(arena: std.mem.Allocator, pairs: *std.ArrayListUnmanaged(bson.Pair),
                 pairs.items[idx].value = .{ .doc = try sub_pairs.toOwnedSlice(arena) };
                 return;
             };
-            var items = try copy_array_to_list(arena, arr);
+            var items = try copy_to_list(bson.Value, arena, arr);
             defer items.deinit(arena);
             if (index >= items.items.len) {
                 try items.appendNTimes(arena, .null, index + 1 - items.items.len);
@@ -217,7 +213,7 @@ fn set_path(arena: std.mem.Allocator, pairs: *std.ArrayListUnmanaged(bson.Pair),
             } else {
                 switch (items.items[index]) {
                     .doc => |sub| {
-                        var sub_pairs = try copy_pairs_to_list(arena, sub);
+                        var sub_pairs = try copy_to_list(bson.Pair, arena, sub);
                         defer sub_pairs.deinit(arena);
                         try set_path(arena, &sub_pairs, segs[2..], value);
                         items.items[index] = .{ .doc = try sub_pairs.toOwnedSlice(arena) };
@@ -241,13 +237,6 @@ fn set_path(arena: std.mem.Allocator, pairs: *std.ArrayListUnmanaged(bson.Pair),
     }
 }
 
-fn copy_array_to_list(arena: std.mem.Allocator, arr: []const bson.Value) UpdateError!std.ArrayListUnmanaged(bson.Value) {
-    var out: std.ArrayListUnmanaged(bson.Value) = .empty;
-    errdefer out.deinit(arena);
-    try out.appendSlice(arena, arr);
-    return out;
-}
-
 fn unset_path(arena: std.mem.Allocator, pairs: *std.ArrayListUnmanaged(bson.Pair), segs: []const []const u8) void {
     if (segs.len == 1) {
         if (find_pair(pairs.items, segs[0])) |idx| {
@@ -258,7 +247,7 @@ fn unset_path(arena: std.mem.Allocator, pairs: *std.ArrayListUnmanaged(bson.Pair
     const idx = find_pair(pairs.items, segs[0]) orelse return;
     switch (pairs.items[idx].value) {
         .doc => |sub| {
-            var sub_pairs = copy_pairs_to_list(arena, sub) catch return;
+            var sub_pairs = copy_to_list(bson.Pair, arena, sub) catch return;
             unset_path(arena, &sub_pairs, segs[1..]);
             pairs.items[idx].value = .{ .doc = sub_pairs.items };
         },
@@ -273,14 +262,7 @@ fn pull_matches(arena: std.mem.Allocator, condition: bson.Value, elem: bson.Valu
                 .doc => |pairs| pairs,
                 else => return false,
             };
-            var all_operators = cond_pairs.len > 0;
+            if (query.all_operator_keys(cond_pairs)) {
-            for (cond_pairs) |p| {
-                if (p.key.len == 0 or p.key[0] != '$') {
-                    all_operators = false;
-                    break;
-                }
-            }
-            if (all_operators) {
                 // Operator condition against the element's value at each
                 // operator's field — treat element doc as the doc.
                 var ok = true;
@@ -297,8 +279,7 @@ fn pull_matches(arena: std.mem.Allocator, condition: bson.Value, elem: bson.Valu
     }
 }
 
-fn numeric_add(arena: std.mem.Allocator, a: bson.Value, b: bson.Value) UpdateError!bson.Value {
+fn numeric_add(a: bson.Value, b: bson.Value) UpdateError!bson.Value {
-    _ = arena;
     if (a == .double or b == .double) {
         const sum: f64 = @floatCast(a.as_f128() + b.as_f128());
         return .{ .double = sum };

src/wire.zig

@@ -13,9 +13,6 @@ pub const op_code_reply: i32 = 2001;
 pub const max_message_size: usize = 48 * 1024 * 1024;
 pub const max_bson_object_size: i32 = 16 * 1024 * 1024;
 
-pub const flag_checksum_present: u32 = 1 << 1;
-pub const flag_more_to_come: u32 = 1 << 0;
-
 pub const Seq = struct {
     name: []const u8,
     docs: []bson.Document,
@@ -175,6 +172,37 @@ pub const Message = struct {
             else => null,
         };
     }
+
+    /// Documents of a batch argument (`documents`, `updates`, `deletes`).
+    /// Drivers send them either as an OP_MSG document sequence or as an array
+    /// inside the command body; callers should not have to care which. The
+    /// result borrows this message's storage — never deinit the documents.
+    pub fn batch(self: *Message, name: []const u8) BatchError![]const bson.Document {
+        for (self.seqs) |seq| {
+            if (std.mem.eql(u8, seq.name, name) and seq.docs.len > 0) return seq.docs;
+        }
+        const arr = switch (bson.get_pair(self.body.pairs, name) orelse return error.MissingBatch) {
+            .array => |a| a,
+            else => return error.BatchNotArray,
+        };
+        const docs = try self.arena.allocator().alloc(bson.Document, arr.len);
+        for (arr, 0..) |item, i| {
+            docs[i] = switch (item) {
+                // Pairs are borrowed from the body document, which owns the
+                // arena; these views must never be deinited.
+                .doc => |pairs| .{ .arena = undefined, .pairs = pairs },
+                else => return error.BatchElementNotDoc,
+            };
+        }
+        return docs;
+    }
+};
+
+pub const BatchError = error{
+    MissingBatch, // no sequence and no body field of that name
+    BatchNotArray, // body field is present but not an array
+    BatchElementNotDoc, // an array element is not a document
+    OutOfMemory,
 };
 
 /// Builder for a command reply document. Strings for keys and values must
@@ -221,6 +249,32 @@ pub const Reply = struct {
     }
 };
 
+/// Write the 16-byte message header with a zero length placeholder. Returns
+/// the offset `end_message` needs to patch the length in.
+fn begin_message(
+    gpa: std.mem.Allocator,
+    out: *std.ArrayListUnmanaged(u8),
+    request_id: u32,
+    response_to: u32,
+    op_code: i32,
+) !usize {
+    const len_pos = out.items.len;
+    var header: [16]u8 = undefined;
+    std.mem.writeInt(u32, header[0..4], 0, .little); // patched by end_message
+    std.mem.writeInt(u32, header[4..8], request_id, .little);
+    std.mem.writeInt(u32, header[8..12], response_to, .little);
+    std.mem.writeInt(i32, header[12..16], op_code, .little);
+    try out.appendSlice(gpa, &header);
+    return len_pos;
+}
+
+/// Patch in the total length of the message started at `len_pos`.
+fn end_message(out: *std.ArrayListUnmanaged(u8), len_pos: usize) !void {
+    const total = out.items.len - len_pos;
+    if (total > std.math.maxInt(u32)) return error.MessageTooLarge;
+    std.mem.writeInt(u32, out.items[len_pos..][0..4], @intCast(total), .little);
+}
+
 /// Serialize an OP_MSG reply: header + flags + single body section.
 pub fn write_message(
     gpa: std.mem.Allocator,
@@ -230,18 +284,13 @@ pub fn write_message(
     body: []const bson.Pair,
     out: *std.ArrayListUnmanaged(u8),
 ) !void {
-    const len_pos = out.items.len;
+    const len_pos = try begin_message(gpa, out, request_id, response_to, op_code_msg);
-    var header: [20]u8 = undefined;
+    var flag_bytes: [4]u8 = undefined;
-    std.mem.writeInt(u32, header[4..8], request_id, .little);
+    std.mem.writeInt(u32, &flag_bytes, flags, .little);
-    std.mem.writeInt(u32, header[8..12], response_to, .little);
+    try out.appendSlice(gpa, &flag_bytes);
-    std.mem.writeInt(i32, header[12..16], op_code_msg, .little);
-    std.mem.writeInt(u32, header[16..20], flags, .little);
-    try out.appendSlice(gpa, &header);
     try out.append(gpa, 0x00); // single body section
     try bson.write_doc(body, gpa, out);
-    const total = out.items.len - len_pos;
+    try end_message(out, len_pos);
-    if (total > std.math.maxInt(u32)) return error.MessageTooLarge;
-    std.mem.writeInt(u32, out.items[len_pos..][0..4], @intCast(total), .little);
 }
 
 /// Serialize an OP_REPLY (legacy): header + responseFlags | cursorID |
@@ -254,19 +303,12 @@ pub fn write_reply_query(
     body: []const bson.Pair,
     out: *std.ArrayListUnmanaged(u8),
 ) !void {
-    const len_pos = out.items.len;
+    const len_pos = try begin_message(gpa, out, request_id, response_to, op_code_reply);
-    var header: [16]u8 = undefined;
-    std.mem.writeInt(u32, header[4..8], request_id, .little);
-    std.mem.writeInt(u32, header[8..12], response_to, .little);
-    std.mem.writeInt(i32, header[12..16], op_code_reply, .little);
-    try out.appendSlice(gpa, &header);
     var reply_fields: [20]u8 = [_]u8{0} ** 20; // responseFlags | cursorID | startingFrom | numberReturned
     std.mem.writeInt(i32, reply_fields[16..20], 1, .little); // numberReturned = 1
     try out.appendSlice(gpa, &reply_fields);
     try bson.write_doc(body, gpa, out);
-    const total = out.items.len - len_pos;
+    try end_message(out, len_pos);
-    if (total > std.math.maxInt(u32)) return error.MessageTooLarge;
-    std.mem.writeInt(u32, out.items[len_pos..][0..4], @intCast(total), .little);
 }
 
 // ---------------------------------------------------------------------------