const std = @import("../std.zig");
const debug_mode = @import("builtin").mode == .Debug;
pub const MemoryPoolError = error{OutOfMemory};
pub fn MemoryPool(comptime Item: type) type {
return MemoryPoolAligned(Item, @alignOf(Item));
}
pub fn MemoryPoolAligned(comptime Item: type, comptime alignment: u29) type {
if (@alignOf(Item) == alignment) {
return MemoryPoolExtra(Item, .{});
} else {
return MemoryPoolExtra(Item, .{ .alignment = alignment });
}
}
pub const Options = struct {
alignment: ?u29 = null,
growable: bool = true,
};
pub fn MemoryPoolExtra(comptime Item: type, comptime pool_options: Options) type {
return struct {
const Pool = @This();
pub const item_size = std.math.max(@sizeOf(Node), @sizeOf(Item));
pub const item_alignment = std.math.max(@alignOf(Node), pool_options.alignment orelse 0);
const Node = struct {
next: ?*@This(),
};
const NodePtr = *align(item_alignment) Node;
const ItemPtr = *align(item_alignment) Item;
arena: std.heap.ArenaAllocator,
free_list: ?NodePtr = null,
pub fn init(allocator: std.mem.Allocator) Pool {
return .{ .arena = std.heap.ArenaAllocator.init(allocator) };
}
pub fn initPreheated(allocator: std.mem.Allocator, initial_size: usize) MemoryPoolError!Pool {
var pool = init(allocator);
errdefer pool.deinit();
var i: usize = 0;
while (i < initial_size) : (i += 1) {
const raw_mem = try pool.allocNew();
const free_node = @ptrCast(NodePtr, raw_mem);
free_node.* = Node{
.next = pool.free_list,
};
pool.free_list = free_node;
}
return pool;
}
pub fn deinit(pool: *Pool) void {
pool.arena.deinit();
pool.* = undefined;
}
pub fn reset(pool: *Pool) void {
const allocator = pool.arena.child_allocator;
pool.arena.deinit();
pool.arena = std.heap.ArenaAllocator.init(allocator);
pool.free_list = null;
}
pub fn create(pool: *Pool) !ItemPtr {
const node = if (pool.free_list) |item| blk: {
pool.free_list = item.next;
break :blk item;
} else if (pool_options.growable)
@ptrCast(NodePtr, try pool.allocNew())
else
return error.OutOfMemory;
const ptr = @ptrCast(ItemPtr, node);
ptr.* = undefined;
return ptr;
}
pub fn destroy(pool: *Pool, ptr: ItemPtr) void {
ptr.* = undefined;
const node = @ptrCast(NodePtr, ptr);
node.* = Node{
.next = pool.free_list,
};
pool.free_list = node;
}
fn allocNew(pool: *Pool) MemoryPoolError!*align(item_alignment) [item_size]u8 {
const mem = try pool.arena.allocator().alignedAlloc(u8, item_alignment, item_size);
return mem[0..item_size];
}
};
}
test "memory pool: basic" {
var pool = MemoryPool(u32).init(std.testing.allocator);
defer pool.deinit();
const p1 = try pool.create();
const p2 = try pool.create();
const p3 = try pool.create();
try std.testing.expect(p1 != p2);
try std.testing.expect(p1 != p3);
try std.testing.expect(p2 != p3);
pool.destroy(p2);
const p4 = try pool.create();
try std.testing.expect(p2 == p4);
}
test "memory pool: preheating (success)" {
var pool = try MemoryPool(u32).initPreheated(std.testing.allocator, 4);
defer pool.deinit();
_ = try pool.create();
_ = try pool.create();
_ = try pool.create();
}
test "memory pool: preheating (failure)" {
var failer = std.testing.FailingAllocator.init(std.testing.allocator, 0);
try std.testing.expectError(error.OutOfMemory, MemoryPool(u32).initPreheated(failer.allocator(), 5));
}
test "memory pool: growable" {
var pool = try MemoryPoolExtra(u32, .{ .growable = false }).initPreheated(std.testing.allocator, 4);
defer pool.deinit();
_ = try pool.create();
_ = try pool.create();
_ = try pool.create();
_ = try pool.create();
try std.testing.expectError(error.OutOfMemory, pool.create());
}