const std = @import("std");

/// A container for all existing instances of T.
/// This is like a dedicated allocator for the type, T, with a refcount system.
pub fn Store(comptime T: type) type {
    const VTable = struct {
        udata: ?*anyopaque = null,
        deinitItem: ?*fn (?*anyopaque, *T) void = null,
    };
    return struct {
        ///
        const Item = T;

        /// A struct which holds an item and its reference count.
        const Slot = struct {
            item  : Item,
            refcnt: usize,

            pub fn ref(self: *@This()) void {
                self.refcnt += 1;
            }
            pub fn unref(self: *@This()) void {
                self.refcnt -= 1;
            }
        };

        ///
        const Error = error {
            DetectedItemLeak,
        };

        ///
        alloc : std.mem.Allocator,

        ///
        vtable: VTable,

        ///
        slots : std.ArrayList(Slot),

        ///
        pub fn init(alloc: std.mem.Allocator, vtable: VTable) @This() {
            return .{
                .alloc  = alloc,
                .vtable = vtable,
                .slots  = std.ArrayList(Slot).init(alloc),
            };
        }
        ///
        pub fn deinit(self: *@This()) !void {
            defer self.slots.deinit();

            try self.collectGarbage();
            if (self.slots.items.len > 0) {
                return Error.DetectedItemLeak;
            }
        }

        /// Returns a slot which contains newly-allocated item.
        /// - A reference count of the returned slot is 1
        /// - The returned slot and its contents are alive till the reference count becomes 0
        pub fn add(self: *@This(), item: Item) !*Slot {
            for (self.slots.items) |slot| {
                if (slot.refcnt == 0) {
                    self.deinitItem(&slot.item);
                    slot.* = .{ .item = item, .refcnt = 1, };
                    return slot;
                }
            }

            const newSlot = try self.alloc.create(Slot);
            errdefer self.alloc.destroy(newSlot);

            newSlot.* = .{ .item = item, .refcnt = 1, };
            try self.slots.append(newSlot);
            return newSlot;
        }

        /// Releases memory of slots whose reference count is 0.
        /// Usually no need to call this function because such slots will be reused.
        pub fn collectGarbage(self: *@This()) !void {
            var slots = self.slots.items;

            var set: usize = 0;
            var get: usize = 0;
            while (set < slots.len) {
                if (slots[set].refcnt == 0) {
                    self.deinitItem(&slots[set].item);
                    self.alloc.destroy(slots[set]);

                    get += 1;
                    if (get >= slots.len) {
                        break;
                    }
                    slots[set] = slots[get];
                } else {
                    set += 1;
                    get += 1;
                }
            }
            try self.slots.resize(set);
        }

        fn deinitItem(self: *@This(), item: *Item) void {
            if (self.vtable.deinitItem != null) {
                self.vtable.deinitItem.?(self.vtable.udata, item);
            }
        }
    };
}

test "add new item" {
    const Item = u16;

    var sut = Store(Item).init(std.testing.allocator, .{});
    defer sut.deinit() catch unreachable;

    {
        var slot = try sut.add(123);
        defer slot.unref();

        // check initial value
        try std.testing.expect(slot.refcnt == 1);
        try std.testing.expect(slot.item   == 123);

        // check refcnt handling
        slot.ref();
        try std.testing.expect(slot.refcnt == 2);
        slot.unref();
        try std.testing.expect(slot.refcnt == 1);
    }
}
test "reuse an expired slot by adding right after removal" {
    const Item = u16;

    var sut = Store(Item).init(std.testing.allocator, .{});
    defer sut.deinit() catch unreachable;

    var address: usize = undefined;
    {
        var slot = try sut.add(123);
        defer slot.unref();
        address = @intFromPtr(slot);
        try std.testing.expect(slot.item == 123);
    }
    {
        var slot = try sut.add(456);
        defer slot.unref();
        try std.testing.expect(address == @intFromPtr(slot));
        try std.testing.expect(slot.item == 456);
    }
}
test "reusing doesn't happen by adding twice" {
    const Item = u16;

    var sut = Store(Item).init(std.testing.allocator, .{});
    defer sut.deinit() catch unreachable;

    var slot1 = try sut.add(123);
    defer slot1.unref();

    var slot2 = try sut.add(456);
    defer slot2.unref();

    try std.testing.expect(slot1 != slot2);
    try std.testing.expect(slot1.item == 123);
    try std.testing.expect(slot2.item == 456);
}
test "chaotic addition and removals" {
    const Item = u16;

    var sut = Store(Item).init(std.testing.allocator, .{});
    defer sut.deinit() catch unreachable;

    var slots = std.ArrayList(*Store(Item).Slot).init(std.testing.allocator);
    defer {
        for (slots.items) |slot| { slot.unref(); }
        slots.deinit();
    }

    // add 0~300
    for (0..300) |i| {
        const x: u16 = @intCast(i);
        try slots.append(try sut.add(x));
    }

    // removes other than multiples of 3
    for (0..100) |i| {
        slots.orderedRemove(300 - (i+1)*3 + 2).unref();
        slots.orderedRemove(300 - (i+1)*3 + 1).unref();
    }

    // add multiples of 3 from 300 to 600
    for (0..100) |i| {
        const x: u16 = @intCast(i);
        try slots.append(try sut.add(300 + x*3));
    }

    // check if the list is composed of multiples of 3 from 0~600
    for (0..200) |i| {
        const x: u16 = @intCast(i);
        try std.testing.expectEqual(slots.items[i].item, x*3);
    }
}