From 1a5386ff00608d1322a6b9dd3a2b32da4941e55b Mon Sep 17 00:00:00 2001
From: "A.Shakhmatov" <aleksey.shakhmatov@mvideo.ru>
Date: Sun, 9 Aug 2026 22:37:18 +0300
Subject: [PATCH] commands: the stages that rewrite a document

`$addFields`, `$set`, `$unset`, `$replaceRoot` and `$unwind`. The corpus goes
0 pass / 24 fail to 21 / 3, and the three left are `$project`'s computed
fields, renames and nested inclusions, which want the `query.project` that
`find` shares and are their own change.

**The design review was wrong about what this needed, and the corpus is what
settled it.** Tier 2 was scoped as a per-stage iterator on the grounds that
`$unwind` is 1->N and "there is no way to express that in a window over the
input". That was true of the window as it stood, and stopped being true the
moment `$project` was made to rebuild the stream instead of moving bounds over
it -- a stage that rebuilds can emit as many documents as it likes, or none. So
all five share one shape: read the window, build a new list, replace the
stream. No iterator, no rewrite.

What the recording settled, and what a hand-written test would have got wrong:

  - `$addFields` whose expression resolves to nothing leaves the field out
    entirely rather than setting it to null -- so `set_path` is only reached
    when there is a value, and `eval_expr`'s absent/null distinction earns its
    keep a second time.
  - `$addFields: {"n.z": 1}` sets the nested path and keeps its siblings, and
    an existing field is replaced *where it stands*, which is what makes the
    stage "add or overwrite" rather than "append".
  - `$unwind` drops a document whose field is missing or an empty array, keeps
    one whose field is not an array *whole*, and numbers `includeArrayIndex`
    from zero. Three separate behaviours where one guess would have covered
    them all wrongly.
  - `$replaceRoot` of a missing path and of a non-document are the same error,
    40228.

`ReplaceRootNotDocument` joins `EvalError` rather than being reported at the
stage: it is a failure only a document can produce, which is the line that set
already draws.

191/191 unit tests in ReleaseFast and ReleaseSafe, 83/83 fuzz, e2e 49, e2e3 16,
e2e4 17, e2e6 72, e2e7 86, crud corpus unchanged at 201/90/196.
---
 src/commands.zig               | 297 ++++++++++++++++++++++++++++++++-
 tests/spec/aggregate/README.md |  11 +-
 2 files changed, 302 insertions(+), 6 deletions(-)

