index_pool.zig

pub const std = @import("std");

pub fn IndexPool(comptime T: type) type {
    return struct {
        const Self = @This();

        pub const IdRange = struct {
            start: T,
            end: T,

            pub fn equals(self: IdRange, other: IdRange) bool {
                return self.start == other.start and self.end == other.end;
            }
        };

        allocator: std.mem.Allocator,
        avaliable: std.ArrayList(IdRange) = .empty,
        reserved: T = 0,

        pub fn pull(self: *Self) !?T {
            if (self.avaliable.items.len > 0) {
                var entry = &self.avaliable.items[self.avaliable.items.len - 1];
                if (entry.end == entry.start) {
                    const result = entry.start;
                    _ = self.avaliable.pop();
                    return result;
                }
                const res = entry.end;
                entry.end -= 1;
                return res;
            }
            return null;
        }

        pub fn init(allocator: std.mem.Allocator, reserved: T) !Self {
            var ranges = try std.ArrayList(IdRange).initCapacity(allocator, 256);
            if (reserved > 0) {
                try ranges.append(allocator, .{ .start = 0, .end = reserved - 1 });
            }
            return .{
                .allocator = allocator,
                .avaliable = ranges,
                .reserved = reserved,
            };
        }

        pub fn return_index(self: *Self, index: T) !void {
            if (self.avaliable.items.len == 0) {
                try self.avaliable.append(self.allocator, .{ .start = index, .end = index });
                return;
            }
            var lower: usize = 0;
            var upper: usize = self.avaliable.items.len;
            var mid: usize = lower;
            while (lower != upper) {
                mid = lower + ((upper - lower) / 2);
                if (index >= self.avaliable.items[mid].start and index <= self.avaliable.items[mid].end) {
                    std.debug.assert(false); // double free
                    return;
                } else if (index < self.avaliable.items[mid].start) {
                    upper = mid;
                } else {
                    lower = mid + 1;
                }
            }
            if (index < self.avaliable.items[mid].start) {
                if (self.avaliable.items[mid].start - index == 1) {
                    self.avaliable.items[mid].start -= 1;
                    if (mid > 0 and self.avaliable.items[mid - 1].end + 1 == self.avaliable.items[mid].start) {
                        self.avaliable.items[mid - 1].end = self.avaliable.items[mid].end; // merge
                        _ = self.avaliable.orderedRemove(mid);
                        mid -= 1;
                    }
                } else {
                    try self.avaliable.insert(self.allocator, mid, .{ .start = index, .end = index });
                }
            } else if (index > self.avaliable.items[mid].end) {
                if (index - self.avaliable.items[mid].end == 1) {
                    self.avaliable.items[mid].end += 1;
                    if (mid + 1 < self.avaliable.items.len and self.avaliable.items[mid + 1].start - 1 == self.avaliable.items[mid].end) {
                        self.avaliable.items[mid].end = self.avaliable.items[mid + 1].end; // merge
                        _ = self.avaliable.orderedRemove(mid + 1);
                    }
                } else {
                    try self.avaliable.insert(self.allocator, mid + 1, .{ .start = index, .end = index });
                }
            }
        }

        pub fn resize(self: *Self, new_size: T) !void {
            for (self.reserved..new_size) |i| {
                try self.return_index(@intCast(i));
            }
            self.reserved = new_size;
        }

        pub fn debug(self: *Self) void {
            std.debug.print("IndexPool Debug: reserved = {d}\n", .{self.reserved});
            for (self.avaliable.items, 0..) |range, idx| {
                std.debug.print("  Range {d}: [{d}, {d}]\n", .{idx, range.start, range.end});
            }
        }

        pub fn deinit(self: *Self) void {
            self.avaliable.deinit(self.allocator);
        }
    };
}

// Unit Tests
test "IndexPool basic initialization" {
    const testing = std.testing;
    var pool = try IndexPool(u32).init(testing.allocator, 10);
    defer pool.deinit();

    try testing.expect(pool.reserved == 10);
    try testing.expect(pool.avaliable.items.len == 1);
    try testing.expect(pool.avaliable.items[0].start == 0);
    try testing.expect(pool.avaliable.items[0].end == 9);
}

test "IndexPool pull indices in sequence" {
    const testing = std.testing;
    var pool = try IndexPool(u32).init(testing.allocator, 5);
    defer pool.deinit();

    // Pull all available indices
    const id1 = try pool.pull();
    try testing.expect(id1.? == 4); // Should pull from end first

    const id2 = try pool.pull();
    try testing.expect(id2.? == 3);

    const id3 = try pool.pull();
    try testing.expect(id3.? == 2);

    const id4 = try pool.pull();
    try testing.expect(id4.? == 1);

    const id5 = try pool.pull();
    try testing.expect(id5.? == 0);

    // Pool should be empty now
    const id6 = try pool.pull();
    try testing.expect(id6 == null);
}

test "IndexPool return and pull cycle" {
    const testing = std.testing;
    var pool = try IndexPool(u32).init(testing.allocator, 5);
    defer pool.deinit();

    // Pull some indices
    const id1 = try pool.pull();
    const id2 = try pool.pull();
    const id3 = try pool.pull();

    try testing.expect(id1.? == 4);
    try testing.expect(id2.? == 3);
    try testing.expect(id3.? == 2);

    // Return an index
    try pool.return_index(3);

    // Pull should now return the returned index
    const id4 = try pool.pull();
    try testing.expect(id4.? == 3);
}

