From fd0fb26aba57ba14676dafc4c294bb5d9e75a649 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Felix=20Quei=C3=9Fner?= Date: Tue, 3 Jan 2023 21:15:20 +0100 Subject: [PATCH] Implements std.ArenaAllocator.reset() (#12590) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Felix "xq" Queißner --- lib/std/heap/arena_allocator.zig | 132 +++++++++++++++++++++++++++++++ 1 file changed, 132 insertions(+) diff --git a/lib/std/heap/arena_allocator.zig b/lib/std/heap/arena_allocator.zig index ec9783185f..7e40c1f53b 100644 --- a/lib/std/heap/arena_allocator.zig +++ b/lib/std/heap/arena_allocator.zig @@ -41,6 +41,8 @@ pub const ArenaAllocator = struct { } pub fn deinit(self: ArenaAllocator) void { + // NOTE: When changing this, make sure `reset()` is adjusted accordingly! + var it = self.state.buffer_list.first; while (it) |node| { // this has to occur before the free because the free frees node @@ -50,6 +52,113 @@ pub const ArenaAllocator = struct { } } + pub const ResetMode = union(enum) { + /// Releases all allocated memory in the arena. + free_all, + /// This will pre-heat the arena for future allocations by allocating a + /// large enough buffer for all previously done allocations. + /// 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 that after the + /// biggest operation, no memory allocations are performed anymore. + retain_capacity, + /// This is the same as `retain_capacity`, but the memory will be shrunk to + /// this value if it exceeds the limit. + retain_with_limit: usize, + }; + /// Queries the current memory use of this arena. + /// This will **not** include the storage required for internal keeping. + pub fn queryCapacity(self: ArenaAllocator) usize { + var size: usize = 0; + var it = self.state.buffer_list.first; + while (it) |node| : (it = node.next) { + // Compute the actually allocated size excluding the + // linked list node. + size += node.data.len - @sizeOf(BufNode); + } + return size; + } + /// Resets the arena allocator and frees all allocated memory. + /// + /// `mode` defines how the currently allocated memory is handled. + /// See the variant documentation for `ResetMode` for the effects of each mode. + /// + /// The function will return whether the reset operation was successful or not. + /// If the reallocation failed `false` is returned. The arena will still be fully + /// functional in that case, all memory is released. Future allocations just might + /// be slower. + /// + /// NOTE: If `mode` is `free_mode`, the function will always return `true`. + pub fn reset(self: *ArenaAllocator, mode: ResetMode) bool { + // Some words on the implementation: + // The reset function can be implemented with two basic approaches: + // - Counting how much bytes were allocated since the last reset, and storing that + // information in State. This will make reset fast and alloc only a teeny tiny bit + // slower. + // - Counting how much bytes were allocated by iterating the chunk linked list. This + // will make reset slower, but alloc() keeps the same speed when reset() as if reset() + // would not exist. + // + // The second variant was chosen for implementation, as with more and more calls to reset(), + // the function will get faster and faster. At one point, the complexity of the function + // will drop to amortized O(1), as we're only ever having a single chunk that will not be + // reallocated, and we're not even touching the backing allocator anymore. + // + // Thus, only the first hand full of calls to reset() will actually need to iterate the linked + // list, all future calls are just taking the first node, and only resetting the `end_index` + // value. + const current_capacity = if (mode != .free_all) + @sizeOf(BufNode) + self.queryCapacity() // we need at least space for exactly one node + the current capacity + else + 0; + if (mode == .free_all or current_capacity == 0) { + // just reset when we don't have anything to reallocate + self.deinit(); + self.state = State{}; + return true; + } + const total_size = switch (mode) { + .retain_capacity => current_capacity, + .retain_with_limit => |limit| std.math.min(limit, current_capacity), + .free_all => unreachable, + }; + // Free all nodes except for the last one + var it = self.state.buffer_list.first; + const maybe_first_node = while (it) |node| { + // this has to occur before the free because the free frees node + const next_it = node.next; + if (next_it == null) + break node; + self.child_allocator.free(node.data); + it = next_it; + } else null; + std.debug.assert(maybe_first_node == null or maybe_first_node.?.next == null); + // reset the state before we try resizing the buffers, so we definitly have reset the arena to 0. + self.state.end_index = 0; + if (maybe_first_node) |first_node| { + // perfect, no need to invoke the child_allocator + if (first_node.data.len == total_size) + return true; + const align_bits = std.math.log2_int(usize, @alignOf(BufNode)); + if (self.child_allocator.rawResize(first_node.data, align_bits, total_size, @returnAddress())) { + // successful resize + first_node.data.len = total_size; + } else { + // manual realloc + const new_ptr = self.child_allocator.rawAlloc(total_size, align_bits, @returnAddress()) orelse { + // we failed to preheat the arena properly, signal this to the user. + return false; + }; + self.child_allocator.rawFree(first_node.data, align_bits, @returnAddress()); + const node = @ptrCast(*BufNode, @alignCast(@alignOf(BufNode), new_ptr)); + node.* = BufNode{ + .data = new_ptr[0..total_size], + }; + self.state.buffer_list.first = node; + } + } + return true; + } + fn createNode(self: *ArenaAllocator, prev_len: usize, minimum_size: usize) ?*BufNode { const actual_min_size = minimum_size + (@sizeOf(BufNode) + 16); const big_enough_len = prev_len + actual_min_size; @@ -137,3 +246,26 @@ pub const ArenaAllocator = struct { } } }; + +test "ArenaAllocator (reset with preheating)" { + var arena_allocator = ArenaAllocator.init(std.testing.allocator); + defer arena_allocator.deinit(); + // provides some variance in the allocated data + var rng_src = std.rand.DefaultPrng.init(19930913); + const random = rng_src.random(); + var rounds: usize = 25; + while (rounds > 0) { + rounds -= 1; + _ = arena_allocator.reset(.retain_capacity); + var alloced_bytes: usize = 0; + var total_size: usize = random.intRangeAtMost(usize, 256, 16384); + while (alloced_bytes < total_size) { + const size = random.intRangeAtMost(usize, 16, 256); + const alignment = 32; + const slice = try arena_allocator.allocator().alignedAlloc(u8, alignment, size); + try std.testing.expect(std.mem.isAligned(@ptrToInt(slice.ptr), alignment)); + try std.testing.expectEqual(size, slice.len); + alloced_bytes += slice.len; + } + } +}