diff --git a/src/commands.zig b/src/commands.zig
index 06d9465..62320c2 100644
--- a/src/commands.zig
+++ b/src/commands.zig
@@ -116,6 +116,8 @@ pub const ErrorCode = enum(i32) {
     location_switch_no_default = 40069,
     location_divide_by_zero = 4848401,
     location_non_numeric_arithmetic = 7157723,
+    location_replace_root_not_document = 40228,
+    location_unwind_bad_path = 28818,
     location_unknown_group_operator = 15952,
     location_group_needs_id = 15955,
     location_accumulator_not_object = 40234,
@@ -2339,6 +2341,41 @@ fn cmd_aggregate(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void {
             in_trees = true;
             start = 0;
             end = trees.items.len;
+        } else if (std.mem.eql(u8, stage_name, "$addFields") or std.mem.eql(u8, stage_name, "$set") or
+            std.mem.eql(u8, stage_name, "$unset") or std.mem.eql(u8, stage_name, "$replaceRoot") or
+            std.mem.eql(u8, stage_name, "$unwind"))
+        {
+            // The stages that rewrite a document, all one shape: read the
+            // window, build a new list, replace the stream. The design review
+            // expected these to need a per-stage iterator because `$unwind` is
+            // 1->N -- true of the window as it stood, and no longer true once a
+            // stage rebuilds the list rather than moving bounds over it.
+            const arena = reply.arena_alloc();
+            var built: std.ArrayListUnmanaged(*const bson.Document) = .empty;
+            errdefer built.deinit(ctx.gpa);
+            const rewrite = try compile_rewrite(ctx, reply, arena, stage_name, stage[0].value) orelse return;
+            var w: usize = start;
+            while (w < end) : (w += 1) {
+                const doc = if (in_trees) trees.items[w] else try doc_tree(arena, coll, offs.items[w]);
+                const ec: EvalCtx = .{
+                    .arena = arena,
+                    .coll = coll,
+                    .src = .{ .docs = (&doc)[0..1] },
+                    .i = 0,
+                };
+                apply_rewrite(ec, rewrite, doc, &built, ctx.gpa) catch |err| {
+                    try report_eval_error(reply, err);
+                    return;
+                };
+            }
+            offs.deinit(ctx.gpa);
+            offs = .empty;
+            trees.deinit(ctx.gpa);
+            trees = built;
+            built = .empty;
+            in_trees = true;
+            start = 0;
+            end = trees.items.len;
         } else if (std.mem.eql(u8, stage_name, "$out") or std.mem.eql(u8, stage_name, "$merge")) {
             const is_out = std.mem.eql(u8, stage_name, "$out");
             // MongoDB requires either to be last, and this server needs it too:
@@ -2962,9 +2999,262 @@ fn compile_switch(reply: *wire.Reply, arena: std.mem.Allocator, spec: bson.Value
 
 /// A failure only a document can produce, so it cannot be caught at compile
 /// time. Both codes measured against mongod.
-const EvalError = error{ DivideByZero, SwitchNoDefault, NonNumericArithmetic } ||
+const EvalError = error{ DivideByZero, SwitchNoDefault, NonNumericArithmetic, ReplaceRootNotDocument } ||
     std.mem.Allocator.Error || error{ EndOfStream, Overflow, InvalidBson };
 
+/// A compiled document-rewriting stage.
+const Rewrite = union(enum) {
+    /// `$addFields` and `$set`, which are the same stage under two names.
+    add_fields: []const Expr.Field,
+    unset: []const []const u8,
+    replace_root: Expr,
+    unwind: UnwindSpec,
+};
+
+fn compile_rewrite(
+    ctx: *Context,
+    reply: *wire.Reply,
+    arena: std.mem.Allocator,
+    name: []const u8,
+    spec: bson.Value,
+) !?Rewrite {
+    _ = ctx;
+    if (std.mem.eql(u8, name, "$unwind")) {
+        return Rewrite{ .unwind = (try unwind_spec(reply, spec)) orelse return null };
+    }
+    if (std.mem.eql(u8, name, "$unset")) {
+        return Rewrite{ .unset = (try unset_paths(reply, arena, spec)) orelse return null };
+    }
+    if (std.mem.eql(u8, name, "$replaceRoot")) {
+        const d = doc_arg(spec) orelse {
+            try bad_value(reply, "$replaceRoot requires a document");
+            return null;
+        };
+        const new_root = bson.get_pair(d, "newRoot") orelse {
+            try bad_value(reply, "$replaceRoot requires newRoot");
+            return null;
+        };
+        return Rewrite{ .replace_root = (try compile_expr(reply, arena, new_root, "$replaceRoot's newRoot")) orelse return null };
+    }
+    const d = doc_arg(spec) orelse {
+        try bad_value(reply, "$addFields requires a document");
+        return null;
+    };
+    const fields = try arena.alloc(Expr.Field, d.len);
+    for (d, 0..) |p, i| {
+        const sub = (try compile_expr(reply, arena, p.value, "an $addFields value")) orelse return null;
+        fields[i] = .{ .key = p.key, .value = sub };
+    }
+    return Rewrite{ .add_fields = fields };
+}
+
+/// Put one document through a rewrite, appending whatever it produces. A stage
+/// may emit no documents (`$unwind` of an empty array) or several.
+fn apply_rewrite(
+    ec: EvalCtx,
+    rewrite: Rewrite,
+    doc: *const bson.Document,
+    out: *std.ArrayListUnmanaged(*const bson.Document),
+    gpa: std.mem.Allocator,
+) EvalError!void {
+    switch (rewrite) {
+        .add_fields => |fields| {
+            var pairs: []const bson.Pair = doc.pairs;
+            for (fields) |f| {
+                // A value that resolves to nothing leaves the field out
+                // entirely, rather than setting it to null.
+                const v = (try eval_expr(ec, f.value)) orelse continue;
+                pairs = try set_path(ec.arena, pairs, f.key, v);
+            }
+            try out.append(gpa, try tree_of(ec.arena, pairs));
+        },
+        .unset => |paths| {
+            var pairs: []const bson.Pair = doc.pairs;
+            for (paths) |path| pairs = try unset_path(ec.arena, pairs, path);
+            try out.append(gpa, try tree_of(ec.arena, pairs));
+        },
+        .replace_root => |e| {
+            const v = (try eval_expr(ec, e)) orelse return error.ReplaceRootNotDocument;
+            const pairs = switch (v) {
+                .doc => |p| p,
+                else => return error.ReplaceRootNotDocument,
+            };
+            try out.append(gpa, try tree_of(ec.arena, pairs));
+        },
+        .unwind => |spec| {
+            const found = try eval_expr(ec, .{ .path = spec.path });
+            const items: []const bson.Value = switch (found orelse bson.Value.null) {
+                .array => |a| a,
+                // A non-array is kept whole, and a missing or null field is
+                // dropped unless the caller asked to keep it.
+                .null => if (spec.keep_empty) &.{} else return,
+                else => {
+                    try out.append(gpa, doc);
+                    return;
+                },
+            };
+            if (items.len == 0) {
+                if (!spec.keep_empty) return;
+                var pairs = try unset_path(ec.arena, doc.pairs, spec.path);
+                if (spec.index_field) |f| pairs = try set_path(ec.arena, pairs, f, .null);
+                try out.append(gpa, try tree_of(ec.arena, pairs));
+                return;
+            }
+            for (items, 0..) |item, i| {
+                var pairs = try set_path(ec.arena, doc.pairs, spec.path, item);
+                if (spec.index_field) |f| {
+                    pairs = try set_path(ec.arena, pairs, f, .{ .int64 = @intCast(i) });
+                }
+                try out.append(gpa, try tree_of(ec.arena, pairs));
+            }
+        },
+    }
+}
+
+fn tree_of(arena: std.mem.Allocator, pairs: []const bson.Pair) EvalError!*const bson.Document {
+    const d = try arena.create(bson.Document);
+    d.* = .{ .arena = undefined, .pairs = pairs };
+    return d;
+}
+
+/// A document with `path` set to `v`, creating the intermediate documents a
+/// dotted path names. Siblings are kept and an existing field is replaced where
+/// it stands, which is what makes `$addFields` "add or overwrite" rather than
+/// "append".
+fn set_path(
+    arena: std.mem.Allocator,
+    pairs: []const bson.Pair,
+    path: []const u8,
+    v: bson.Value,
+) EvalError![]const bson.Pair {
+    const dot = std.mem.indexOfScalar(u8, path, '.');
+    const head = if (dot) |d| path[0..d] else path;
+    const value: bson.Value = if (dot) |d| blk: {
+        const existing = switch (bson.get_pair(pairs, head) orelse bson.Value.null) {
+            .doc => |sub| sub,
+            // A non-document in the way is replaced by one, as mongod does.
+            else => &.{},
+        };
+        break :blk .{ .doc = try set_path(arena, existing, path[d + 1 ..], v) };
+    } else v;
+
+    for (pairs, 0..) |p, i| {
+        if (!std.mem.eql(u8, p.key, head)) continue;
+        const out = try arena.alloc(bson.Pair, pairs.len);
+        @memcpy(out, pairs);
+        out[i] = .{ .key = head, .value = value };
+        return out;
+    }
+    const out = try arena.alloc(bson.Pair, pairs.len + 1);
+    @memcpy(out[0..pairs.len], pairs);
+    out[pairs.len] = .{ .key = head, .value = value };
+    return out;
+}
+
+/// The same document without `path`. A path naming nothing changes nothing.
+fn unset_path(
+    arena: std.mem.Allocator,
+    pairs: []const bson.Pair,
+    path: []const u8,
+) EvalError![]const bson.Pair {
+    const dot = std.mem.indexOfScalar(u8, path, '.');
+    const head = if (dot) |d| path[0..d] else path;
+    var out: std.ArrayListUnmanaged(bson.Pair) = .empty;
+    try out.ensureTotalCapacity(arena, pairs.len);
+    for (pairs) |p| {
+        if (!std.mem.eql(u8, p.key, head)) {
+            out.appendAssumeCapacity(p);
+            continue;
+        }
+        const d = dot orelse continue; // a leaf: drop it
+        const sub = switch (p.value) {
+            .doc => |x| x,
+            else => {
+                out.appendAssumeCapacity(p);
+                continue;
+            },
+        };
+        out.appendAssumeCapacity(.{
+            .key = p.key,
+            .value = .{ .doc = try unset_path(arena, sub, path[d + 1 ..]) },
+        });
+    }
+    return out.items;
+}
+
+/// The paths a `$unset` names: one string, or an array of them.
+fn unset_paths(reply: *wire.Reply, arena: std.mem.Allocator, v: bson.Value) !?[]const []const u8 {
+    switch (v) {
+        .string => |str| {
+            const one = try arena.alloc([]const u8, 1);
+            one[0] = str;
+            return one;
+        },
+        .array => |items| {
+            const out = try arena.alloc([]const u8, items.len);
+            for (items, 0..) |item, i| {
+                out[i] = switch (item) {
+                    .string => |str| str,
+                    else => {
+                        try bad_value(reply, "$unset takes field names");
+                        return null;
+                    },
+                };
+            }
+            return out;
+        },
+        else => {
+            try bad_value(reply, "$unset takes a field name or an array of them");
+            return null;
+        },
+    }
+}
+
+/// What a `$unwind` was asked to do. Both spellings land here: the bare
+/// `"$path"` string and the document form.
+const UnwindSpec = struct {
+    path: []const u8,
+    keep_empty: bool = false,
+    index_field: ?[]const u8 = null,
+};
+
+fn unwind_spec(reply: *wire.Reply, v: bson.Value) !?UnwindSpec {
+    const raw: bson.Value = switch (v) {
+        .doc => |d| bson.get_pair(d, "path") orelse {
+            try bad_value(reply, "$unwind requires a path");
+            return null;
+        },
+        else => v,
+    };
+    const str = switch (raw) {
+        .string => |x| x,
+        else => {
+            try bad_value(reply, "$unwind requires a string path");
+            return null;
+        },
+    };
+    if (str.len == 0 or str[0] != '$') {
+        const detail = try std.fmt.allocPrint(
+            reply.arena_alloc(),
+            "path option to $unwind stage should be prefixed with a '$': {s}",
+            .{str},
+        );
+        try reply.put_error(@intFromEnum(ErrorCode.location_unwind_bad_path), "Location28818", detail);
+        return null;
+    }
+    var spec: UnwindSpec = .{ .path = str[1..] };
+    if (v == .doc) {
+        if (bson.get_pair(v.doc, "preserveNullAndEmptyArrays")) |b| spec.keep_empty = query.truthy(b);
+        if (bson.get_pair(v.doc, "includeArrayIndex")) |f| {
+            spec.index_field = switch (f) {
+                .string => |x| x,
+                else => null,
+            };
+        }
+    }
+    return spec;
+}
+
 /// Turn a failure only a document could produce into the reply mongod gives.
 fn report_eval_error(reply: *wire.Reply, err: EvalError) !void {
     switch (err) {
@@ -2978,6 +3268,11 @@ fn report_eval_error(reply: *wire.Reply, err: EvalError) !void {
             "Location40069",
             "$switch could not find a matching branch for an input, and no default was specified.",
         ),
+        error.ReplaceRootNotDocument => try reply.put_error(
+            @intFromEnum(ErrorCode.location_replace_root_not_document),
+            "Location40228",
+            "$replaceRoot requires a document as its newRoot",
+        ),
         error.NonNumericArithmetic => try reply.put_error(
             @intFromEnum(ErrorCode.location_non_numeric_arithmetic),
             "Location7157723",
diff --git a/tests/spec/aggregate/README.md b/tests/spec/aggregate/README.md
index 08c9cfb..829f3a4 100644
--- a/tests/spec/aggregate/README.md
+++ b/tests/spec/aggregate/README.md
@@ -64,7 +64,7 @@ accumulators in:
 ```
 group-accumulators.json   19 pass    0 fail   0 skip
 expressions.json          27 pass    0 fail   0 skip
-document-stages.json       0 pass   24 fail   0 skip
+document-stages.json      21 pass    3 fail   0 skip
 ```
 
 `group-accumulators` found its first real disagreement on the way to 18:
@@ -73,10 +73,11 @@ that counted documents rather than numbers would have passed every test
 anybody would think to write by hand.
 
 `expressions.json` was recorded before the evaluator was written and read
-1 pass / 26 fail against it; it is green now. `document-stages.json` is the
-next tier's spec, recorded and not implemented -- `$addFields`/`$set`,
-`$unset`, `$replaceRoot`, `$unwind` and `$project`'s computed fields, none of
-which this server has, so it starts at zero.
+1 pass / 26 fail against it; it is green now. `document-stages.json` was recorded at
+0 pass / 24 fail; `$addFields`/`$set`, `$unset`, `$replaceRoot` and `$unwind`
+took it to 21. The three left are `$project`'s computed fields, its renames and
+its nested inclusions, which want the shared `query.project` and are their own
+change.
 
 What recording *that* settled: