| ... | @@ -0,0 +1,181 @@ |
| 1 | const std = @import("../std.zig"); |
| 2 | |
| 3 | const debug_mode = @import("builtin").mode == .Debug; |
| 4 | |
| 5 | pub const MemoryPoolError = error{OutOfMemory}; |
| 6 | |
| 7 | /// A memory pool that can allocate objects of a single type very quickly. |
| 8 | /// Use this when you need to allocate a lot of objects of the same type, |
| 9 | /// because It outperforms general purpose allocators. |
| 10 | pub fn MemoryPool(comptime Item: type) type { |
| 11 | return MemoryPoolAligned(Item, @alignOf(Item)); |
| 12 | } |
| 13 | |
| 14 | /// A memory pool that can allocate objects of a single type very quickly. |
| 15 | /// Use this when you need to allocate a lot of objects of the same type, |
| 16 | /// because It outperforms general purpose allocators. |
| 17 | pub fn MemoryPoolAligned(comptime Item: type, comptime alignment: u29) type { |
| 18 | if (@alignOf(Item) == alignment) { |
| 19 | return MemoryPoolExtra(Item, .{}); |
| 20 | } else { |
| 21 | return MemoryPoolExtra(Item, .{ .alignment = alignment }); |
| 22 | } |
| 23 | } |
| 24 | |
| 25 | pub const Options = struct { |
| 26 | /// The alignment of the memory pool items. Use `null` for natural alignment. |
| 27 | alignment: ?u29 = null, |
| 28 | |
| 29 | /// If `true`, the memory pool can allocate additional items after a initial setup. |
| 30 | /// If `false`, the memory pool will not allocate further after a call to `initPreheated`. |
| 31 | growable: bool = true, |
| 32 | }; |
| 33 | |
| 34 | /// A memory pool that can allocate objects of a single type very quickly. |
| 35 | /// Use this when you need to allocate a lot of objects of the same type, |
| 36 | /// because It outperforms general purpose allocators. |
| 37 | pub fn MemoryPoolExtra(comptime Item: type, comptime pool_options: Options) type { |
| 38 | return struct { |
| 39 | const Pool = @This(); |
| 40 | |
| 41 | /// Size of the memory pool items. This is not necessarily the same |
| 42 | /// as `@sizeOf(Item)` as the pool also uses the items for internal means. |
| 43 | pub const item_size = std.math.max(@sizeOf(Node), @sizeOf(Item)); |
| 44 | |
| 45 | /// Alignment of the memory pool items. This is not necessarily the same |
| 46 | /// as `@alignOf(Item)` as the pool also uses the items for internal means. |
| 47 | pub const item_alignment = std.math.max(@alignOf(Node), pool_options.alignment orelse 0); |
| 48 | |
| 49 | const Node = struct { |
| 50 | next: ?*@This(), |
| 51 | }; |
| 52 | const NodePtr = *align(item_alignment) Node; |
| 53 | const ItemPtr = *align(item_alignment) Item; |
| 54 | |
| 55 | arena: std.heap.ArenaAllocator, |
| 56 | free_list: ?NodePtr = null, |
| 57 | |
| 58 | /// Creates a new memory pool. |
| 59 | pub fn init(allocator: std.mem.Allocator) Pool { |
| 60 | return .{ .arena = std.heap.ArenaAllocator.init(allocator) }; |
| 61 | } |
| 62 | |
| 63 | /// Creates a new memory pool and pre-allocates `initial_size` items. |
| 64 | /// This allows the up to `initial_size` active allocations before a |
| 65 | /// `OutOfMemory` error happens when calling `create()`. |
| 66 | pub fn initPreheated(allocator: std.mem.Allocator, initial_size: usize) MemoryPoolError!Pool { |
| 67 | var pool = init(allocator); |
| 68 | errdefer pool.deinit(); |
| 69 | |
| 70 | var i: usize = 0; |
| 71 | while (i < initial_size) : (i += 1) { |
| 72 | const raw_mem = try pool.allocNew(); |
| 73 | const free_node = @ptrCast(NodePtr, raw_mem); |
| 74 | free_node.* = Node{ |
| 75 | .next = pool.free_list, |
| 76 | }; |
| 77 | pool.free_list = free_node; |
| 78 | } |
| 79 | |
| 80 | return pool; |
| 81 | } |
| 82 | |
| 83 | /// Destroys the memory pool and frees all allocated memory. |
| 84 | pub fn deinit(pool: *Pool) void { |
| 85 | pool.arena.deinit(); |
| 86 | pool.* = undefined; |
| 87 | } |
| 88 | |
| 89 | /// Resets the memory pool and destroys all allocated items. |
| 90 | /// This can be used to batch-destroy all objects without invalidating the memory pool. |
| 91 | pub fn reset(pool: *Pool) void { |
| 92 | // TODO: Potentially store all allocated objects in a list as well, allowing to |
| 93 | // just move them into the free list instead of actually releasing the memory. |
| 94 | const allocator = pool.arena.child_allocator; |
| 95 | |
| 96 | // TODO: Replace with "pool.arena.reset()" when implemented. |
| 97 | pool.arena.deinit(); |
| 98 | pool.arena = std.heap.ArenaAllocator.init(allocator); |
| 99 | |
| 100 | pool.free_list = null; |
| 101 | } |
| 102 | |
| 103 | /// Creates a new item and adds it to the memory pool. |
| 104 | pub fn create(pool: *Pool) !ItemPtr { |
| 105 | const node = if (pool.free_list) |item| blk: { |
| 106 | pool.free_list = item.next; |
| 107 | break :blk item; |
| 108 | } else if (pool_options.growable) |
| 109 | @ptrCast(NodePtr, try pool.allocNew()) |
| 110 | else |
| 111 | return error.OutOfMemory; |
| 112 | |
| 113 | const ptr = @ptrCast(ItemPtr, node); |
| 114 | ptr.* = undefined; |
| 115 | return ptr; |
| 116 | } |
| 117 | |
| 118 | /// Destroys a previously created item. |
| 119 | /// Only pass items to `ptr` that were previously created with `create()` of the same memory pool! |
| 120 | pub fn destroy(pool: *Pool, ptr: ItemPtr) void { |
| 121 | ptr.* = undefined; |
| 122 | |
| 123 | const node = @ptrCast(NodePtr, ptr); |
| 124 | node.* = Node{ |
| 125 | .next = pool.free_list, |
| 126 | }; |
| 127 | pool.free_list = node; |
| 128 | } |
| 129 | |
| 130 | fn allocNew(pool: *Pool) MemoryPoolError!*align(item_alignment) [item_size]u8 { |
| 131 | const mem = try pool.arena.allocator().alignedAlloc(u8, item_alignment, item_size); |
| 132 | return mem[0..item_size]; // coerce slice to array pointer |
| 133 | } |
| 134 | }; |
| 135 | } |
| 136 | |
| 137 | test "memory pool: basic" { |
| 138 | var pool = MemoryPool(u32).init(std.testing.allocator); |
| 139 | defer pool.deinit(); |
| 140 | |
| 141 | const p1 = try pool.create(); |
| 142 | const p2 = try pool.create(); |
| 143 | const p3 = try pool.create(); |
| 144 | |
| 145 | // Assert uniqueness |
| 146 | try std.testing.expect(p1 != p2); |
| 147 | try std.testing.expect(p1 != p3); |
| 148 | try std.testing.expect(p2 != p3); |
| 149 | |
| 150 | pool.destroy(p2); |
| 151 | const p4 = try pool.create(); |
| 152 | |
| 153 | // Assert memory resuse |
| 154 | try std.testing.expect(p2 == p4); |
| 155 | } |
| 156 | |
| 157 | test "memory pool: preheating (success)" { |
| 158 | var pool = try MemoryPool(u32).initPreheated(std.testing.allocator, 4); |
| 159 | defer pool.deinit(); |
| 160 | |
| 161 | _ = try pool.create(); |
| 162 | _ = try pool.create(); |
| 163 | _ = try pool.create(); |
| 164 | } |
| 165 | |
| 166 | test "memory pool: preheating (failure)" { |
| 167 | var failer = std.testing.FailingAllocator.init(std.testing.allocator, 0); |
| 168 | try std.testing.expectError(error.OutOfMemory, MemoryPool(u32).initPreheated(failer.allocator(), 5)); |
| 169 | } |
| 170 | |
| 171 | test "memory pool: growable" { |
| 172 | var pool = try MemoryPoolExtra(u32, .{ .growable = false }).initPreheated(std.testing.allocator, 4); |
| 173 | defer pool.deinit(); |
| 174 | |
| 175 | _ = try pool.create(); |
| 176 | _ = try pool.create(); |
| 177 | _ = try pool.create(); |
| 178 | _ = try pool.create(); |
| 179 | |
| 180 | try std.testing.expectError(error.OutOfMemory, pool.create()); |
| 181 | } |
Andrew Kelley <andrew@ziglang.org> 2023-01-03 14:53:54-05:00
GitHub <noreply@github.com> 2023-01-03 14:53:54-05:00
Signed by PGP key 4AEE18F83AFDEB23
lib