| ... | ... | @@ -41,6 +41,8 @@ pub const ArenaAllocator = struct { |
| 41 | 41 | } |
| 42 | 42 | |
| 43 | 43 | pub fn deinit(self: ArenaAllocator) void { |
| 44 | // NOTE: When changing this, make sure `reset()` is adjusted accordingly! |
| 45 | |
| 44 | 46 | var it = self.state.buffer_list.first; |
| 45 | 47 | while (it) |node| { |
| 46 | 48 | // this has to occur before the free because the free frees node |
| ... | ... | @@ -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 | 162 | fn createNode(self: *ArenaAllocator, prev_len: usize, minimum_size: usize) ?*BufNode { |
| 54 | 163 | const actual_min_size = minimum_size + (@sizeOf(BufNode) + 16); |
| 55 | 164 | const big_enough_len = prev_len + actual_min_size; |
| ... | ... | @@ -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 | } |