MultiforaDB/src/stress.zig

// Dev stress test for the index B+tree (not part of the build):
//    zig run -O ReleaseFast src/stress.zig
// Bulk-builds 30k entries, inserts 20k more one at a time (splits at depth 2),
// deletes 16.6k randomly (empty-leaf cascades), drains everything, and checks
// lookups against a brute-force model throughout.
const std = @import("std");
const index = @import("index.zig");
const bson = @import("bson.zig");

fn doc_of(gpa: std.mem.Allocator, pairs: []const bson.Pair) ![]u8 {
    var out: std.ArrayListUnmanaged(u8) = .empty;
    defer out.deinit(gpa);
    try bson.write_doc(pairs, gpa, &out);
    return out.toOwnedSlice(gpa);
}

const Fact = struct { a: i32, b: i32 };

fn check_range(
    gpa: std.mem.Allocator,
    ix: *const index.Index,
    prefix: bson.Value,
    lo: ?bson.Value,
    hi: ?bson.Value,
    facts: []const Fact,
    alive: []const bool,
) !void {
    var out: std.ArrayListUnmanaged([]const u8) = .empty;
    defer out.deinit(gpa);
    try ix.lookup_range(gpa, &.{prefix}, lo, true, hi, false, &out);
    var expected: usize = 0;
    for (facts, 0..) |f, fi| {
        if (!alive[fi]) continue;
        if (f.a != prefix.int32) continue;
        if (lo) |l| if (f.b < l.int32) continue;
        if (hi) |h| if (f.b >= h.int32) continue;
        expected += 1;
    }
    if (out.items.len != expected) {
        std.debug.print("MISMATCH: prefix={d} lo={?d} hi={?d}: got {d}, want {d}\n", .{
            prefix.int32,
            if (lo) |l| l.int32 else null,
            if (hi) |h| h.int32 else null,
            out.items.len,
            expected,
        });
        std.process.exit(1);
    }
}

