authorgravatar for xq@random-projects.netFelix "xq" Queißner <xq@random-projects.net> 2023-01-03 21:15:20+01:00
committergravatar for noreply@github.comGitHub <noreply@github.com> 2023-01-03 15:15:20-05:00
logfd0fb26aba57ba14676dafc4c294bb5d9e75a649
tree58d8630f2483084b3a77d3f81fcc5386293f48d1
parent8d64e5283a7cc51ece36e32b050ede6d7282d1b1
signaturebadge-question-mark Signed by PGP key 4AEE18F83AFDEB23

Implements std.ArenaAllocator.reset() (#12590)

Co-authored-by: Felix "xq" Queißner <xq@random-projects.net>

1 files changed, 132 insertions(+), 0 deletions(-)

lib/std/heap/arena_allocator.zig+132
...@@ -41,6 +41,8 @@ pub const ArenaAllocator = struct {...@@ -41,6 +41,8 @@ pub const ArenaAllocator = struct {
41 }41 }
4242
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 node48 // 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 }
5254
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
250test "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}