const std = @import("std");

// Wrapper allocator that counts what the largest amount allocated at any time was from backing_allocator.
pub const MaxAllocator = struct {
    backing_allocator: std.mem.Allocator,
    currently_allocated: usize = 0,
    max_allocated: usize = 0,

    pub fn allocator(self: *@This()) std.mem.Allocator {
        return std.mem.Allocator{
            .ptr = self,
            .vtable = &.{
                .alloc = rawAlloc,
                .resize = rawResize,
                .free = rawFree,
            },
        };
    }

    fn rawAlloc(ctx: *anyopaque, len: usize, ptr_align: u8, ret_addr: usize) ?[*]u8 {
        var self: *@This() = @alignCast(@ptrCast(ctx));
        var res = self.backing_allocator.rawAlloc(len, ptr_align, ret_addr);
        if (res != null) {
            self.currently_allocated += len;
            self.max_allocated = @max(self.currently_allocated, self.max_allocated);
        }
        return res;
    }

    fn rawResize(ctx: *anyopaque, buf: []u8, buf_align: u8, new_len: usize, ret_addr: usize) bool {
        var self: *@This() = @alignCast(@ptrCast(ctx));
        const res = self.backing_allocator.rawResize(buf, buf_align, new_len, ret_addr);
        if (res) {
            self.currently_allocated -= buf.len;
            self.currently_allocated += new_len;
            self.max_allocated = @max(self.currently_allocated, self.max_allocated);
        }
        return res;
    }

    fn rawFree(ctx: *anyopaque, buf: []u8, buf_align: u8, ret_addr: usize) void {
        var self: *@This() = @alignCast(@ptrCast(ctx));
        self.backing_allocator.rawFree(buf, buf_align, ret_addr);
        self.currently_allocated -= buf.len;
    }
};

pub fn maxAllocator(backing_allocator: std.mem.Allocator) MaxAllocator {
    return MaxAllocator{ .backing_allocator = backing_allocator };
}

test {
    const buflen = 2048;
    var buf = [_]u8{0} ** buflen;
    var fba = std.heap.FixedBufferAllocator.init(&buf);
    var ma = maxAllocator(fba.allocator());
    var a = ma.allocator();
    var x = try a.alloc(u8, 100);
    try std.testing.expectEqual(@as(usize, 100), ma.max_allocated);
    try std.testing.expectEqual(@as(usize, 100), ma.currently_allocated);
    var y = try a.alloc(u8, 10);
    try std.testing.expectEqual(@as(usize, 110), ma.max_allocated);
    try std.testing.expectEqual(@as(usize, 110), ma.currently_allocated);
    a.free(x);
    var z = try a.alloc(u8, 50);
    try std.testing.expectEqual(@as(usize, 110), ma.max_allocated);
    try std.testing.expectEqual(@as(usize, 60), ma.currently_allocated);

    try std.testing.expect(a.resize(z, 1024));
    z.len = 1024;
    try std.testing.expectEqual(@as(usize, 1034), ma.max_allocated);
    try std.testing.expectEqual(@as(usize, 1034), ma.currently_allocated);

    a.free(y);
    a.free(z);

    try std.testing.expectEqual(@as(usize, 1034), ma.max_allocated);
    try std.testing.expectEqual(@as(usize, 0), ma.currently_allocated);
}