pub fn main() !void {
    const gpa = std.heap.page_allocator;
    var prng = std.Random.DefaultPrng.init(0xBADCAFE);
    const rand = prng.random();

    const N = 30_000;
    var ids: std.ArrayListUnmanaged([]u8) = .empty;
    var facts: std.ArrayListUnmanaged(Fact) = .empty;
    var alive: std.ArrayListUnmanaged(bool) = .empty;
    var pairs: [2]bson.Pair = undefined;

    // 1. Bulk build an index over N docs.
    var keys = [_]index.IndexKey{ .{ .path = "a", .descending = false }, .{ .path = "b", .descending = false } };
    var ix = try index.Index.init(gpa, "ab", &keys, false, false, null);
    for (0..N) |i| {
        const id = try std.fmt.allocPrint(gpa, "id{d}", .{i});
        try ids.append(gpa, id);
        const a = rand.intRangeAtMost(i32, 0, 99);
        const b = rand.intRangeAtMost(i32, 0, 999);
        try facts.append(gpa, .{ .a = a, .b = b });
        try alive.append(gpa, true);
        pairs[0] = .{ .key = "a", .value = .{ .int32 = a } };
        pairs[1] = .{ .key = "b", .value = .{ .int32 = b } };
        const d = try doc_of(gpa, &pairs);
        try ix.append_doc_entries(gpa, d, id);
    }
    _ = try ix.finish_bulk(gpa, false);
    std.debug.print("bulk: count={d} leaves={d} depth={d} nodes={d}\n", .{
        ix.count(),
        ix.leaf_count,
        ix.depth,
        ix.nodes.items.len,
    });
    if (ix.count() != N) return error.BadCount;

    // Random range checks against brute force.
    for (0..500) |_| {
        const a = rand.intRangeAtMost(i32, 0, 99);
        const lo_v = rand.intRangeAtMost(i32, -10, 1009);
        const hi_v = rand.intRangeAtMost(i32, -10, 1009);
        const use_lo = rand.boolean();
        const use_hi = rand.boolean();
        try check_range(gpa, &ix, .{ .int32 = a }, if (use_lo) .{ .int32 = lo_v } else null, if (use_hi) .{ .int32 = hi_v } else null, facts.items, alive.items);
    }

    // 2. Incremental inserts, random order (splits + rebalancing-free path).
    const M = 20_000;
    for (0..M) |i| {
        const id = try std.fmt.allocPrint(gpa, "new{d}", .{i});
        try ids.append(gpa, id);
        const a = rand.intRangeAtMost(i32, 0, 99);
        const b = rand.intRangeAtMost(i32, 0, 999);
        try facts.append(gpa, .{ .a = a, .b = b });
        try alive.append(gpa, true);
        pairs[0] = .{ .key = "a", .value = .{ .int32 = a } };
        pairs[1] = .{ .key = "b", .value = .{ .int32 = b } };
        const d = try doc_of(gpa, &pairs);
        _ = try ix.add_doc(gpa, d, id, false);
    }
    std.debug.print("after inserts: count={d} leaves={d} depth={d} nodes={d}\n", .{
        ix.count(),
        ix.leaf_count,
        ix.depth,
        ix.nodes.items.len,
    });
    if (ix.count() != N + M) return error.BadCount;
    for (0..500) |_| {
        const a = rand.intRangeAtMost(i32, 0, 99);
        const lo_v = rand.intRangeAtMost(i32, -10, 1009);
        const hi_v = rand.intRangeAtMost(i32, -10, 1009);
        const use_lo = rand.boolean();
        const use_hi = rand.boolean();
        try check_range(gpa, &ix, .{ .int32 = a }, if (use_lo) .{ .int32 = lo_v } else null, if (use_hi) .{ .int32 = hi_v } else null, facts.items, alive.items);
    }

    // 3. Delete every 3rd doc in random order (empty-leaf cascades,
    //    one-child internals).
    var order: std.ArrayListUnmanaged(usize) = .empty;
    for (0..N + M) |i| if (i % 3 == 0) try order.append(gpa, i);
    rand.shuffle(usize, order.items);
    for (order.items) |i| {
        const a = facts.items[i].a;
        const b = facts.items[i].b;
        pairs[0] = .{ .key = "a", .value = .{ .int32 = a } };
        pairs[1] = .{ .key = "b", .value = .{ .int32 = b } };
        const d = try doc_of(gpa, &pairs);
        ix.remove_doc(gpa, d, ids.items[i]);
        alive.items[i] = false;
    }
    std.debug.print("after deletes: count={d} leaves={d} depth={d} nodes={d}\n", .{
        ix.count(),
        ix.leaf_count,
        ix.depth,
        ix.nodes.items.len,
    });
    if (ix.count() != (N + M) - (N + M) / 3 - 1) return error.BadCount;
    for (0..500) |_| {
        const a = rand.intRangeAtMost(i32, 0, 99);
        const lo_v = rand.intRangeAtMost(i32, -10, 1009);
        const hi_v = rand.intRangeAtMost(i32, -10, 1009);
        const use_lo = rand.boolean();
        const use_hi = rand.boolean();
        try check_range(gpa, &ix, .{ .int32 = a }, if (use_lo) .{ .int32 = lo_v } else null, if (use_hi) .{ .int32 = hi_v } else null, facts.items, alive.items);
    }

    // 4. Equality lookups still exact.
    for (0..300) |_| {
        const a = rand.intRangeAtMost(i32, 0, 99);
        var out: std.ArrayListUnmanaged([]const u8) = .empty;
        defer out.deinit(gpa);
        try ix.lookup_eq(gpa, &.{.{ .int32 = a }}, &out);
        var expected: usize = 0;
        for (facts.items, 0..) |f, fi| {
            if (alive.items[fi] and f.a == a) expected += 1;
        }
        if (out.items.len != expected) {
            std.debug.print("EQ MISMATCH a={d}: got {d} want {d}\n", .{
                a,
                out.items.len,
                expected,
            });
            return error.BadCount;
        }
    }

    // 5. Delete everything (empty-leaf cascades, one-child internals),
    //    then verify the tree still works for fresh inserts.
    var live: std.ArrayListUnmanaged(usize) = .empty;
    defer live.deinit(gpa);
    for (facts.items, 0..) |_, i| if (alive.items[i]) try live.append(gpa, i);
    rand.shuffle(usize, live.items);
    var remaining = live.items.len;
    for (live.items) |i| {
        const a = facts.items[i].a;
        const b = facts.items[i].b;
        pairs[0] = .{ .key = "a", .value = .{ .int32 = a } };
        pairs[1] = .{ .key = "b", .value = .{ .int32 = b } };
        const d = try doc_of(gpa, &pairs);
        ix.remove_doc(gpa, d, ids.items[i]);
        remaining -= 1;
        if (ix.count() != remaining) {
            std.debug.print("COUNT MISMATCH during drain: {d} != {d}\n", .{
                ix.count(),
                remaining,
            });
            return error.BadCount;
        }
    }
    std.debug.print("after full drain: count={d} leaves={d} depth={d} nodes={d}\n", .{
        ix.count(),
        ix.leaf_count,
        ix.depth,
        ix.nodes.items.len,
    });
    if (ix.count() != 0) return error.BadCount;
    // The drained tree still accepts and finds entries.
    pairs[0] = .{ .key = "a", .value = .{ .int32 = 7 } };
    pairs[1] = .{ .key = "b", .value = .{ .int32 = 42 } };
    const d2 = try doc_of(gpa, &pairs);
    defer gpa.free(d2);
    _ = try ix.add_doc(gpa, d2, "final", false);
    var out: std.ArrayListUnmanaged([]const u8) = .empty;
    defer out.deinit(gpa);
    try ix.lookup_eq(gpa, &.{.{ .int32 = 7 }}, &out);
    if (out.items.len != 1) return error.BadCount;
    var it = ix.iter();
    if (it.next() == null) return error.BadCount;
    if (it.next() != null) return error.BadCount;

    std.debug.print("STRESS OK\n", .{});
}