| ... | @@ -41,6 +41,8 @@ pub const ArenaAllocator = struct { | ... | @@ -41,6 +41,8 @@ pub const ArenaAllocator = struct { |
| 41 | } | 41 | } |
| 42 | | 42 | |
| 43 | pub fn deinit(self: ArenaAllocator) void { | 43 | pub fn deinit(self: ArenaAllocator) void { |
| | 44 | // NOTE: When changing this, make sure `reset()` is adjusted accordingly! |
| | 45 | |
| 44 | var it = self.state.buffer_list.first; | 46 | var it = self.state.buffer_list.first; |
| 45 | while (it) |node| { | 47 | while (it) |node| { |
| 46 | // this has to occur before the free because the free frees node | 48 | // this has to occur before the free because the free frees node |
| ... | @@ -50,6 +52,113 @@ pub const ArenaAllocator = struct { | ... | @@ -50,6 +52,113 @@ pub const ArenaAllocator = struct { |
| 50 | } | 52 | } |
| 51 | } | 53 | } |
| 52 | | 54 | |
| | 55 | pub const ResetMode = union(enum) { |
| | 56 | /// Releases all allocated memory in the arena. |
| | 57 | free_all, |
| | 58 | /// This will pre-heat the arena for future allocations by allocating a |
| | 59 | /// large enough buffer for all previously done allocations. |
| | 60 | /// Preheating will speed up the allocation process by invoking the backing allocator |
| | 61 | /// less often than before. If `reset()` is used in a loop, this means that after the |
| | 62 | /// biggest operation, no memory allocations are performed anymore. |
| | 63 | retain_capacity, |
| | 64 | /// This is the same as `retain_capacity`, but the memory will be shrunk to |
| | 65 | /// this value if it exceeds the limit. |
| | 66 | retain_with_limit: usize, |
| | 67 | }; |
| | 68 | /// Queries the current memory use of this arena. |
| | 69 | /// This will **not** include the storage required for internal keeping. |
| | 70 | pub fn queryCapacity(self: ArenaAllocator) usize { |
| | 71 | var size: usize = 0; |
| | 72 | var it = self.state.buffer_list.first; |
| | 73 | while (it) |node| : (it = node.next) { |
| | 74 | // Compute the actually allocated size excluding the |
| | 75 | // linked list node. |
| | 76 | size += node.data.len - @sizeOf(BufNode); |
| | 77 | } |
| | 78 | return size; |
| | 79 | } |
| | 80 | /// Resets the arena allocator and frees all allocated memory. |
| | 81 | /// |
| | 82 | /// `mode` defines how the currently allocated memory is handled. |
| | 83 | /// See the variant documentation for `ResetMode` for the effects of each mode. |
| | 84 | /// |
| | 85 | /// The function will return whether the reset operation was successful or not. |
| | 86 | /// If the reallocation failed `false` is returned. The arena will still be fully |
| | 87 | /// functional in that case, all memory is released. Future allocations just might |
| | 88 | /// be slower. |
| | 89 | /// |
| | 90 | /// NOTE: If `mode` is `free_mode`, the function will always return `true`. |
| | 91 | pub fn reset(self: *ArenaAllocator, mode: ResetMode) bool { |
| | 92 | // Some words on the implementation: |
| | 93 | // The reset function can be implemented with two basic approaches: |
| | 94 | // - Counting how much bytes were allocated since the last reset, and storing that |
| | 95 | // information in State. This will make reset fast and alloc only a teeny tiny bit |
| | 96 | // slower. |
| | 97 | // - Counting how much bytes were allocated by iterating the chunk linked list. This |
| | 98 | // will make reset slower, but alloc() keeps the same speed when reset() as if reset() |
| | 99 | // would not exist. |
| | 100 | // |
| | 101 | // The second variant was chosen for implementation, as with more and more calls to reset(), |
| | 102 | // the function will get faster and faster. At one point, the complexity of the function |
| | 103 | // will drop to amortized O(1), as we're only ever having a single chunk that will not be |
| | 104 | // reallocated, and we're not even touching the backing allocator anymore. |
| | 105 | // |
| | 106 | // Thus, only the first hand full of calls to reset() will actually need to iterate the linked |
| | 107 | // list, all future calls are just taking the first node, and only resetting the `end_index` |
| | 108 | // value. |
| | 109 | const current_capacity = if (mode != .free_all) |
| | 110 | @sizeOf(BufNode) + self.queryCapacity() // we need at least space for exactly one node + the current capacity |
| | 111 | else |
| | 112 | 0; |
| | 113 | if (mode == .free_all or current_capacity == 0) { |
| | 114 | // just reset when we don't have anything to reallocate |
| | 115 | self.deinit(); |
| | 116 | self.state = State{}; |
| | 117 | return true; |
| | 118 | } |
| | 119 | const total_size = switch (mode) { |
| | 120 | .retain_capacity => current_capacity, |
| | 121 | .retain_with_limit => |limit| std.math.min(limit, current_capacity), |
| | 122 | .free_all => unreachable, |
| | 123 | }; |
| | 124 | // Free all nodes except for the last one |
| | 125 | var it = self.state.buffer_list.first; |
| | 126 | const maybe_first_node = while (it) |node| { |
| | 127 | // this has to occur before the free because the free frees node |
| | 128 | const next_it = node.next; |
| | 129 | if (next_it == null) |
| | 130 | break node; |
| | 131 | self.child_allocator.free(node.data); |
| | 132 | it = next_it; |
| | 133 | } else null; |
| | 134 | std.debug.assert(maybe_first_node == null or maybe_first_node.?.next == null); |
| | 135 | // reset the state before we try resizing the buffers, so we definitly have reset the arena to 0. |
| | 136 | self.state.end_index = 0; |
| | 137 | if (maybe_first_node) |first_node| { |
| | 138 | // perfect, no need to invoke the child_allocator |
| | 139 | if (first_node.data.len == total_size) |
| | 140 | return true; |
| | 141 | const align_bits = std.math.log2_int(usize, @alignOf(BufNode)); |
| | 142 | if (self.child_allocator.rawResize(first_node.data, align_bits, total_size, @returnAddress())) { |
| | 143 | // successful resize |
| | 144 | first_node.data.len = total_size; |
| | 145 | } else { |
| | 146 | // manual realloc |
| | 147 | const new_ptr = self.child_allocator.rawAlloc(total_size, align_bits, @returnAddress()) orelse { |
| | 148 | // we failed to preheat the arena properly, signal this to the user. |
| | 149 | return false; |
| | 150 | }; |
| | 151 | self.child_allocator.rawFree(first_node.data, align_bits, @returnAddress()); |
| | 152 | const node = @ptrCast(*BufNode, @alignCast(@alignOf(BufNode), new_ptr)); |
| | 153 | node.* = BufNode{ |
| | 154 | .data = new_ptr[0..total_size], |
| | 155 | }; |
| | 156 | self.state.buffer_list.first = node; |
| | 157 | } |
| | 158 | } |
| | 159 | return true; |
| | 160 | } |
| | 161 | |
| 53 | fn createNode(self: *ArenaAllocator, prev_len: usize, minimum_size: usize) ?*BufNode { | 162 | fn createNode(self: *ArenaAllocator, prev_len: usize, minimum_size: usize) ?*BufNode { |
| 54 | const actual_min_size = minimum_size + (@sizeOf(BufNode) + 16); | 163 | const actual_min_size = minimum_size + (@sizeOf(BufNode) + 16); |
| 55 | const big_enough_len = prev_len + actual_min_size; | 164 | const big_enough_len = prev_len + actual_min_size; |
| ... | @@ -137,3 +246,26 @@ pub const ArenaAllocator = struct { | ... | @@ -137,3 +246,26 @@ pub const ArenaAllocator = struct { |
| 137 | } | 246 | } |
| 138 | } | 247 | } |
| 139 | }; | 248 | }; |
| | 249 | |
| | 250 | test "ArenaAllocator (reset with preheating)" { |
| | 251 | var arena_allocator = ArenaAllocator.init(std.testing.allocator); |
| | 252 | defer arena_allocator.deinit(); |
| | 253 | // provides some variance in the allocated data |
| | 254 | var rng_src = std.rand.DefaultPrng.init(19930913); |
| | 255 | const random = rng_src.random(); |
| | 256 | var rounds: usize = 25; |
| | 257 | while (rounds > 0) { |
| | 258 | rounds -= 1; |
| | 259 | _ = arena_allocator.reset(.retain_capacity); |
| | 260 | var alloced_bytes: usize = 0; |
| | 261 | var total_size: usize = random.intRangeAtMost(usize, 256, 16384); |
| | 262 | while (alloced_bytes < total_size) { |
| | 263 | const size = random.intRangeAtMost(usize, 16, 256); |
| | 264 | const alignment = 32; |
| | 265 | const slice = try arena_allocator.allocator().alignedAlloc(u8, alignment, size); |
| | 266 | try std.testing.expect(std.mem.isAligned(@ptrToInt(slice.ptr), alignment)); |
| | 267 | try std.testing.expectEqual(size, slice.len); |
| | 268 | alloced_bytes += slice.len; |
| | 269 | } |
| | 270 | } |
| | 271 | } |