diff --git a/lib/std/heap/memory_pool.zig b/lib/std/heap/memory_pool.zig index 2b201f2b5422..52783ecb83cd 100644 --- a/lib/std/heap/memory_pool.zig +++ b/lib/std/heap/memory_pool.zig @@ -39,22 +39,24 @@ pub fn MemoryPoolExtra(comptime Item: type, comptime pool_options: Options) type return struct { const Pool = @This(); - /// Size of the memory pool items. This is not necessarily the same - /// as `@sizeOf(Item)` as the pool also uses the items for internal means. + /// Size of the memory pool items. This may be larger than `@sizeOf(Item)` + /// as the pool also uses the items for internal means. pub const item_size = @max(@sizeOf(Node), @sizeOf(Item)); // This needs to be kept in sync with Node. const node_alignment: Alignment = .of(*anyopaque); - /// Alignment of the memory pool items. This is not necessarily the same - /// as `@alignOf(Item)` as the pool also uses the items for internal means. + /// Alignment of the memory pool items. This may be larger than `@alignOf(Item)` + /// as the pool also uses the items for internal means. pub const item_alignment: Alignment = node_alignment.max(pool_options.alignment orelse .of(Item)); - const Node = struct { - next: ?*align(item_alignment.toByteUnits()) @This(), - }; - const NodePtr = *align(item_alignment.toByteUnits()) Node; - const ItemPtr = *align(item_alignment.toByteUnits()) Item; + const Node = struct { next: ?*align(unit_al_bytes) @This() }; + const Byte = std.meta.Int(.unsigned, std.mem.byte_size_in_bits); + const Unit = [item_alignment.forward(item_size)]Byte; + const unit_al_bytes = item_alignment.toByteUnits(); + + const ItemPtr = *align(unit_al_bytes) Item; + const NodePtr = *align(unit_al_bytes) Node; arena: std.heap.ArenaAllocator, free_list: ?NodePtr = null, @@ -84,18 +86,29 @@ pub fn MemoryPoolExtra(comptime Item: type, comptime pool_options: Options) type /// This allows up to `size` active allocations before an /// `OutOfMemory` error might happen when calling `create()`. pub fn preheat(pool: *Pool, size: usize) MemoryPoolError!void { - var i: usize = 0; - while (i < size) : (i += 1) { - const raw_mem = try pool.allocNew(); - const free_node = @as(NodePtr, @ptrCast(raw_mem)); - free_node.* = Node{ - .next = pool.free_list, - }; + const raw_mem = try pool.allocNew(size); + const uni_slc = raw_mem[0..size]; + for (uni_slc) |*unit| { + const free_node: NodePtr = @ptrCast(unit); + free_node.next = pool.free_list; pool.free_list = free_node; } } - pub const ResetMode = std.heap.ArenaAllocator.ResetMode; + pub const ResetMode = union(enum) { + /// Releases all allocated memory in the arena. + free_all, + /// This will pre-heat the memory pool for future allocations by allocating a + /// large enough buffer to accomodate the highest amount of actively allocated items + /// in the past. Preheating will speed up the allocation process by invoking the + /// backing allocator less often than before. If `reset()` is used in a loop, this + /// means if the highest amount of actively allocated items is never being surpassed, + /// no memory allocations are performed anymore. + retain_capacity, + /// This is the same as `retain_capacity`, but the memory will be shrunk to + /// only hold at most this value of items. + retain_with_limit: usize, + }; /// Resets the memory pool and destroys all allocated items. /// This can be used to batch-destroy all objects without invalidating the memory pool. @@ -107,27 +120,34 @@ pub fn MemoryPoolExtra(comptime Item: type, comptime pool_options: Options) type /// /// NOTE: If `mode` is `free_all`, the function will always return `true`. pub fn reset(pool: *Pool, mode: ResetMode) bool { - // TODO: Potentially store all allocated objects in a list as well, allowing to - // just move them into the free list instead of actually releasing the memory. - - const reset_successful = pool.arena.reset(mode); - + const ArenaResetMode = std.heap.ArenaAllocator.ResetMode; + const arena_mode = switch (mode) { + .free_all => .free_all, + .retain_capacity => .retain_capacity, + .retain_with_limit => |limit| ArenaResetMode{ .retain_with_limit = limit * item_size }, + }; pool.free_list = null; - - return reset_successful; + if (!pool.arena.reset(arena_mode)) return false; + // When the backing arena allocator is being reset to + // a capacity greater than 0, then its internals consists + // of a *single* buffer node of said capacity. This means, + // we can safely pre-heat without causing additional allocations. + const arena_capacity = pool.arena.queryCapacity() / item_size; + if (arena_capacity != 0) pool.preheat(arena_capacity) catch unreachable; + return true; } /// Creates a new item and adds it to the memory pool. pub fn create(pool: *Pool) !ItemPtr { - const node = if (pool.free_list) |item| blk: { + const node_ptr: NodePtr = if (pool.free_list) |item| blk: { pool.free_list = item.next; break :blk item; } else if (pool_options.growable) - @as(NodePtr, @ptrCast(try pool.allocNew())) + @ptrCast(try pool.allocNew(1)) else return error.OutOfMemory; - const ptr = @as(ItemPtr, @ptrCast(node)); + const ptr: ItemPtr = @ptrCast(node_ptr); ptr.* = undefined; return ptr; } @@ -136,17 +156,14 @@ pub fn MemoryPoolExtra(comptime Item: type, comptime pool_options: Options) type /// Only pass items to `ptr` that were previously created with `create()` of the same memory pool! pub fn destroy(pool: *Pool, ptr: ItemPtr) void { ptr.* = undefined; - - const node = @as(NodePtr, @ptrCast(ptr)); - node.* = Node{ - .next = pool.free_list, - }; - pool.free_list = node; + const node_ptr: NodePtr = @ptrCast(ptr); + node_ptr.next = pool.free_list; + pool.free_list = node_ptr; } - fn allocNew(pool: *Pool) MemoryPoolError!*align(item_alignment.toByteUnits()) [item_size]u8 { - const mem = try pool.arena.allocator().alignedAlloc(u8, item_alignment, item_size); - return mem[0..item_size]; // coerce slice to array pointer + fn allocNew(pool: *Pool, n: usize) MemoryPoolError![*]align(unit_al_bytes) Unit { + const mem = try pool.arena.allocator().alignedAlloc(Unit, item_alignment, n); + return mem.ptr; } }; } @@ -220,3 +237,25 @@ test "greater than pointer manual alignment" { const foo: *align(16) Foo = try pool.create(); _ = foo; } + +test "reset" { + const pool_extra = MemoryPoolExtra(u32, .{ .growable = false }); + var pool = try pool_extra.initPreheated(std.testing.allocator, 3); + defer pool.deinit(); + + try std.testing.expect(pool.create() != error.OutOfMemory); + try std.testing.expect(pool.create() != error.OutOfMemory); + try std.testing.expect(pool.create() != error.OutOfMemory); + try std.testing.expect(pool.create() == error.OutOfMemory); + + try std.testing.expect(pool.reset(.{ .retain_with_limit = 2 })); + + try std.testing.expect(pool.create() != error.OutOfMemory); + try std.testing.expect(pool.create() != error.OutOfMemory); + try std.testing.expect(pool.create() == error.OutOfMemory); + + try std.testing.expect(pool.reset(.{ .retain_with_limit = 1 })); + + try std.testing.expect(pool.create() != error.OutOfMemory); + try std.testing.expect(pool.create() == error.OutOfMemory); +}