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 {...@@ -78,7 +78,7 @@ fn countListCapacity(first_node: ?*Node) usize {
78 while (it) |node| : (it = node.next) {78 while (it) |node| : (it = node.next) {
79 // Compute the actually allocated size excluding the79 // Compute the actually allocated size excluding the
80 // linked list node.80 // linked list node.
81 capacity += node.size - @sizeOf(Node);81 capacity += node.size.toInt() - @sizeOf(Node);
82 }82 }
83 return capacity;83 return capacity;
84}84}
...@@ -164,7 +164,10 @@ pub fn reset(arena: *ArenaAllocator, mode: ResetMode) bool {...@@ -164,7 +164,10 @@ pub fn reset(arena: *ArenaAllocator, mode: ResetMode) bool {
164 };164 };
165 const allocated_slice = node.allocatedSliceUnsafe();165 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)) {
168 arena.child_allocator.rawFree(allocated_slice, .of(Node), @returnAddress());171 arena.child_allocator.rawFree(allocated_slice, .of(Node), @returnAddress());
169 first_node_ptr.* = null;172 first_node_ptr.* = null;
170 continue;173 continue;
...@@ -173,19 +176,17 @@ pub fn reset(arena: *ArenaAllocator, mode: ResetMode) bool {...@@ -173,19 +176,17 @@ pub fn reset(arena: *ArenaAllocator, mode: ResetMode) bool {
173 node.end_index = 0;176 node.end_index = 0;
174 first_node_ptr.* = node;177 first_node_ptr.* = node;
175178
176 const adjusted_capacity: usize = mem.alignForward(usize, new_capacity, 2);179 if (allocated_slice.len == new_size) {
177
178 if (allocated_slice.len - @sizeOf(Node) == adjusted_capacity) {
179 // perfect, no need to invoke the child_allocator180 // perfect, no need to invoke the child_allocator
180 continue;181 continue;
181 }182 }
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())) {
184 // successful resize185 // successful resize
185 node.size = adjusted_capacity;186 node.size = .fromInt(new_size);
186 } else {187 } else {
187 // manual realloc188 // 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 {
189 // we failed to preheat the arena properly, signal this to the user.190 // we failed to preheat the arena properly, signal this to the user.
190 ok = false;191 ok = false;
191 continue;192 continue;
...@@ -193,7 +194,7 @@ pub fn reset(arena: *ArenaAllocator, mode: ResetMode) bool {...@@ -193,7 +194,7 @@ pub fn reset(arena: *ArenaAllocator, mode: ResetMode) bool {
193 arena.child_allocator.rawFree(allocated_slice, .of(Node), @returnAddress());194 arena.child_allocator.rawFree(allocated_slice, .of(Node), @returnAddress());
194 const new_first_node: *Node = @ptrCast(@alignCast(new_ptr));195 const new_first_node: *Node = @ptrCast(@alignCast(new_ptr));
195 new_first_node.* = .{196 new_first_node.* = .{
196 .size = adjusted_capacity,197 .size = .fromInt(new_size),
197 .end_index = 0,198 .end_index = 0,
198 .next = null,199 .next = null,
199 };200 };
...@@ -224,30 +225,49 @@ const Node = struct {...@@ -224,30 +225,49 @@ const Node = struct {
224 /// accessing it.225 /// accessing it.
225 next: ?*Node,226 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
229 fn loadBuf(node: *Node) []u8 {250 fn loadBuf(node: *Node) []u8 {
230 // monotonic is fine since `size` can only ever grow, so the buffer returned251 // monotonic is fine since `size` can only ever grow, so the buffer returned
231 // by this function is always valid memory.252 // by this function is always valid memory.
232 const size = @atomicLoad(usize, &node.size, .monotonic);253 const size = @atomicLoad(Size, &node.size, .monotonic);
233 return @as([*]u8, @ptrCast(node))[0 .. size & ~resize_bit][@sizeOf(Node)..];254 return @as([*]u8, @ptrCast(node))[0..size.toInt()][@sizeOf(Node)..];
234 }255 }
235256
236 /// Returns allocated slice or `null` if node is already (being) resized.257 /// Returns allocated slice or `null` if node is already (being) resized.
237 fn beginResize(node: *Node) ?[]u8 {258 fn beginResize(node: *Node) ?[]u8 {
238 const size = @atomicRmw(usize, &node.size, .Or, resize_bit, .acquire); // syncs with release in `endResize`259 const size = @atomicRmw(Size, &node.size, .Or, .{ .resizing = true }, .acquire); // syncs with release in `endResize`
239 if (size & resize_bit != 0) return null;260 if (size.resizing) return null;
240 return @as([*]u8, @ptrCast(node))[0..size];261 return @as([*]u8, @ptrCast(node))[0..size.toInt()];
241 }262 }
242263
243 fn endResize(node: *Node, size: usize) void {264 fn endResize(node: *Node, size: usize) void {
244 assert(size & resize_bit == 0);265 return @atomicStore(Size, &node.size, .fromInt(size), .release); // syncs with acquire in `beginResize`
245 return @atomicStore(usize, &node.size, size, .release); // syncs with acquire in `beginResize`
246 }266 }
247267
248 /// Not threadsafe.268 /// Not threadsafe.
249 fn allocatedSliceUnsafe(node: *Node) []u8 {269 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()];
251 }271 }
252};272};
253273
...@@ -408,7 +428,7 @@ fn alloc(ctx: *anyopaque, n: usize, alignment: Alignment, ret_addr: usize) ?[*]u...@@ -408,7 +428,7 @@ fn alloc(ctx: *anyopaque, n: usize, alignment: Alignment, ret_addr: usize) ?[*]u
408 it = node.next;428 it = node.next;
409 }) {429 }) {
410 last_free = node;430 last_free = node;
411 assert(node.size & Node.resize_bit == 0);431 assert(!node.size.resizing);
412 const buf = node.allocatedSliceUnsafe()[@sizeOf(Node)..];432 const buf = node.allocatedSliceUnsafe()[@sizeOf(Node)..];
413 const aligned_index = alignedIndex(buf.ptr, 0, alignment);433 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...@@ -433,7 +453,7 @@ fn alloc(ctx: *anyopaque, n: usize, alignment: Alignment, ret_addr: usize) ?[*]u
433 const new_size = mem.alignForward(usize, @sizeOf(Node) + aligned_index + n, 2);453 const new_size = mem.alignForward(usize, @sizeOf(Node) + aligned_index + n, 2);
434454
435 if (arena.child_allocator.rawResize(allocated_slice, .of(Node), new_size, @returnAddress())) {455 if (arena.child_allocator.rawResize(allocated_slice, .of(Node), new_size, @returnAddress())) {
436 node.size = new_size;456 node.size = .fromInt(new_size);
437 break :candidate .{ node, best_fit_prev };457 break :candidate .{ node, best_fit_prev };
438 }458 }
439 }459 }
...@@ -486,12 +506,11 @@ fn alloc(ctx: *anyopaque, n: usize, alignment: Alignment, ret_addr: usize) ?[*]u...@@ -486,12 +506,11 @@ fn alloc(ctx: *anyopaque, n: usize, alignment: Alignment, ret_addr: usize) ?[*]u
486 const big_enough_size = prev_size + min_size + 16;506 const big_enough_size = prev_size + min_size + 16;
487 break :size mem.alignForward(usize, big_enough_size + big_enough_size / 2, 2);507 break :size mem.alignForward(usize, big_enough_size + big_enough_size / 2, 2);
488 };508 };
489 assert(size & Node.resize_bit == 0);
490 const ptr = arena.child_allocator.rawAlloc(size, .of(Node), @returnAddress()) orelse509 const ptr = arena.child_allocator.rawAlloc(size, .of(Node), @returnAddress()) orelse
491 return null;510 return null;
492 const new_node: *Node = @ptrCast(@alignCast(ptr));511 const new_node: *Node = @ptrCast(@alignCast(ptr));
493 new_node.* = .{512 new_node.* = .{
494 .size = size,513 .size = .fromInt(size),
495 .end_index = undefined, // set below514 .end_index = undefined, // set below
496 .next = undefined, // set below515 .next = undefined, // set below
497 };516 };
...@@ -501,7 +520,7 @@ fn alloc(ctx: *anyopaque, n: usize, alignment: Alignment, ret_addr: usize) ?[*]u...@@ -501,7 +520,7 @@ fn alloc(ctx: *anyopaque, n: usize, alignment: Alignment, ret_addr: usize) ?[*]u
501520
502 const buf = new_node.allocatedSliceUnsafe()[@sizeOf(Node)..];521 const buf = new_node.allocatedSliceUnsafe()[@sizeOf(Node)..];
503 const aligned_index = alignedIndex(buf.ptr, 0, alignment);522 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
506 new_node.end_index = aligned_index + n;525 new_node.end_index = aligned_index + n;
507 new_node.next = first_node;526 new_node.next = first_node;
...@@ -649,6 +668,7 @@ test "reset while retaining a buffer" {...@@ -649,6 +668,7 @@ test "reset while retaining a buffer" {
649 try std.testing.expect(arena_allocator.state.used_list.?.next != null);668 try std.testing.expect(arena_allocator.state.used_list.?.next != null);
650669
651 // This retains the first allocated buffer670 // 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 }));
653 try std.testing.expect(arena_allocator.state.used_list.?.next == null);672 try std.testing.expect(arena_allocator.state.used_list.?.next == null);
673 try std.testing.expectEqual(2, arena_allocator.queryCapacity());
654}674}