test "IndexPool return index creates new range" {
    const testing = std.testing;
    var pool = try IndexPool(u32).init(testing.allocator, 5);
    defer pool.deinit();

    // Pull all indices
    _ = try pool.pull(); // 4
    _ = try pool.pull(); // 3
    _ = try pool.pull(); // 2
    _ = try pool.pull(); // 1
    _ = try pool.pull(); // 0

    // Pool should be empty
    try testing.expect(pool.avaliable.items.len == 0);

    // Return a non-adjacent index
    try pool.return_index(2);

    // Should create a new range
    try testing.expect(pool.avaliable.items.len == 1);
    try testing.expect(pool.avaliable.items[0].start == 2);
    try testing.expect(pool.avaliable.items[0].end == 2);
}

test "IndexPool range merging adjacent indices" {
    const testing = std.testing;
    var pool = try IndexPool(u32).init(testing.allocator, 5);
    defer pool.deinit();

    // Pull all indices
    _ = try pool.pull(); // 4
    _ = try pool.pull(); // 3
    _ = try pool.pull(); // 2
    _ = try pool.pull(); // 1
    _ = try pool.pull(); // 0

    // Return indices that should merge
    try pool.return_index(2);
    try pool.return_index(3); // Should merge with range [2,2] to become [2,3]

    try testing.expect(pool.avaliable.items.len == 1);
    try testing.expect(pool.avaliable.items[0].start == 2);
    try testing.expect(pool.avaliable.items[0].end == 3);

    // Return index 1, should extend the range to [1,3]
    try pool.return_index(1);

    try testing.expect(pool.avaliable.items.len == 1);
    try testing.expect(pool.avaliable.items[0].start == 1);
    try testing.expect(pool.avaliable.items[0].end == 3);
}

test "IndexPool range merging with multiple ranges" {
    const testing = std.testing;
    var pool = try IndexPool(u32).init(testing.allocator, 10);
    defer pool.deinit();

    // Pull all indices
    for (0..10) |_| {
        try testing.expect(try pool.pull() != null);
    }

    // Create separate ranges
    try pool.return_index(2);
    try pool.return_index(5);
    try pool.return_index(8);

    try testing.expect(pool.avaliable.items.len == 3);

    // Return index 3, should merge ranges [2,2] and [5,5] if we return 4 next
    try pool.return_index(3);
    try testing.expect(pool.avaliable.items.len == 3); // Still separate

    try pool.return_index(4);
    // Now [2,2], [3,3], [4,4], and [5,5] should all merge into [2,5]
    try testing.expect(pool.avaliable.items.len == 2); // [2,5] and [8,8]
}

test "IndexPool resize functionality" {
    const testing = std.testing;
    var pool = try IndexPool(u32).init(testing.allocator, 5);
    defer pool.deinit();

    // Pull some indices
    _ = try pool.pull(); // 4
    _ = try pool.pull(); // 3

    // Current available range should be [0,2]
    try testing.expect(pool.avaliable.items.len == 1);
    try testing.expect(pool.avaliable.items[0].start == 0);
    try testing.expect(pool.avaliable.items[0].end == 2);

    // Resize to add more indices
    try pool.resize(8);

    // Should have indices [0,2] and [5,7] available
    // Note: the implementation might merge or keep separate depending on order
    const total_available = blk: {
        var total: u32 = 0;
        for (pool.avaliable.items) |range| {
            total += (range.end - range.start + 1);
        }
        break :blk total;
    };
    try testing.expect(total_available == 6); // 3 original + 3 new = 6 total
}

test "IndexPool empty pool behavior" {
    const testing = std.testing;
    var pool = try IndexPool(u32).init(testing.allocator, 0);
    defer pool.deinit();

    // Should be empty initially
    const id = try pool.pull();
    try testing.expect(id == null);

    // Resize to add indices
    try pool.resize(3);

    // Now should have indices available
    const id1 = try pool.pull();
    const id2 = try pool.pull();
    const id3 = try pool.pull();

    try testing.expect(id1.? == 2);
    try testing.expect(id2.? == 1);
    try testing.expect(id3.? == 0);

    // Should be empty again
    const id4 = try pool.pull();
    try testing.expect(id4 == null);
}

test "IndexPool IdRange equals method" {
    const testing = std.testing;
    const Range = IndexPool(u32).IdRange;

    const range1 = Range{ .start = 1, .end = 5 };
    const range2 = Range{ .start = 1, .end = 5 };
    const range3 = Range{ .start = 1, .end = 6 };
    const range4 = Range{ .start = 2, .end = 5 };

    try testing.expect(range1.equals(range2));
    try testing.expect(!range1.equals(range3));
    try testing.expect(!range1.equals(range4));
}

test "IndexPool with different integer types" {
    const testing = std.testing;
    
    // Test with u8
    var pool_u8 = try IndexPool(u8).init(testing.allocator, 3);
    defer pool_u8.deinit();

    const id1 = try pool_u8.pull();
    try testing.expect(id1.? == 2);

    // Test with usize
    var pool_usize = try IndexPool(usize).init(testing.allocator, 2);
    defer pool_usize.deinit();

    const id2 = try pool_usize.pull();
    try testing.expect(id2.? == 1);
}