FuzzContext

const FuzzContext = struct {
    control_allocator: Allocator,
    sample_allocator: Allocator,

    last_alloc_index: Alloc.Index,
    allocs: [max_alloc_count]Alloc,

    const max_alloc_count = 64;
    const max_action_count = 2 * max_alloc_count;

    const max_alloc_size = 16 << 10;

    const Alloc = struct {
        control_ptr: [*]u8,
        sample_ptr: [*]u8,
        common: packed struct(usize) {
            len: Len,
            alignment: Alignment,
            _: @Int(.unsigned, padding_bits) = 0,
        },

        const Len = enum(@Int(.unsigned, len_bits)) {
            free = (1 << len_bits) - 1,
            _,
        };
        const len_bits = @min(64, @bitSizeOf(usize)) - @bitSizeOf(Alignment);
        const padding_bits = @bitSizeOf(usize) - (len_bits + @bitSizeOf(Alignment));

        const Index = enum(usize) {
            none = std.math.maxInt(usize),
            _,
        };
    };

    const Action = enum {
        alloc,
        resize,
        remap,
        free,
    };

    const Init = struct {
        threaded_instance: *std.Io.Threaded,
        arena_state: *ArenaAllocator.State,
        control_instance: *std.heap.FixedBufferAllocator,
        sample_instance: *FuzzAllocator,
    };

    fn init(
        control_allocator: Allocator,
        sample_allocator: Allocator,
    ) FuzzContext {
        @disableInstrumentation();
        return .{
            .control_allocator = control_allocator,
            .sample_allocator = sample_allocator,
            .last_alloc_index = .none,
            .allocs = undefined,
        };
    }

    fn deinit(ctx: *FuzzContext) void {
        @disableInstrumentation();
        ctx.* = undefined;
    }

    fn check(ctx: *const FuzzContext, n_allocs: usize) !void {
        @disableInstrumentation();
        for (ctx.allocs[0..n_allocs]) |allocation| {
            const len: usize = switch (allocation.common.len) {
                .free => continue,
                _ => |len| @backingInt(len),
            };
            const control = allocation.control_ptr[0..len];
            const sample = allocation.sample_ptr[0..len];
            try std.testing.expectEqualSlices(u8, control, sample);
        }
    }

    fn doOneAlloc(ctx: *FuzzContext, len: usize, alignment: Alignment, index: Alloc.Index) void {
        @disableInstrumentation();
        assert(ctx.allocs[@backingInt(index)].common.len == .free);

        const control_ptr = ctx.control_allocator.rawAlloc(len, alignment, @returnAddress()) orelse
            return;
        const sample_ptr = ctx.sample_allocator.rawAlloc(len, alignment, @returnAddress()) orelse {
            ctx.control_allocator.rawFree(control_ptr[0..len], alignment, @returnAddress());
            return;
        };

        ctx.allocs[@backingInt(index)] = .{
            .control_ptr = control_ptr,
            .sample_ptr = sample_ptr,
            .common = .{
                .len = @fromBackingInt(@intCast(len)),
                .alignment = alignment,
            },
        };

        for (control_ptr[0..len], sample_ptr[0..len], 0..) |*control, *sample, i| {
            control.* = @truncate(i);
            sample.* = @truncate(i);
        }

        @atomicStore(Alloc.Index, &ctx.last_alloc_index, index, .release);
    }
    fn doOneResize(ctx: *FuzzContext, new_len: usize) void {
        @disableInstrumentation();

        const index = @atomicRmw(Alloc.Index, &ctx.last_alloc_index, .Xchg, .none, .acquire);
        if (index == .none) return;

        const allocation = &ctx.allocs[@backingInt(index)];
        assert(allocation.common.len != .free);
        const memory = allocation.sample_ptr[0..@backingInt(allocation.common.len)];
        const alignment = allocation.common.alignment;

        assert(alignment.check(@intFromPtr(allocation.control_ptr)));
        assert(alignment.check(@intFromPtr(allocation.sample_ptr)));

        // Since `resize` is fallible, we have to ensure that `control_allocator`
        // is always successful by reserving the memory we need beforehand.
        const new_control_ptr = ctx.control_allocator.rawAlloc(new_len, alignment, @returnAddress()) orelse
            return;
        if (ctx.sample_allocator.rawResize(memory, alignment, new_len, @returnAddress())) {
            const old_control = allocation.control_ptr[0..memory.len];
            const overlap = @min(memory.len, new_len);
            @memcpy(new_control_ptr[0..overlap], old_control[0..overlap]);
            ctx.control_allocator.rawFree(old_control, alignment, @returnAddress());
        } else {
            ctx.control_allocator.rawFree(new_control_ptr[0..new_len], alignment, @returnAddress());
            return;
        }

        ctx.allocs[@backingInt(index)] = .{
            .control_ptr = new_control_ptr,
            .sample_ptr = memory.ptr,
            .common = .{
                .len = @fromBackingInt(@intCast(new_len)),
                .alignment = alignment,
            },
        };

        if (new_len > memory.len) {
            for (
                allocation.control_ptr[memory.len..new_len],
                allocation.sample_ptr[memory.len..new_len],
                0..,
            ) |*control, *sample, i| {
                control.* = @truncate(i);
                sample.* = @truncate(i);
            }
        }

        @atomicStore(Alloc.Index, &ctx.last_alloc_index, index, .release);
    }
    fn doOneRemap(ctx: *FuzzContext, new_len: usize) void {
        @disableInstrumentation();
        return doOneResize(ctx, new_len);
    }
    fn doOneFree(ctx: *FuzzContext) void {
        @disableInstrumentation();

        const index = @atomicRmw(Alloc.Index, &ctx.last_alloc_index, .Xchg, .none, .acquire);
        if (index == .none) return;

        const allocation = &ctx.allocs[@backingInt(index)];
        assert(allocation.common.len != .free);
        const len: usize = @backingInt(allocation.common.len);
        const alignment = allocation.common.alignment;

        assert(alignment.check(@intFromPtr(allocation.control_ptr)));
        assert(alignment.check(@intFromPtr(allocation.sample_ptr)));

        ctx.control_allocator.rawFree(allocation.control_ptr[0..len], alignment, @returnAddress());
        ctx.sample_allocator.rawFree(allocation.sample_ptr[0..len], alignment, @returnAddress());

        ctx.allocs[@backingInt(index)] = .{
            .control_ptr = undefined,
            .sample_ptr = undefined,
            .common = .{
                .len = .free,
                .alignment = .@"1",
            },
        };
    }
}