master
  1const std = @import("../std.zig");
  2const Allocator = std.mem.Allocator;
  3const Alignment = std.mem.Alignment;
  4const MemoryPool = std.heap.MemoryPool;
  5
  6/// Deprecated.
  7pub fn Managed(comptime Item: type) type {
  8    return ExtraManaged(Item, .{ .alignment = null });
  9}
 10
 11/// A memory pool that can allocate objects of a single type very quickly.
 12/// Use this when you need to allocate a lot of objects of the same type,
 13/// because it outperforms general purpose allocators.
 14/// Allocated items are aligned to `alignment`-byte addresses or `@alignOf(Item)`
 15/// if `alignment` is `null`.
 16/// Functions that potentially allocate memory accept an `Allocator` parameter.
 17pub fn Aligned(comptime Item: type, comptime alignment: Alignment) type {
 18    return Extra(Item, .{ .alignment = alignment });
 19}
 20
 21/// Deprecated.
 22pub fn AlignedManaged(comptime Item: type, comptime alignment: Alignment) type {
 23    return ExtraManaged(Item, .{ .alignment = alignment });
 24}
 25
 26pub const Options = struct {
 27    /// The alignment of the memory pool items. Use `null` for natural alignment.
 28    alignment: ?Alignment = null,
 29
 30    /// If `true`, the memory pool can allocate additional items after a initial setup.
 31    /// If `false`, the memory pool will not allocate further after a call to `initPreheated`.
 32    growable: bool = true,
 33};
 34
 35/// A memory pool that can allocate objects of a single type very quickly.
 36/// Use this when you need to allocate a lot of objects of the same type,
 37/// because it outperforms general purpose allocators.
 38/// Functions that potentially allocate memory accept an `Allocator` parameter.
 39pub fn Extra(comptime Item: type, comptime pool_options: Options) type {
 40    if (pool_options.alignment) |a| {
 41        if (a.compare(.eq, .of(Item))) {
 42            var new_options = pool_options;
 43            new_options.alignment = null;
 44            return Extra(Item, new_options);
 45        }
 46    }
 47    return struct {
 48        const Pool = @This();
 49
 50        arena_state: std.heap.ArenaAllocator.State,
 51        free_list: std.SinglyLinkedList,
 52
 53        /// Size of the memory pool items. This is not necessarily the same
 54        /// as `@sizeOf(Item)` as the pool also uses the items for internal means.
 55        pub const item_size = @max(@sizeOf(Node), @sizeOf(Item));
 56
 57        /// Alignment of the memory pool items. This is not necessarily the same
 58        /// as `@alignOf(Item)` as the pool also uses the items for internal means.
 59        pub const item_alignment: Alignment = .max(pool_options.alignment orelse .of(Item), .of(Node));
 60
 61        const Node = std.SinglyLinkedList.Node;
 62        const ItemPtr = *align(item_alignment.toByteUnits()) Item;
 63
 64        /// A MemoryPool containing no elements.
 65        pub const empty: Pool = .{
 66            .arena_state = .{},
 67            .free_list = .{},
 68        };
 69
 70        /// Creates a new memory pool and pre-allocates `num` items.
 71        /// This allows up to `num` active allocations before an
 72        /// `OutOfMemory` error might happen when calling `create()`.
 73        pub fn initCapacity(allocator: Allocator, num: usize) Allocator.Error!Pool {
 74            var pool: Pool = .empty;
 75            errdefer pool.deinit(allocator);
 76            try pool.addCapacity(allocator, num);
 77            return pool;
 78        }
 79
 80        /// Destroys the memory pool and frees all allocated memory.
 81        pub fn deinit(pool: *Pool, allocator: Allocator) void {
 82            pool.arena_state.promote(allocator).deinit();
 83            pool.* = undefined;
 84        }
 85
 86        pub fn toManaged(pool: Pool, allocator: Allocator) ExtraManaged(Item, pool_options) {
 87            return .{
 88                .allocator = allocator,
 89                .unmanaged = pool,
 90            };
 91        }
 92
 93        /// Pre-allocates `num` items and adds them to the memory pool.
 94        /// This allows at least `num` active allocations before an
 95        /// `OutOfMemory` error might happen when calling `create()`.
 96        pub fn addCapacity(pool: *Pool, allocator: Allocator, num: usize) Allocator.Error!void {
 97            var i: usize = 0;
 98            while (i < num) : (i += 1) {
 99                const memory = try pool.allocNew(allocator);
100                pool.free_list.prepend(@ptrCast(memory));
101            }
102        }
103
104        pub const ResetMode = std.heap.ArenaAllocator.ResetMode;
105
106        /// Resets the memory pool and destroys all allocated items.
107        /// This can be used to batch-destroy all objects without invalidating the memory pool.
108        ///
109        /// The function will return whether the reset operation was successful or not.
110        /// If the reallocation  failed `false` is returned. The pool will still be fully
111        /// functional in that case, all memory is released. Future allocations just might
112        /// be slower.
113        ///
114        /// NOTE: If `mode` is `free_all`, the function will always return `true`.
115        pub fn reset(pool: *Pool, allocator: Allocator, mode: ResetMode) bool {
116            // TODO: Potentially store all allocated objects in a list as well, allowing to
117            //       just move them into the free list instead of actually releasing the memory.
118
119            var arena = pool.arena_state.promote(allocator);
120            defer pool.arena_state = arena.state;
121
122            const reset_successful = arena.reset(mode);
123            pool.free_list = .{};
124
125            return reset_successful;
126        }
127
128        /// Creates a new item and adds it to the memory pool.
129        /// `allocator` may be `undefined` if pool is not `growable`.
130        pub fn create(pool: *Pool, allocator: Allocator) Allocator.Error!ItemPtr {
131            const ptr: ItemPtr = if (pool.free_list.popFirst()) |node|
132                @ptrCast(@alignCast(node))
133            else if (pool_options.growable)
134                @ptrCast(try pool.allocNew(allocator))
135            else
136                return error.OutOfMemory;
137
138            ptr.* = undefined;
139            return ptr;
140        }
141
142        /// Destroys a previously created item.
143        /// Only pass items to `ptr` that were previously created with `create()` of the same memory pool!
144        pub fn destroy(pool: *Pool, ptr: ItemPtr) void {
145            ptr.* = undefined;
146            pool.free_list.prepend(@ptrCast(ptr));
147        }
148
149        fn allocNew(pool: *Pool, allocator: Allocator) Allocator.Error!*align(item_alignment.toByteUnits()) [item_size]u8 {
150            var arena = pool.arena_state.promote(allocator);
151            defer pool.arena_state = arena.state;
152            const memory = try arena.allocator().alignedAlloc(u8, item_alignment, item_size);
153            return memory[0..item_size];
154        }
155    };
156}
157
158/// Deprecated.
159pub fn ExtraManaged(comptime Item: type, comptime pool_options: Options) type {
160    if (pool_options.alignment) |a| {
161        if (a.compare(.eq, .of(Item))) {
162            var new_options = pool_options;
163            new_options.alignment = null;
164            return ExtraManaged(Item, new_options);
165        }
166    }
167    return struct {
168        const Pool = @This();
169
170        allocator: Allocator,
171        unmanaged: Unmanaged,
172
173        pub const Unmanaged = Extra(Item, pool_options);
174        pub const item_size = Unmanaged.item_size;
175        pub const item_alignment = Unmanaged.item_alignment;
176
177        const ItemPtr = Unmanaged.ItemPtr;
178
179        /// Creates a new memory pool.
180        pub fn init(allocator: Allocator) Pool {
181            return Unmanaged.empty.toManaged(allocator);
182        }
183
184        /// Creates a new memory pool and pre-allocates `num` items.
185        /// This allows up to `num` active allocations before an
186        /// `OutOfMemory` error might happen when calling `create()`.
187        pub fn initCapacity(allocator: Allocator, num: usize) Allocator.Error!Pool {
188            return (try Unmanaged.initCapacity(allocator, num)).toManaged(allocator);
189        }
190
191        /// Destroys the memory pool and frees all allocated memory.
192        pub fn deinit(pool: *Pool) void {
193            pool.unmanaged.deinit(pool.allocator);
194            pool.* = undefined;
195        }
196
197        /// Pre-allocates `num` items and adds them to the memory pool.
198        /// This allows at least `num` active allocations before an
199        /// `OutOfMemory` error might happen when calling `create()`.
200        pub fn addCapacity(pool: *Pool, num: usize) Allocator.Error!void {
201            return pool.unmanaged.addCapacity(pool.allocator, num);
202        }
203
204        pub const ResetMode = Unmanaged.ResetMode;
205
206        /// Resets the memory pool and destroys all allocated items.
207        /// This can be used to batch-destroy all objects without invalidating the memory pool.
208        ///
209        /// The function will return whether the reset operation was successful or not.
210        /// If the reallocation  failed `false` is returned. The pool will still be fully
211        /// functional in that case, all memory is released. Future allocations just might
212        /// be slower.
213        ///
214        /// NOTE: If `mode` is `free_all`, the function will always return `true`.
215        pub fn reset(pool: *Pool, mode: ResetMode) bool {
216            return pool.unmanaged.reset(pool.allocator, mode);
217        }
218
219        /// Creates a new item and adds it to the memory pool.
220        pub fn create(pool: *Pool) Allocator.Error!ItemPtr {
221            return pool.unmanaged.create(pool.allocator);
222        }
223
224        /// Destroys a previously created item.
225        /// Only pass items to `ptr` that were previously created with `create()` of the same memory pool!
226        pub fn destroy(pool: *Pool, ptr: ItemPtr) void {
227            return pool.unmanaged.destroy(ptr);
228        }
229
230        fn allocNew(pool: *Pool) Allocator.Error!*align(item_alignment) [item_size]u8 {
231            return pool.unmanaged.allocNew(pool.allocator);
232        }
233    };
234}
235
236test "basic" {
237    const a = std.testing.allocator;
238
239    {
240        var pool: MemoryPool(u32) = .empty;
241        defer pool.deinit(a);
242
243        const p1 = try pool.create(a);
244        const p2 = try pool.create(a);
245        const p3 = try pool.create(a);
246
247        // Assert uniqueness
248        try std.testing.expect(p1 != p2);
249        try std.testing.expect(p1 != p3);
250        try std.testing.expect(p2 != p3);
251
252        pool.destroy(p2);
253        const p4 = try pool.create(a);
254
255        // Assert memory reuse
256        try std.testing.expect(p2 == p4);
257    }
258
259    {
260        var pool: Managed(u32) = .init(std.testing.allocator);
261        defer pool.deinit();
262
263        const p1 = try pool.create();
264        const p2 = try pool.create();
265        const p3 = try pool.create();
266
267        // Assert uniqueness
268        try std.testing.expect(p1 != p2);
269        try std.testing.expect(p1 != p3);
270        try std.testing.expect(p2 != p3);
271
272        pool.destroy(p2);
273        const p4 = try pool.create();
274
275        // Assert memory reuse
276        try std.testing.expect(p2 == p4);
277    }
278}
279
280test "initCapacity (success)" {
281    const a = std.testing.allocator;
282
283    {
284        var pool: MemoryPool(u32) = try .initCapacity(a, 4);
285        defer pool.deinit(a);
286
287        _ = try pool.create(a);
288        _ = try pool.create(a);
289        _ = try pool.create(a);
290    }
291
292    {
293        var pool: Managed(u32) = try .initCapacity(a, 4);
294        defer pool.deinit();
295
296        _ = try pool.create();
297        _ = try pool.create();
298        _ = try pool.create();
299    }
300}
301
302test "initCapacity (failure)" {
303    const failer = std.testing.failing_allocator;
304    try std.testing.expectError(error.OutOfMemory, MemoryPool(u32).initCapacity(failer, 5));
305    try std.testing.expectError(error.OutOfMemory, Managed(u32).initCapacity(failer, 5));
306}
307
308test "growable" {
309    const a = std.testing.allocator;
310
311    {
312        var pool: Extra(u32, .{ .growable = false }) = try .initCapacity(a, 4);
313        defer pool.deinit(a);
314
315        _ = try pool.create(a);
316        _ = try pool.create(a);
317        _ = try pool.create(a);
318        _ = try pool.create(a);
319
320        try std.testing.expectError(error.OutOfMemory, pool.create(a));
321    }
322
323    {
324        var pool: ExtraManaged(u32, .{ .growable = false }) = try .initCapacity(a, 4);
325        defer pool.deinit();
326
327        _ = try pool.create();
328        _ = try pool.create();
329        _ = try pool.create();
330        _ = try pool.create();
331
332        try std.testing.expectError(error.OutOfMemory, pool.create());
333    }
334}
335
336test "greater than pointer default alignment" {
337    const Foo = struct {
338        data: u64 align(16),
339    };
340    const a = std.testing.allocator;
341
342    {
343        var pool: MemoryPool(Foo) = .empty;
344        defer pool.deinit(a);
345
346        const foo: *Foo = try pool.create(a);
347        pool.destroy(foo);
348    }
349
350    {
351        var pool: Managed(Foo) = .init(a);
352        defer pool.deinit();
353
354        const foo: *Foo = try pool.create();
355        pool.destroy(foo);
356    }
357}
358
359test "greater than pointer manual alignment" {
360    const Foo = struct {
361        data: u64,
362    };
363    const a = std.testing.allocator;
364
365    {
366        var pool: Aligned(Foo, .@"16") = .empty;
367        defer pool.deinit(a);
368
369        const foo: *align(16) Foo = try pool.create(a);
370        pool.destroy(foo);
371    }
372
373    {
374        var pool: AlignedManaged(Foo, .@"16") = .init(a);
375        defer pool.deinit();
376
377        const foo: *align(16) Foo = try pool.create();
378        pool.destroy(foo);
379    }
380}