MultiforaDB/src/spill2.zig

// Dev stress test: spilled records through leaf splits and internal levels.
//    zig run -O ReleaseFast src/spill2.zig
// 5000 docs with 2 KiB keys: every record spills; the tree still answers
// exact lookups and deletes half of them.
const std = @import("std");
const index = @import("index.zig");
const pgr = @import("pager.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);
}

/// A throwaway data file for this harness. The B+tree's pages live in the data
/// file now, so a standalone harness has to provide one.
fn harness_pager(gpa: std.mem.Allocator, name: []const u8) !*pgr.Pager {
    const threaded = try gpa.create(std.Io.Threaded);
    threaded.* = .init_single_threaded;
    const io = threaded.io();
    const path = try std.fmt.allocPrint(gpa, ".zig-cache/{s}.data", .{name});
    std.Io.Dir.cwd().deleteFile(io, path) catch {};
    const pg = try gpa.create(pgr.Pager);
    pg.* = try pgr.Pager.open(gpa, io, path, .{});
    return pg;
}

pub fn main() !void {
    const gpa = std.heap.page_allocator;

    var keys = [_]index.IndexKey{.{ .path = "tag", .descending = false }};
    var ix = try index.Index.init(gpa, try harness_pager(gpa, "spill2"), "tag", &keys, false, false, null);

    // 5000 docs, each with a 2 KiB key: spills on every record, forcing
    // leaves and internal nodes to hold overflow references.
    const N = 5000;
    var ids: std.ArrayListUnmanaged(u64) = .empty;
    defer ids.deinit(gpa);
    var pairs: [2]bson.Pair = undefined;
    var buf = try gpa.alloc(u8, 2000);
    defer gpa.free(buf);
    var facts: std.ArrayListUnmanaged(struct { key: []u8 }) = .empty;
    defer {
        for (facts.items) |f| gpa.free(f.key);
        facts.deinit(gpa);
    }
    for (0..N) |i| {
        for (buf) |*c| c.* = 'x';
        // unique suffix
        std.mem.writeInt(u32, buf[0..4], @intCast(i), .little);
        const key = try gpa.dupe(u8, buf);
        try facts.append(gpa, .{ .key = key });
        const id: u64 = @intCast(i + 1);
        try ids.append(gpa, id);
        pairs[0] = .{ .key = "_id", .value = .{ .int32 = @intCast(i) } };
        pairs[1] = .{ .key = "tag", .value = .{ .string = key } };
        const d = try doc_of(gpa, &pairs);
        _ = try ix.add_doc(gpa, d, id, false);
    }
    std.debug.print("count={d} leaves={d} depth={d} overflow={d}\n", .{
        ix.count(),
        ix.leaf_count,
        ix.depth,
        ix.ovf_tail,
    });
    if (ix.count() != N) return error.Bad;

    // Spot-check exact lookups.
    var prng2 = std.Random.DefaultPrng.init(0xabc);
    const rand2 = prng2.random();
    for (0..300) |_| {
        const i = rand2.intRangeAtMost(usize, 0, N - 1);
        var out: std.ArrayListUnmanaged(u64) = .empty;
        defer out.deinit(gpa);
        try ix.lookup_eq(gpa, &.{.{ .string = facts.items[i].key }}, &out);
        if (out.items.len != 1 or out.items[0] != ids.items[i]) {
            std.debug.print("lookup mismatch at {d}\n", .{i});
            return error.Bad;
        }
    }
    // Delete half (random), verify count and no leftover.
    var order: std.ArrayListUnmanaged(usize) = .empty;
    defer order.deinit(gpa);
    for (0..N) |i| if (i % 2 == 0) try order.append(gpa, i);
    rand2.shuffle(usize, order.items);
    var removed: usize = 0;
    for (order.items) |i| {
        pairs[0] = .{ .key = "_id", .value = .{ .int32 = @intCast(i) } };
        pairs[1] = .{ .key = "tag", .value = .{ .string = facts.items[i].key } };
        const d = try doc_of(gpa, &pairs);
        ix.remove_doc(gpa, d, ids.items[i]);
        removed += 1;
        if (ix.count() != N - removed) {
            std.debug.print("count mismatch at {d}: {d} != {d}\n", .{ i, ix.count(), N - removed });
            return error.Bad;
        }
    }
    for (order.items) |i| {
        var out: std.ArrayListUnmanaged(u64) = .empty;
        defer out.deinit(gpa);
        try ix.lookup_eq(gpa, &.{.{ .string = facts.items[i].key }}, &out);
        if (out.items.len != 0) return error.Bad;
    }
    std.debug.print("SPILL2 OK (leaves={d} depth={d})\n", .{ ix.leaf_count, ix.depth });
}