1//! This allocator takes an existing allocator, wraps it, and provides an interface where
2//! you can allocate and then free it all together. Calls to free an individual item only
3//! free the item if it was the most recent allocation, otherwise calls to free do
4//! nothing.
5//!
6//! The `Allocator` implementation provided is threadsafe, given that `child_allocator`
7//! is threadsafe as well.
8const ArenaAllocator = @This();
9
10child_allocator: Allocator,
11state: State,
12
13/// Inner state of ArenaAllocator. Can be stored rather than the entire ArenaAllocator
14/// as a memory-saving optimization.
15///
16/// Default initialization of this struct is deprecated; use `init` instead.
17pub const State = struct {
18 used_list: ?*Node = null,
19 free_list: ?*Node = null,
20
21 pub const init: State = .{
22 .used_list = null,
23 .free_list = null,
24 };
25
26 pub fn promote(state: State, child_allocator: Allocator) ArenaAllocator {
27 return .{
28 .child_allocator = child_allocator,
29 .state = state,
30 };
31 }
32};
33
34pub fn allocator(arena: *ArenaAllocator) Allocator {
35 return .{
36 .ptr = arena,
37 .vtable = &.{
38 .alloc = alloc,
39 .resize = resize,
40 .remap = remap,
41 .free = free,
42 },
43 };
44}
45
46pub fn init(child_allocator: Allocator) ArenaAllocator {
47 return State.init.promote(child_allocator);
48}
49
50/// Not threadsafe.
51pub fn deinit(arena: ArenaAllocator) void {
52 // NOTE: When changing this, make sure `reset()` is adjusted accordingly!
53
54 for ([_]?*Node{ arena.state.used_list, arena.state.free_list }) |first_node| {
55 var it = first_node;
56 while (it) |node| {
57 // this has to occur before the free because the free frees node
58 it = node.next;
59 arena.child_allocator.rawFree(node.allocatedSliceUnsafe(), .of(Node), @returnAddress());
60 }
61 }
62}
63
64/// Queries the current memory use of this arena.
65/// This will **not** include the storage required for internal keeping.
66///
67/// Not threadsafe.
68pub fn queryCapacity(arena: ArenaAllocator) usize {
69 var capacity: usize = 0;
70 for ([_]?*Node{ arena.state.used_list, arena.state.free_list }) |first_node| {
71 capacity += countListCapacity(first_node);
72 }
73 return capacity;
74}
75fn countListCapacity(first_node: ?*Node) usize {
76 var capacity: usize = 0;
77 var it = first_node;
78 while (it) |node| : (it = node.next) {
79 // Compute the actually allocated size excluding the
80 // linked list node.
81 capacity += node.size.toInt() - @sizeOf(Node);
82 }
83 return capacity;
84}
85
86pub const ResetMode = union(enum) {
87 /// Releases all allocated memory in the arena.
88 free_all,
89 /// This will pre-heat the arena for future allocations by allocating a
90 /// large enough buffer for all previously done allocations.
91 /// Preheating will speed up the allocation process by invoking the backing allocator
92 /// less often than before. If `reset()` is used in a loop, this means that after the
93 /// biggest operation, no memory allocations are performed anymore.
94 retain_capacity,
95 /// This is the same as `retain_capacity`, but the memory will be shrunk to
96 /// this value if it exceeds the limit.
97 retain_with_limit: usize,
98};
99/// Resets the arena allocator and frees all allocated memory.
100///
101/// `mode` defines how the currently allocated memory is handled.
102/// See the variant documentation for `ResetMode` for the effects of each mode.
103///
104/// The function will return whether the reset operation was successful or not.
105/// If the reallocation failed `false` is returned. The arena will still be fully
106/// functional in that case, all memory is released. Future allocations just might
107/// be slower.
108///
109/// Not threadsafe.
110///
111/// NOTE: If `mode` is `free_all`, the function will always return `true`.
112pub fn reset(arena: *ArenaAllocator, mode: ResetMode) bool {
113 // Some words on the implementation:
114 // The reset function can be implemented with two basic approaches:
115 // - Counting how much bytes were allocated since the last reset, and storing that
116 // information in State. This will make reset fast and alloc only a teeny tiny bit
117 // slower.
118 // - Counting how much bytes were allocated by iterating the chunk linked list. This
119 // will make reset slower, but alloc() keeps the same speed when reset() as if reset()
120 // would not exist.
121 //
122 // The second variant was chosen for implementation, as with more and more calls to reset(),
123 // the function will get faster and faster. At one point, the complexity of the function
124 // will drop to amortized O(1), as we're only ever having a single chunk that will not be
125 // reallocated, and we're not even touching the backing allocator anymore.
126 //
127 // Thus, only the first hand full of calls to reset() will actually need to iterate the linked
128 // list, all future calls are just taking the first node, and only resetting the `end_index`
129 // value.
130
131 const limit: ?usize = switch (mode) {
132 .retain_capacity => null,
133 .retain_with_limit => |limit| limit,
134 .free_all => 0,
135 };
136 if (limit == 0) {
137 // just reset when we don't have anything to reallocate
138 arena.deinit();
139 arena.state = .init;
140 return true;
141 }
142
143 const used_capacity = countListCapacity(arena.state.used_list);
144 const free_capacity = countListCapacity(arena.state.free_list);
145
146 const new_used_capacity = if (limit) |lim| @min(lim, used_capacity) else used_capacity;
147 const new_free_capacity = if (limit) |lim| @min(lim - new_used_capacity, free_capacity) else free_capacity;
148
149 var ok = true;
150
151 for (
152 [_]*?*Node{ &arena.state.used_list, &arena.state.free_list },
153 [_]usize{ new_used_capacity, new_free_capacity },
154 ) |first_node_ptr, new_capacity| {
155 // Free all nodes except for the last one
156 var it = first_node_ptr.*;
157 const node: *Node = while (it) |node| {
158 // this has to occur before the free because the free frees node
159 it = node.next;
160 if (it == null) break node;
161 arena.child_allocator.rawFree(node.allocatedSliceUnsafe(), .of(Node), @returnAddress());
162 } else {
163 continue;
164 };
165 const allocated_slice = node.allocatedSliceUnsafe();
166
167 // Align backwards to always stay below limit.
168 const new_size = mem.alignBackward(usize, @sizeOf(Node) + new_capacity, 2);
169
170 if (new_size == @sizeOf(Node)) {
171 arena.child_allocator.rawFree(allocated_slice, .of(Node), @returnAddress());
172 first_node_ptr.* = null;
173 continue;
174 }
175
176 node.end_index = 0;
177 first_node_ptr.* = node;
178
179 if (allocated_slice.len == new_size) {
180 // perfect, no need to invoke the child_allocator
181 continue;
182 }
183
184 if (arena.child_allocator.rawResize(allocated_slice, .of(Node), new_size, @returnAddress())) {
185 // successful resize
186 node.size = .fromInt(new_size);
187 } else {
188 // manual realloc
189 const new_ptr = arena.child_allocator.rawAlloc(new_size, .of(Node), @returnAddress()) orelse {
190 // we failed to preheat the arena properly, signal this to the user.
191 ok = false;
192 continue;
193 };
194 arena.child_allocator.rawFree(allocated_slice, .of(Node), @returnAddress());
195 const new_first_node: *Node = @ptrCast(@alignCast(new_ptr));
196 new_first_node.* = .{
197 .size = .fromInt(new_size),
198 .end_index = 0,
199 .next = null,
200 };
201 first_node_ptr.* = new_first_node;
202 }
203 }
204
205 return ok;
206}
207
208/// Concurrent accesses to node pointers generally have to have acquire/release
209/// semantics to guarantee that newly allocated notes are in a valid state when
210/// being inserted into a list. Exceptions are possible, e.g. a cmpxchg loop that
211/// never accesses the node returned on failure can use monotonic semantics on
212/// failure, but must still use release semantics on success to protect the node
213/// it's trying to push.
214const Node = struct {
215 /// Only meant to be accessed indirectly via the methods supplied by this type,
216 /// except if the node is owned by the thread accessing it.
217 /// Must always be an even number to accommodate `resize` bit.
218 size: Size,
219 /// Any increase of `end_index` has to use acquire semantics;
220 /// any decrease of `end_index` that invalidates (formerly) active allocations
221 /// has to use release semantics.
222 /// This guarantees that all accesses to memory that's about to be freed
223 /// happen-before the free is published.
224 /// Since `size` can only grow and never shrink, memory access depending on
225 /// any `end_index` <= any `size` can never be OOB.
226 end_index: usize,
227 /// This field should only be accessed if the node is owned by the thread
228 /// accessing it.
229 next: ?*Node,
230
231 const Size = packed struct(usize) {
232 resizing: bool,
233 _: @Int(.unsigned, @bitSizeOf(usize) - 1) = 0,
234
235 fn fromInt(int: usize) Size {
236 assert(int >= @sizeOf(Node));
237 const size: Size = @bitCast(int);
238 assert(!size.resizing);
239 return size;
240 }
241
242 fn toInt(size: Size) usize {
243 var int = size;
244 int.resizing = false;
245 return @bitCast(int);
246 }
247
248 comptime {
249 assert(Size{ .resizing = true } == @as(Size, @bitCast(@as(usize, 1))));
250 }
251 };
252
253 fn loadBuf(node: *Node) []u8 {
254 // `size` can only ever grow, so the buffer returned by this function is
255 // always valid memory.
256 const size = @atomicLoad(Size, &node.size, .monotonic);
257 return @as([*]u8, @ptrCast(node))[0..size.toInt()][@sizeOf(Node)..];
258 }
259
260 /// Returns allocated slice or `null` if node is already (being) resized.
261 fn beginResize(node: *Node) ?[]u8 {
262 const size = @atomicRmw(Size, &node.size, .Or, .{ .resizing = true }, .acquire); // syncs with release in `endResize`
263 if (size.resizing) return null;
264 return @as([*]u8, @ptrCast(node))[0..size.toInt()];
265 }
266
267 fn endResize(node: *Node, size: usize, prev_size: usize) void {
268 assert(size >= prev_size); // nodes must not shrink
269 assert(@atomicLoad(Size, &node.size, .unordered).toInt() == prev_size);
270 return @atomicStore(Size, &node.size, .fromInt(size), .release); // syncs with acquire in `beginResize`
271 }
272
273 /// Not threadsafe.
274 fn allocatedSliceUnsafe(node: *Node) []u8 {
275 return @as([*]u8, @ptrCast(node))[0..node.size.toInt()];
276 }
277};
278
279fn loadFirstNode(arena: *ArenaAllocator) ?*Node {
280 return @atomicLoad(?*Node, &arena.state.used_list, .acquire); // syncs with release in successful `tryPushNode`
281}
282
283const PushResult = union(enum) {
284 success,
285 failure: ?*Node,
286};
287fn tryPushNode(arena: *ArenaAllocator, node: *Node) PushResult {
288 assert(node != node.next);
289 if (@cmpxchgStrong( // strong because retrying means discarding a fitting node -> expensive
290 ?*Node,
291 &arena.state.used_list,
292 node.next,
293 node,
294 .release, // syncs with acquire in failure path or `loadFirstNode`
295 .acquire, // syncs with release in success path
296 )) |old_node| {
297 return .{ .failure = old_node };
298 } else {
299 return .success;
300 }
301}
302
303fn stealFreeList(arena: *ArenaAllocator) ?*Node {
304 // We don't need acq_rel here because we're always swapping in `null`, so
305 // there's no node we'd need to release.
306 return @atomicRmw(?*Node, &arena.state.free_list, .Xchg, null, .acquire); // syncs with release in `pushFreeList`
307}
308
309fn pushFreeList(arena: *ArenaAllocator, first: *Node, last: *Node) void {
310 assert(first != last.next);
311 assert(first != first.next);
312 assert(last != last.next);
313 while (@cmpxchgWeak(
314 ?*Node,
315 &arena.state.free_list,
316 last.next,
317 first,
318 .release, // syncs with acquire in `stealFreeList`
319 .monotonic, // we never access any fields of `old_free_list`, we only care about the pointer
320 )) |old_free_list| {
321 last.next = old_free_list;
322 }
323}
324
325fn alignedIndex(buf_ptr: [*]u8, end_index: usize, alignment: Alignment) usize {
326 // Wrapping arithmetic to avoid overflows since `end_index` isn't bounded by
327 // `size`. This is always ok since the max alignment in byte units is also
328 // the max value of `usize` so wrapped values are correctly aligned anyway.
329 return alignment.forward(@intFromPtr(buf_ptr) +% end_index) -% @intFromPtr(buf_ptr);
330}
331
332fn alloc(ctx: *anyopaque, n: usize, alignment: Alignment, ret_addr: usize) ?[*]u8 {
333 const arena: *ArenaAllocator = @ptrCast(@alignCast(ctx));
334 _ = ret_addr;
335
336 assert(n > 0);
337
338 var cur_first_node = arena.loadFirstNode();
339
340 var cur_new_node: ?*Node = null;
341 defer if (cur_new_node) |node| {
342 node.next = null; // optimize for empty free list
343 arena.pushFreeList(node, node);
344 };
345
346 retry: while (true) {
347 const first_node: ?*Node, const prev_size: usize = first_node: {
348 const node = cur_first_node orelse break :first_node .{ null, 0 };
349 const buf = node.loadBuf();
350
351 // To avoid using a CAS loop in the hot path we atomically increase
352 // `end_index` by a large enough amount to be able to always provide
353 // the required alignment within the reserved memory. To recover the
354 // space this potentially wastes we try to subtract the 'overshoot'
355 // with a single cmpxchg afterwards, which may fail.
356
357 const alignable = n + alignment.toByteUnits() - 1;
358 const end_index = @atomicRmw(usize, &node.end_index, .Add, alignable, .acquire); // acquire any memory that may have been freed
359 const aligned_index = alignedIndex(buf.ptr, end_index, alignment);
360 assert(end_index + alignable >= aligned_index + n);
361 if (end_index + alignable != aligned_index + n) {
362 _ = @cmpxchgStrong(
363 usize,
364 &node.end_index,
365 end_index + alignable,
366 aligned_index + n,
367 .monotonic, // no need to release alignment padding; there's no one accessing it!
368 .monotonic,
369 );
370 }
371
372 if (aligned_index + n > buf.len) break :first_node .{ node, buf.len };
373 return buf[aligned_index..][0..n].ptr;
374 };
375
376 resize: {
377 // Before attempting to get our hands on a new node, we try to resize
378 // the one we're currently holding. This is an exclusive operation;
379 // if another thread is already in this section we can never resize.
380
381 const node = first_node orelse break :resize;
382 const allocated_slice = node.beginResize() orelse break :resize;
383 var size = allocated_slice.len;
384 defer node.endResize(size, allocated_slice.len);
385
386 const buf = allocated_slice[@sizeOf(Node)..];
387 const end_index = @atomicLoad(usize, &node.end_index, .monotonic);
388 const aligned_index = alignedIndex(buf.ptr, end_index, alignment);
389 const new_size = mem.alignForward(usize, @sizeOf(Node) + aligned_index + n, 2);
390
391 if (new_size <= allocated_slice.len) {
392 // A `resize` or `free` call managed to sneak in and we need to
393 // guarantee that `size` is only ever increased; retry!
394 continue :retry;
395 }
396
397 if (arena.child_allocator.rawResize(allocated_slice, .of(Node), new_size, @returnAddress())) {
398 size = new_size;
399
400 // strong because a spurious failure could result in suboptimal
401 // usage of this node
402 if (null == @cmpxchgStrong(
403 usize,
404 &node.end_index,
405 end_index,
406 aligned_index + n,
407 .acquire, // acquire any memory that may have been freed
408 .monotonic,
409 )) {
410 const new_buf = allocated_slice.ptr[0..new_size][@sizeOf(Node)..];
411 return new_buf[aligned_index..][0..n].ptr;
412 }
413 }
414 }
415
416 // We need a new node! First, we search `free_list` for one that's big
417 // enough, if we don't find one there we fall back to allocating a new
418 // node with `child_allocator` (if we haven't already done that!).
419
420 from_free_list: {
421 // We 'steal' the entire free list to operate on it without other
422 // threads getting up into our business.
423 // This is a rather pragmatic approach, but since the free list isn't
424 // used very frequently it's fine performance-wise, even under load.
425 // Also this avoids the ABA problem; stealing the list with an atomic
426 // swap doesn't introduce any potentially stale `next` pointers.
427
428 const free_list = arena.stealFreeList() orelse break :from_free_list;
429
430 const first_free: *Node, const last_free: *Node, const node: *Node, const prev: ?*Node = find: {
431 var best_fit_prev: ?*Node = null;
432 var best_fit: ?*Node = null;
433 var best_fit_diff: usize = std.math.maxInt(usize);
434
435 var it_prev: ?*Node = null;
436 var it: ?*Node = free_list;
437 while (it) |node| : ({
438 it_prev = node;
439 it = node.next;
440 }) {
441 assert(!node.size.resizing);
442 const buf = node.allocatedSliceUnsafe()[@sizeOf(Node)..];
443 const aligned_index = alignedIndex(buf.ptr, 0, alignment);
444
445 const diff = aligned_index + n -| buf.len;
446 if (diff < best_fit_diff) {
447 best_fit_prev = it_prev;
448 best_fit = node;
449 best_fit_diff = diff;
450 }
451 }
452
453 break :find .{ free_list, it_prev.?, best_fit.?, best_fit_prev };
454 };
455
456 const aligned_index, const need_resize = aligned_index: {
457 const buf = node.allocatedSliceUnsafe()[@sizeOf(Node)..];
458 const aligned_index = alignedIndex(buf.ptr, 0, alignment);
459 break :aligned_index .{ aligned_index, aligned_index + n > buf.len };
460 };
461
462 if (need_resize) {
463 // Ideally we want to use all nodes in `free_list` eventually,
464 // so even if none fit we'll try to resize the one that was the
465 // closest to being large enough.
466 const new_size = mem.alignForward(usize, @sizeOf(Node) + aligned_index + n, 2);
467 if (arena.child_allocator.rawResize(node.allocatedSliceUnsafe(), .of(Node), new_size, @returnAddress())) {
468 node.size = .fromInt(new_size);
469 } else {
470 arena.pushFreeList(first_free, last_free);
471 break :from_free_list; // we couldn't find a fitting free node
472 }
473 }
474
475 const buf = node.allocatedSliceUnsafe()[@sizeOf(Node)..];
476 const old_next = node.next;
477
478 node.end_index = aligned_index + n;
479 node.next = first_node;
480
481 switch (arena.tryPushNode(node)) {
482 .success => {
483 // Finish removing node from free list.
484 if (prev) |p| p.next = old_next;
485
486 // Push remaining stolen free list back onto `arena.state.free_list`.
487 const new_first_free = if (node == first_free) old_next else first_free;
488 const new_last_free = if (node == last_free) prev else last_free;
489 if (new_first_free) |first| {
490 const last = new_last_free.?;
491 arena.pushFreeList(first, last);
492 }
493
494 return buf[aligned_index..][0..n].ptr;
495 },
496 .failure => |old_first_node| {
497 // restore free list to as we found it
498 node.next = old_next;
499 arena.pushFreeList(first_free, last_free);
500
501 cur_first_node = old_first_node;
502 continue :retry; // there's a new first node; retry!
503 },
504 }
505 }
506
507 const new_node: *Node = new_node: {
508 if (cur_new_node) |new_node| {
509 break :new_node new_node;
510 } else {
511 @branchHint(.cold);
512 }
513
514 const size: Node.Size = size: {
515 const min_size = @sizeOf(Node) + alignment.toByteUnits() + n;
516 const big_enough_size = prev_size + min_size + 16;
517 const size = mem.alignForward(usize, big_enough_size + big_enough_size / 2, 2);
518 break :size .fromInt(size);
519 };
520 const ptr = arena.child_allocator.rawAlloc(size.toInt(), .of(Node), @returnAddress()) orelse
521 return null;
522 const new_node: *Node = @ptrCast(@alignCast(ptr));
523 new_node.* = .{
524 .size = size,
525 .end_index = undefined, // set below
526 .next = undefined, // set below
527 };
528 cur_new_node = new_node;
529 break :new_node new_node;
530 };
531
532 const buf = new_node.allocatedSliceUnsafe()[@sizeOf(Node)..];
533 const aligned_index = alignedIndex(buf.ptr, 0, alignment);
534 assert(new_node.size.toInt() >= @sizeOf(Node) + aligned_index + n);
535
536 new_node.end_index = aligned_index + n;
537 new_node.next = first_node;
538
539 switch (arena.tryPushNode(new_node)) {
540 .success => {
541 cur_new_node = null;
542 return buf[aligned_index..][0..n].ptr;
543 },
544 .failure => |old_first_node| {
545 cur_first_node = old_first_node;
546 },
547 }
548 }
549}
550
551fn resize(ctx: *anyopaque, memory: []u8, alignment: Alignment, new_len: usize, ret_addr: usize) bool {
552 const arena: *ArenaAllocator = @ptrCast(@alignCast(ctx));
553 _ = alignment;
554 _ = ret_addr;
555
556 assert(memory.len > 0);
557 assert(new_len > 0);
558
559 const node = arena.loadFirstNode().?;
560 const buf_ptr = @as([*]u8, @ptrCast(node)) + @sizeOf(Node);
561
562 const cur_end_index = @atomicLoad(usize, &node.end_index, .monotonic);
563
564 if (buf_ptr + cur_end_index != memory.ptr + memory.len) {
565 // It's not the most recent allocation, so it cannot be expanded,
566 // but it's fine if they want to make it smaller.
567 return new_len <= memory.len;
568 }
569
570 if (new_len <= memory.len) {
571 const new_end_index = cur_end_index - (memory.len - new_len);
572 assert(buf_ptr + new_end_index == memory.ptr + new_len);
573
574 _ = @cmpxchgStrong(
575 usize,
576 &node.end_index,
577 cur_end_index,
578 new_end_index,
579 .release, // release freed memory
580 .monotonic,
581 );
582 return true; // Shrinking allocations should always succeed.
583 }
584
585 // Saturating arithmetic because `end_index` is not guaranteed to be `<= size`.
586 // The allocation we're trying to resize *could* belong to a different node!
587 if (node.loadBuf().len -| cur_end_index >= new_len - memory.len) {
588 const new_end_index = cur_end_index + (new_len - memory.len);
589 assert(buf_ptr + new_end_index == memory.ptr + new_len);
590
591 return null == @cmpxchgStrong(
592 usize,
593 &node.end_index,
594 cur_end_index,
595 new_end_index,
596 .acquire, // acquire any memory that may have been freed
597 .monotonic,
598 );
599 }
600
601 return false;
602}
603
604fn remap(ctx: *anyopaque, memory: []u8, alignment: Alignment, new_len: usize, ret_addr: usize) ?[*]u8 {
605 return if (resize(ctx, memory, alignment, new_len, ret_addr)) memory.ptr else null;
606}
607
608fn free(ctx: *anyopaque, memory: []u8, alignment: Alignment, ret_addr: usize) void {
609 const arena: *ArenaAllocator = @ptrCast(@alignCast(ctx));
610 _ = alignment;
611 _ = ret_addr;
612
613 assert(memory.len > 0);
614
615 const node = arena.loadFirstNode().?;
616 const buf_ptr = @as([*]u8, @ptrCast(node)) + @sizeOf(Node);
617
618 const cur_end_index = @atomicLoad(usize, &node.end_index, .monotonic);
619
620 if (buf_ptr + cur_end_index != memory.ptr + memory.len) {
621 // Not the most recent allocation; we cannot free it.
622 return;
623 }
624
625 const new_end_index = cur_end_index - memory.len;
626 assert(buf_ptr + new_end_index == memory.ptr);
627
628 _ = @cmpxchgStrong(
629 usize,
630 &node.end_index,
631 cur_end_index,
632 new_end_index,
633 .release, // release freed memory
634 .monotonic,
635 );
636}
637
638const std = @import("std");
639const assert = std.debug.assert;
640const mem = std.mem;
641const Allocator = std.mem.Allocator;
642const Alignment = std.mem.Alignment;
643
644test "reset with preheating" {
645 var arena_allocator = ArenaAllocator.init(std.testing.allocator);
646 defer arena_allocator.deinit();
647 // provides some variance in the allocated data
648 var rng_src = std.Random.DefaultPrng.init(std.testing.random_seed);
649 const random = rng_src.random();
650 var rounds: usize = 25;
651 while (rounds > 0) {
652 rounds -= 1;
653 _ = arena_allocator.reset(.retain_capacity);
654 var alloced_bytes: usize = 0;
655 const total_size: usize = random.intRangeAtMost(usize, 256, 16384);
656 while (alloced_bytes < total_size) {
657 const size = random.intRangeAtMost(usize, 16, 256);
658 const alignment: Alignment = .@"32";
659 const slice = try arena_allocator.allocator().alignedAlloc(u8, alignment, size);
660 try std.testing.expect(alignment.check(@intFromPtr(slice.ptr)));
661 try std.testing.expectEqual(size, slice.len);
662 alloced_bytes += slice.len;
663 }
664 }
665}
666
667test "reset while retaining a buffer" {
668 var arena_allocator = ArenaAllocator.init(std.testing.allocator);
669 defer arena_allocator.deinit();
670 const a = arena_allocator.allocator();
671
672 // Create two internal buffers
673 _ = try a.alloc(u8, 1);
674 try std.testing.expect(arena_allocator.state.used_list != null);
675 while (arena_allocator.state.used_list.?.next == null) {
676 _ = try a.alloc(u8, 1000);
677 }
678
679 // This retains the first allocated buffer
680 try std.testing.expect(arena_allocator.reset(.{ .retain_with_limit = 2 }));
681 try std.testing.expect(arena_allocator.state.used_list.?.next == null);
682 try std.testing.expectEqual(2, arena_allocator.queryCapacity());
683}
684
685test "fuzz multi threaded" {
686 @disableInstrumentation();
687 if (@import("builtin").single_threaded) return error.SkipZigTest;
688
689 const gpa = std.heap.smp_allocator;
690
691 var io_instance: std.Io.Threaded = .init(gpa, .{});
692 defer io_instance.deinit();
693
694 var arena_state: ArenaAllocator.State = .init;
695 // No need to deinit arena_state, all allocations are in `sample_buffer`!
696
697 const buffer_size = FuzzContext.max_alloc_count * FuzzContext.max_alloc_size;
698
699 const control_buffer = try gpa.alloc(u8, buffer_size);
700 defer gpa.free(control_buffer);
701 var control_instance: std.heap.FixedBufferAllocator = .init(control_buffer);
702
703 const sample_buffer = try gpa.alloc(u8, buffer_size);
704 defer gpa.free(sample_buffer);
705 var sample_instance: FuzzAllocator = .init(sample_buffer);
706
707 try std.testing.fuzz(FuzzContext.Init{
708 .threaded_instance = &io_instance,
709 .arena_state = &arena_state,
710 .control_instance = &control_instance,
711 .sample_instance = &sample_instance,
712 }, fuzzMultiThreaded, .{});
713}
714
715fn fuzzMultiThreaded(fuzz_init: FuzzContext.Init, smith: *std.testing.Smith) anyerror!void {
716 @disableInstrumentation();
717 const testing = std.testing;
718 const io = fuzz_init.threaded_instance.io();
719
720 fuzz_init.sample_instance.prepareFailures(smith);
721
722 const control_allocator = fuzz_init.control_instance.threadSafeAllocator();
723 const sample_child_allocator = fuzz_init.sample_instance.allocator();
724
725 var arena_instance = fuzz_init.arena_state.*.promote(sample_child_allocator);
726 defer fuzz_init.arena_state.* = arena_instance.state;
727
728 var ctx: FuzzContext = .init(
729 control_allocator,
730 arena_instance.allocator(),
731 );
732 defer ctx.deinit();
733
734 var group: std.Io.Group = .init;
735 defer group.cancel(io);
736
737 var n_allocs: usize = 0;
738 var n_actions: usize = 0;
739 while (!smith.eosWeightedSimple(99, 1) and n_actions < FuzzContext.max_action_count) {
740 errdefer comptime unreachable;
741
742 const weights: []const testing.Smith.Weight = if (n_allocs == FuzzContext.max_alloc_count)
743 &.{
744 .value(FuzzContext.Action, .resize, 1),
745 .value(FuzzContext.Action, .remap, 1),
746 .value(FuzzContext.Action, .free, 1),
747 }
748 else
749 &.{
750 .value(FuzzContext.Action, .resize, 1),
751 .value(FuzzContext.Action, .remap, 1),
752 .value(FuzzContext.Action, .free, 1),
753 .value(FuzzContext.Action, .alloc, 3),
754 };
755 switch (smith.valueWeighted(FuzzContext.Action, weights)) {
756 .alloc => {
757 const alloc_index = n_allocs;
758 n_allocs += 1;
759 ctx.allocs[alloc_index].common.len = .free;
760 group.concurrent(io, FuzzContext.doOneAlloc, .{
761 &ctx,
762 nextLen(smith),
763 smith.valueRangeAtMost(
764 Alignment,
765 .@"1",
766 .fromByteUnits(2 * std.heap.page_size_max),
767 ),
768 @fromBackingInt(@intCast(alloc_index)),
769 }) catch unreachable;
770 },
771 .resize => group.concurrent(io, FuzzContext.doOneResize, .{ &ctx, nextLen(smith) }) catch unreachable,
772 .remap => group.concurrent(io, FuzzContext.doOneRemap, .{ &ctx, nextLen(smith) }) catch unreachable,
773 .free => group.concurrent(io, FuzzContext.doOneFree, .{&ctx}) catch unreachable,
774 }
775 n_actions += 1;
776 }
777
778 try group.await(io);
779 try ctx.check(n_allocs);
780
781 // This also covers the `deinit` logic since `free_all` uses it internally.
782
783 const old_capacity = arena_instance.queryCapacity();
784 const reset_mode: ResetMode = switch (smith.value(@typeInfo(ResetMode).@"union".tag_type.?)) {
785 .free_all => .free_all,
786 .retain_capacity => .retain_capacity,
787 .retain_with_limit => .{ .retain_with_limit = smith.value(usize) },
788 };
789 const ok = arena_instance.reset(reset_mode);
790 const new_capacity = arena_instance.queryCapacity();
791 switch (reset_mode) {
792 .free_all => {
793 try testing.expect(ok);
794 try testing.expectEqual(0, new_capacity);
795 fuzz_init.sample_instance.reset();
796 },
797 .retain_with_limit => |limit| if (ok) try testing.expect(new_capacity <= limit),
798 .retain_capacity => if (ok) try testing.expectEqual(old_capacity, new_capacity),
799 }
800
801 fuzz_init.control_instance.reset();
802}
803fn nextLen(smith: *std.testing.Smith) @typeInfo(FuzzContext.Alloc.Len).@"enum".tag_type {
804 @disableInstrumentation();
805 const BackingInt = @typeInfo(FuzzContext.Alloc.Len).@"enum".tag_type;
806 return smith.valueRangeAtMost(BackingInt, 1, FuzzContext.max_alloc_size);
807}
808
809const FuzzContext = struct {
810 control_allocator: Allocator,
811 sample_allocator: Allocator,
812
813 last_alloc_index: Alloc.Index,
814 allocs: [max_alloc_count]Alloc,
815
816 const max_alloc_count = 64;
817 const max_action_count = 2 * max_alloc_count;
818
819 const max_alloc_size = 16 << 10;
820
821 const Alloc = struct {
822 control_ptr: [*]u8,
823 sample_ptr: [*]u8,
824 common: packed struct(usize) {
825 len: Len,
826 alignment: Alignment,
827 _: @Int(.unsigned, padding_bits) = 0,
828 },
829
830 const Len = enum(@Int(.unsigned, len_bits)) {
831 free = (1 << len_bits) - 1,
832 _,
833 };
834 const len_bits = @min(64, @bitSizeOf(usize)) - @bitSizeOf(Alignment);
835 const padding_bits = @bitSizeOf(usize) - (len_bits + @bitSizeOf(Alignment));
836
837 const Index = enum(usize) {
838 none = std.math.maxInt(usize),
839 _,
840 };
841 };
842
843 const Action = enum {
844 alloc,
845 resize,
846 remap,
847 free,
848 };
849
850 const Init = struct {
851 threaded_instance: *std.Io.Threaded,
852 arena_state: *ArenaAllocator.State,
853 control_instance: *std.heap.FixedBufferAllocator,
854 sample_instance: *FuzzAllocator,
855 };
856
857 fn init(
858 control_allocator: Allocator,
859 sample_allocator: Allocator,
860 ) FuzzContext {
861 @disableInstrumentation();
862 return .{
863 .control_allocator = control_allocator,
864 .sample_allocator = sample_allocator,
865 .last_alloc_index = .none,
866 .allocs = undefined,
867 };
868 }
869
870 fn deinit(ctx: *FuzzContext) void {
871 @disableInstrumentation();
872 ctx.* = undefined;
873 }
874
875 fn check(ctx: *const FuzzContext, n_allocs: usize) !void {
876 @disableInstrumentation();
877 for (ctx.allocs[0..n_allocs]) |allocation| {
878 const len: usize = switch (allocation.common.len) {
879 .free => continue,
880 _ => |len| @backingInt(len),
881 };
882 const control = allocation.control_ptr[0..len];
883 const sample = allocation.sample_ptr[0..len];
884 try std.testing.expectEqualSlices(u8, control, sample);
885 }
886 }
887
888 fn doOneAlloc(ctx: *FuzzContext, len: usize, alignment: Alignment, index: Alloc.Index) void {
889 @disableInstrumentation();
890 assert(ctx.allocs[@backingInt(index)].common.len == .free);
891
892 const control_ptr = ctx.control_allocator.rawAlloc(len, alignment, @returnAddress()) orelse
893 return;
894 const sample_ptr = ctx.sample_allocator.rawAlloc(len, alignment, @returnAddress()) orelse {
895 ctx.control_allocator.rawFree(control_ptr[0..len], alignment, @returnAddress());
896 return;
897 };
898
899 ctx.allocs[@backingInt(index)] = .{
900 .control_ptr = control_ptr,
901 .sample_ptr = sample_ptr,
902 .common = .{
903 .len = @fromBackingInt(@intCast(len)),
904 .alignment = alignment,
905 },
906 };
907
908 for (control_ptr[0..len], sample_ptr[0..len], 0..) |*control, *sample, i| {
909 control.* = @truncate(i);
910 sample.* = @truncate(i);
911 }
912
913 @atomicStore(Alloc.Index, &ctx.last_alloc_index, index, .release);
914 }
915 fn doOneResize(ctx: *FuzzContext, new_len: usize) void {
916 @disableInstrumentation();
917
918 const index = @atomicRmw(Alloc.Index, &ctx.last_alloc_index, .Xchg, .none, .acquire);
919 if (index == .none) return;
920
921 const allocation = &ctx.allocs[@backingInt(index)];
922 assert(allocation.common.len != .free);
923 const memory = allocation.sample_ptr[0..@backingInt(allocation.common.len)];
924 const alignment = allocation.common.alignment;
925
926 assert(alignment.check(@intFromPtr(allocation.control_ptr)));
927 assert(alignment.check(@intFromPtr(allocation.sample_ptr)));
928
929 // Since `resize` is fallible, we have to ensure that `control_allocator`
930 // is always successful by reserving the memory we need beforehand.
931 const new_control_ptr = ctx.control_allocator.rawAlloc(new_len, alignment, @returnAddress()) orelse
932 return;
933 if (ctx.sample_allocator.rawResize(memory, alignment, new_len, @returnAddress())) {
934 const old_control = allocation.control_ptr[0..memory.len];
935 const overlap = @min(memory.len, new_len);
936 @memcpy(new_control_ptr[0..overlap], old_control[0..overlap]);
937 ctx.control_allocator.rawFree(old_control, alignment, @returnAddress());
938 } else {
939 ctx.control_allocator.rawFree(new_control_ptr[0..new_len], alignment, @returnAddress());
940 return;
941 }
942
943 ctx.allocs[@backingInt(index)] = .{
944 .control_ptr = new_control_ptr,
945 .sample_ptr = memory.ptr,
946 .common = .{
947 .len = @fromBackingInt(@intCast(new_len)),
948 .alignment = alignment,
949 },
950 };
951
952 if (new_len > memory.len) {
953 for (
954 allocation.control_ptr[memory.len..new_len],
955 allocation.sample_ptr[memory.len..new_len],
956 0..,
957 ) |*control, *sample, i| {
958 control.* = @truncate(i);
959 sample.* = @truncate(i);
960 }
961 }
962
963 @atomicStore(Alloc.Index, &ctx.last_alloc_index, index, .release);
964 }
965 fn doOneRemap(ctx: *FuzzContext, new_len: usize) void {
966 @disableInstrumentation();
967 return doOneResize(ctx, new_len);
968 }
969 fn doOneFree(ctx: *FuzzContext) void {
970 @disableInstrumentation();
971
972 const index = @atomicRmw(Alloc.Index, &ctx.last_alloc_index, .Xchg, .none, .acquire);
973 if (index == .none) return;
974
975 const allocation = &ctx.allocs[@backingInt(index)];
976 assert(allocation.common.len != .free);
977 const len: usize = @backingInt(allocation.common.len);
978 const alignment = allocation.common.alignment;
979
980 assert(alignment.check(@intFromPtr(allocation.control_ptr)));
981 assert(alignment.check(@intFromPtr(allocation.sample_ptr)));
982
983 ctx.control_allocator.rawFree(allocation.control_ptr[0..len], alignment, @returnAddress());
984 ctx.sample_allocator.rawFree(allocation.sample_ptr[0..len], alignment, @returnAddress());
985
986 ctx.allocs[@backingInt(index)] = .{
987 .control_ptr = undefined,
988 .sample_ptr = undefined,
989 .common = .{
990 .len = .free,
991 .alignment = .@"1",
992 },
993 };
994 }
995};
996
997const FuzzAllocator = struct {
998 fba: std.heap.FixedBufferAllocator,
999 spurious_failures: [256]u8,
1000 index: u8,
1001
1002 fn init(buffer: []u8) FuzzAllocator {
1003 @disableInstrumentation();
1004 return .{
1005 .fba = .init(buffer),
1006 .spurious_failures = undefined, // set with `preprepareFailures`
1007 .index = 0,
1008 };
1009 }
1010
1011 fn prepareFailures(fa: *FuzzAllocator, smith: *std.testing.Smith) void {
1012 @disableInstrumentation();
1013 const bool_weights: []const std.testing.Smith.Weight = &.{
1014 .value(u8, 0, 10),
1015 .value(u8, 1, 1),
1016 };
1017 smith.bytesWeighted(&fa.spurious_failures, bool_weights);
1018 fa.index = 0;
1019 }
1020
1021 fn reset(fa: *FuzzAllocator) void {
1022 @disableInstrumentation();
1023 fa.fba.reset();
1024 }
1025
1026 fn allocator(fa: *FuzzAllocator) Allocator {
1027 @disableInstrumentation();
1028 return .{
1029 .ptr = fa,
1030 .vtable = &.{
1031 .alloc = FuzzAllocator.alloc,
1032 .resize = FuzzAllocator.resize,
1033 .remap = FuzzAllocator.remap,
1034 .free = FuzzAllocator.free,
1035 },
1036 };
1037 }
1038
1039 fn alloc(ctx: *anyopaque, len: usize, alignment: Alignment, ret_addr: usize) ?[*]u8 {
1040 @disableInstrumentation();
1041 const fa: *FuzzAllocator = @ptrCast(@alignCast(ctx));
1042 _ = ret_addr;
1043
1044 const index = @atomicRmw(u8, &fa.index, .Add, 1, .monotonic);
1045 if (fa.spurious_failures[index] != 0) return null;
1046 return fa.fba.threadSafeAllocator().rawAlloc(len, alignment, @returnAddress());
1047 }
1048
1049 fn resize(ctx: *anyopaque, memory: []u8, alignment: Alignment, new_len: usize, ret_addr: usize) bool {
1050 @disableInstrumentation();
1051 const fa: *FuzzAllocator = @ptrCast(@alignCast(ctx));
1052 _ = ret_addr;
1053
1054 const index = @atomicRmw(u8, &fa.index, .Add, 1, .monotonic);
1055 if (fa.spurious_failures[index] != 0) return false;
1056 return fa.fba.threadSafeAllocator().rawResize(memory, alignment, new_len, @returnAddress());
1057 }
1058
1059 fn remap(ctx: *anyopaque, memory: []u8, alignment: Alignment, new_len: usize, ret_addr: usize) ?[*]u8 {
1060 @disableInstrumentation();
1061 const fa: *FuzzAllocator = @ptrCast(@alignCast(ctx));
1062 _ = ret_addr;
1063
1064 const index = @atomicRmw(u8, &fa.index, .Add, 1, .monotonic);
1065 if (fa.spurious_failures[index] != 0) return null;
1066 return fa.fba.threadSafeAllocator().rawRemap(memory, alignment, new_len, @returnAddress());
1067 }
1068
1069 fn free(ctx: *anyopaque, memory: []u8, alignment: Alignment, ret_addr: usize) void {
1070 @disableInstrumentation();
1071 const fa: *FuzzAllocator = @ptrCast(@alignCast(ctx));
1072 _ = ret_addr;
1073 return fa.fba.threadSafeAllocator().rawFree(memory, alignment, @returnAddress());
1074 }
1075};