1const std = @import("../std.zig");
2const Allocator = std.mem.Allocator;
3const Alignment = std.mem.Alignment;
4const MemoryPool = std.heap.MemoryPool;
5
6/// A memory pool that can allocate objects of a single type very quickly.
7/// Use this when you need to allocate a lot of objects of the same type,
8/// because it outperforms general purpose allocators.
9/// Allocated items are aligned to `alignment`-byte addresses or `@alignOf(Item)`
10/// if `alignment` is `null`.
11/// Functions that potentially allocate memory accept an `Allocator` parameter.
12pub fn Aligned(comptime Item: type, comptime alignment: Alignment) type {
13 return Extra(Item, .{ .alignment = alignment });
14}
15
16pub const Options = struct {
17 /// The alignment of the memory pool items. Use `null` for natural alignment.
18 alignment: ?Alignment = null,
19
20 /// If `true`, the memory pool can allocate additional items after a initial setup.
21 /// If `false`, the memory pool will not allocate further after a call to `initPreheated`.
22 growable: bool = true,
23};
24
25/// A memory pool that can allocate objects of a single type very quickly.
26/// Use this when you need to allocate a lot of objects of the same type,
27/// because it outperforms general purpose allocators.
28/// Functions that potentially allocate memory accept an `Allocator` parameter.
29pub fn Extra(comptime Item: type, comptime pool_options: Options) type {
30 if (pool_options.alignment) |a| {
31 if (a.compare(.eq, .of(Item))) {
32 var new_options = pool_options;
33 new_options.alignment = null;
34 return Extra(Item, new_options);
35 }
36 }
37 return struct {
38 const Pool = @This();
39
40 arena_state: std.heap.ArenaAllocator.State,
41 free_list: std.SinglyLinkedList,
42
43 /// Size of the memory pool items. This is not necessarily the same
44 /// as `@sizeOf(Item)` as the pool also uses the items for internal means.
45 pub const item_size = @max(@sizeOf(Node), @sizeOf(Item));
46
47 /// Alignment of the memory pool items. This is not necessarily the same
48 /// as `@alignOf(Item)` as the pool also uses the items for internal means.
49 pub const item_alignment: Alignment = .max(pool_options.alignment orelse .of(Item), .of(Node));
50
51 const Node = std.SinglyLinkedList.Node;
52 const ItemPtr = *align(item_alignment.toByteUnits()) Item;
53
54 /// A MemoryPool containing no elements.
55 pub const empty: Pool = .{
56 .arena_state = .{},
57 .free_list = .{},
58 };
59
60 /// Creates a new memory pool and pre-allocates `num` items.
61 /// This allows up to `num` active allocations before an
62 /// `OutOfMemory` error might happen when calling `create()`.
63 pub fn initCapacity(allocator: Allocator, num: usize) Allocator.Error!Pool {
64 var pool: Pool = .empty;
65 errdefer pool.deinit(allocator);
66 try pool.addCapacity(allocator, num);
67 return pool;
68 }
69
70 /// Destroys the memory pool and frees all allocated memory.
71 pub fn deinit(pool: *Pool, allocator: Allocator) void {
72 pool.arena_state.promote(allocator).deinit();
73 pool.* = undefined;
74 }
75
76 /// Pre-allocates `num` items and adds them to the memory pool.
77 /// This allows at least `num` active allocations before an
78 /// `OutOfMemory` error might happen when calling `create()`.
79 pub fn addCapacity(pool: *Pool, allocator: Allocator, num: usize) Allocator.Error!void {
80 var i: usize = 0;
81 while (i < num) : (i += 1) {
82 const memory = try pool.allocNew(allocator);
83 pool.free_list.prepend(@ptrCast(memory));
84 }
85 }
86
87 pub const ResetMode = std.heap.ArenaAllocator.ResetMode;
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 ///
92 /// The function will return whether the reset operation was successful or not.
93 /// If the reallocation failed `false` is returned. The pool will still be fully
94 /// functional in that case, all memory is released. Future allocations just might
95 /// be slower.
96 ///
97 /// NOTE: If `mode` is `free_all`, the function will always return `true`.
98 pub fn reset(pool: *Pool, allocator: Allocator, mode: ResetMode) bool {
99 // TODO: Potentially store all allocated objects in a list as well, allowing to
100 // just move them into the free list instead of actually releasing the memory.
101
102 var arena = pool.arena_state.promote(allocator);
103 defer pool.arena_state = arena.state;
104
105 const reset_successful = arena.reset(mode);
106 pool.free_list = .{};
107
108 return reset_successful;
109 }
110
111 /// Creates a new item and adds it to the memory pool.
112 /// `allocator` may be `undefined` if pool is not `growable`.
113 pub fn create(pool: *Pool, allocator: Allocator) Allocator.Error!ItemPtr {
114 const ptr: ItemPtr = if (pool.free_list.popFirst()) |node|
115 @ptrCast(@alignCast(node))
116 else if (pool_options.growable)
117 @ptrCast(try pool.allocNew(allocator))
118 else
119 return error.OutOfMemory;
120
121 ptr.* = undefined;
122 return ptr;
123 }
124
125 /// Destroys a previously created item.
126 /// Only pass items to `ptr` that were previously created with `create()` of the same memory pool!
127 pub fn destroy(pool: *Pool, ptr: ItemPtr) void {
128 ptr.* = undefined;
129 pool.free_list.prepend(@ptrCast(ptr));
130 }
131
132 fn allocNew(pool: *Pool, allocator: Allocator) Allocator.Error!*align(item_alignment.toByteUnits()) [item_size]u8 {
133 var arena = pool.arena_state.promote(allocator);
134 defer pool.arena_state = arena.state;
135 const memory = try arena.allocator().alignedAlloc(u8, item_alignment, item_size);
136 return memory[0..item_size];
137 }
138 };
139}
140
141test "basic" {
142 const a = std.testing.allocator;
143
144 var pool: MemoryPool(u32) = .empty;
145 defer pool.deinit(a);
146
147 const p1 = try pool.create(a);
148 const p2 = try pool.create(a);
149 const p3 = try pool.create(a);
150
151 // Assert uniqueness
152 try std.testing.expect(p1 != p2);
153 try std.testing.expect(p1 != p3);
154 try std.testing.expect(p2 != p3);
155
156 pool.destroy(p2);
157 const p4 = try pool.create(a);
158
159 // Assert memory reuse
160 try std.testing.expect(p2 == p4);
161}
162
163test "initCapacity (success)" {
164 const a = std.testing.allocator;
165
166 var pool: MemoryPool(u32) = try .initCapacity(a, 4);
167 defer pool.deinit(a);
168
169 _ = try pool.create(a);
170 _ = try pool.create(a);
171 _ = try pool.create(a);
172}
173
174test "initCapacity (failure)" {
175 const failer = std.testing.failing_allocator;
176 try std.testing.expectError(error.OutOfMemory, MemoryPool(u32).initCapacity(failer, 5));
177}
178
179test "growable" {
180 const a = std.testing.allocator;
181
182 var pool: Extra(u32, .{ .growable = false }) = try .initCapacity(a, 4);
183 defer pool.deinit(a);
184
185 _ = try pool.create(a);
186 _ = try pool.create(a);
187 _ = try pool.create(a);
188 _ = try pool.create(a);
189
190 try std.testing.expectError(error.OutOfMemory, pool.create(a));
191}
192
193test "greater than pointer default alignment" {
194 const Foo = struct {
195 data: u64 align(16),
196 };
197 const a = std.testing.allocator;
198
199 var pool: MemoryPool(Foo) = .empty;
200 defer pool.deinit(a);
201
202 const foo: *Foo = try pool.create(a);
203 pool.destroy(foo);
204}
205
206test "greater than pointer manual alignment" {
207 const Foo = struct {
208 data: u64,
209 };
210 const a = std.testing.allocator;
211
212 var pool: Aligned(Foo, .@"16") = .empty;
213 defer pool.deinit(a);
214
215 const foo: *align(16) Foo = try pool.create(a);
216 pool.destroy(foo);
217}