authorgravatar for justus@klausecker.deJustus Klausecker <justus@klausecker.de> 2026-02-26 15:31:03+01:00
committergravatar for justus@klausecker.deJustus Klausecker <justus@klausecker.de> 2026-02-26 15:40:48+01:00
logde4112395714b43890aa1e9d1e7dd952e8947288
tree35566c638e59179349f37b2ef2c30620a657a446
parent2fa2300ba46f35a9073475bd24a8aa9180e0b228

std.heap.ArenaAllocator: fix `reset` creating undersized nodes

Previously resetting with `retain_capacity < @sizeOf(Node)` would create an invalid node. This is now fixed, plus `Node.size` now has its own `Size` type that provides additional safety via assertions to prevent bugs like this in the future.

1 files changed, 44 insertions(+), 24 deletions(-)

lib/std/heap/ArenaAllocator.zig+44-24
......@@ -78,7 +78,7 @@ fn countListCapacity(first_node: ?*Node) usize {
7878 while (it) |node| : (it = node.next) {
7979 // Compute the actually allocated size excluding the
8080 // linked list node.
81 capacity += node.size - @sizeOf(Node);
81 capacity += node.size.toInt() - @sizeOf(Node);
8282 }
8383 return capacity;
8484}
......@@ -164,7 +164,10 @@ pub fn reset(arena: *ArenaAllocator, mode: ResetMode) bool {
164164 };
165165 const allocated_slice = node.allocatedSliceUnsafe();
166166
167 if (new_capacity == 0) {
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)) {
168171 arena.child_allocator.rawFree(allocated_slice, .of(Node), @returnAddress());
169172 first_node_ptr.* = null;
170173 continue;
......@@ -173,19 +176,17 @@ pub fn reset(arena: *ArenaAllocator, mode: ResetMode) bool {
173176 node.end_index = 0;
174177 first_node_ptr.* = node;
175178
176 const adjusted_capacity: usize = mem.alignForward(usize, new_capacity, 2);
177
178 if (allocated_slice.len - @sizeOf(Node) == adjusted_capacity) {
179 if (allocated_slice.len == new_size) {
179180 // perfect, no need to invoke the child_allocator
180181 continue;
181182 }
182183
183 if (arena.child_allocator.rawResize(allocated_slice, .of(Node), adjusted_capacity, @returnAddress())) {
184 if (arena.child_allocator.rawResize(allocated_slice, .of(Node), new_size, @returnAddress())) {
184185 // successful resize
185 node.size = adjusted_capacity;
186 node.size = .fromInt(new_size);
186187 } else {
187188 // manual realloc
188 const new_ptr = arena.child_allocator.rawAlloc(adjusted_capacity, .of(Node), @returnAddress()) orelse {
189 const new_ptr = arena.child_allocator.rawAlloc(new_size, .of(Node), @returnAddress()) orelse {
189190 // we failed to preheat the arena properly, signal this to the user.
190191 ok = false;
191192 continue;
......@@ -193,7 +194,7 @@ pub fn reset(arena: *ArenaAllocator, mode: ResetMode) bool {
193194 arena.child_allocator.rawFree(allocated_slice, .of(Node), @returnAddress());
194195 const new_first_node: *Node = @ptrCast(@alignCast(new_ptr));
195196 new_first_node.* = .{
196 .size = adjusted_capacity,
197 .size = .fromInt(new_size),
197198 .end_index = 0,
198199 .next = null,
199200 };
......@@ -224,30 +225,49 @@ const Node = struct {
224225 /// accessing it.
225226 next: ?*Node,
226227
227 const resize_bit: usize = 1;
228 const Size = packed struct(usize) {
229 resizing: bool,
230 _: @Int(.unsigned, @bitSizeOf(usize) - 1) = 0,
231
232 fn fromInt(int: usize) Size {
233 assert(int >= @sizeOf(Node));
234 const size: Size = @bitCast(int);
235 assert(!size.resizing);
236 return size;
237 }
238
239 fn toInt(size: Size) usize {
240 var int = size;
241 int.resizing = false;
242 return @bitCast(int);
243 }
244
245 comptime {
246 assert(Size{ .resizing = true } == @as(Size, @bitCast(@as(usize, 1))));
247 }
248 };
228249
229250 fn loadBuf(node: *Node) []u8 {
230251 // monotonic is fine since `size` can only ever grow, so the buffer returned
231252 // by this function is always valid memory.
232 const size = @atomicLoad(usize, &node.size, .monotonic);
233 return @as([*]u8, @ptrCast(node))[0 .. size & ~resize_bit][@sizeOf(Node)..];
253 const size = @atomicLoad(Size, &node.size, .monotonic);
254 return @as([*]u8, @ptrCast(node))[0..size.toInt()][@sizeOf(Node)..];
234255 }
235256
236257 /// Returns allocated slice or `null` if node is already (being) resized.
237258 fn beginResize(node: *Node) ?[]u8 {
238 const size = @atomicRmw(usize, &node.size, .Or, resize_bit, .acquire); // syncs with release in `endResize`
239 if (size & resize_bit != 0) return null;
240 return @as([*]u8, @ptrCast(node))[0..size];
259 const size = @atomicRmw(Size, &node.size, .Or, .{ .resizing = true }, .acquire); // syncs with release in `endResize`
260 if (size.resizing) return null;
261 return @as([*]u8, @ptrCast(node))[0..size.toInt()];
241262 }
242263
243264 fn endResize(node: *Node, size: usize) void {
244 assert(size & resize_bit == 0);
245 return @atomicStore(usize, &node.size, size, .release); // syncs with acquire in `beginResize`
265 return @atomicStore(Size, &node.size, .fromInt(size), .release); // syncs with acquire in `beginResize`
246266 }
247267
248268 /// Not threadsafe.
249269 fn allocatedSliceUnsafe(node: *Node) []u8 {
250 return @as([*]u8, @ptrCast(node))[0 .. node.size & ~resize_bit];
270 return @as([*]u8, @ptrCast(node))[0..node.size.toInt()];
251271 }
252272};
253273
......@@ -408,7 +428,7 @@ fn alloc(ctx: *anyopaque, n: usize, alignment: Alignment, ret_addr: usize) ?[*]u
408428 it = node.next;
409429 }) {
410430 last_free = node;
411 assert(node.size & Node.resize_bit == 0);
431 assert(!node.size.resizing);
412432 const buf = node.allocatedSliceUnsafe()[@sizeOf(Node)..];
413433 const aligned_index = alignedIndex(buf.ptr, 0, alignment);
414434
......@@ -433,7 +453,7 @@ fn alloc(ctx: *anyopaque, n: usize, alignment: Alignment, ret_addr: usize) ?[*]u
433453 const new_size = mem.alignForward(usize, @sizeOf(Node) + aligned_index + n, 2);
434454
435455 if (arena.child_allocator.rawResize(allocated_slice, .of(Node), new_size, @returnAddress())) {
436 node.size = new_size;
456 node.size = .fromInt(new_size);
437457 break :candidate .{ node, best_fit_prev };
438458 }
439459 }
......@@ -486,12 +506,11 @@ fn alloc(ctx: *anyopaque, n: usize, alignment: Alignment, ret_addr: usize) ?[*]u
486506 const big_enough_size = prev_size + min_size + 16;
487507 break :size mem.alignForward(usize, big_enough_size + big_enough_size / 2, 2);
488508 };
489 assert(size & Node.resize_bit == 0);
490509 const ptr = arena.child_allocator.rawAlloc(size, .of(Node), @returnAddress()) orelse
491510 return null;
492511 const new_node: *Node = @ptrCast(@alignCast(ptr));
493512 new_node.* = .{
494 .size = size,
513 .size = .fromInt(size),
495514 .end_index = undefined, // set below
496515 .next = undefined, // set below
497516 };
......@@ -501,7 +520,7 @@ fn alloc(ctx: *anyopaque, n: usize, alignment: Alignment, ret_addr: usize) ?[*]u
501520
502521 const buf = new_node.allocatedSliceUnsafe()[@sizeOf(Node)..];
503522 const aligned_index = alignedIndex(buf.ptr, 0, alignment);
504 assert(new_node.size >= @sizeOf(Node) + aligned_index + n);
523 assert(new_node.size.toInt() >= @sizeOf(Node) + aligned_index + n);
505524
506525 new_node.end_index = aligned_index + n;
507526 new_node.next = first_node;
......@@ -649,6 +668,7 @@ test "reset while retaining a buffer" {
649668 try std.testing.expect(arena_allocator.state.used_list.?.next != null);
650669
651670 // This retains the first allocated buffer
652 try std.testing.expect(arena_allocator.reset(.{ .retain_with_limit = 1 }));
671 try std.testing.expect(arena_allocator.reset(.{ .retain_with_limit = 2 }));
653672 try std.testing.expect(arena_allocator.state.used_list.?.next == null);
673 try std.testing.expectEqual(2, arena_allocator.queryCapacity());
654674}