diff --git a/lib/std/heap/ArenaAllocator.zig b/lib/std/heap/ArenaAllocator.zig index 5901e6474fd88750a9592ab35934053b81b3fa4b..a9fea912841df99be85feae311cde1340b9d3088 100644 --- a/lib/std/heap/ArenaAllocator.zig +++ b/lib/std/heap/ArenaAllocator.zig @@ -78,7 +78,7 @@ fn countListCapacity(first_node: ?*Node) usize { while (it) |node| : (it = node.next) { // Compute the actually allocated size excluding the // linked list node. - capacity += node.size - @sizeOf(Node); + capacity += node.size.toInt() - @sizeOf(Node); } return capacity; } @@ -164,7 +164,10 @@ pub fn reset(arena: *ArenaAllocator, mode: ResetMode) bool { }; const allocated_slice = node.allocatedSliceUnsafe(); - if (new_capacity == 0) { + // Align backwards to always stay below limit. + const new_size = mem.alignBackward(usize, @sizeOf(Node) + new_capacity, 2); + + if (new_size == @sizeOf(Node)) { arena.child_allocator.rawFree(allocated_slice, .of(Node), @returnAddress()); first_node_ptr.* = null; continue; @@ -173,19 +176,17 @@ pub fn reset(arena: *ArenaAllocator, mode: ResetMode) bool { node.end_index = 0; first_node_ptr.* = node; - const adjusted_capacity: usize = mem.alignForward(usize, new_capacity, 2); - - if (allocated_slice.len - @sizeOf(Node) == adjusted_capacity) { + if (allocated_slice.len == new_size) { // perfect, no need to invoke the child_allocator continue; } - if (arena.child_allocator.rawResize(allocated_slice, .of(Node), adjusted_capacity, @returnAddress())) { + if (arena.child_allocator.rawResize(allocated_slice, .of(Node), new_size, @returnAddress())) { // successful resize - node.size = adjusted_capacity; + node.size = .fromInt(new_size); } else { // manual realloc - const new_ptr = arena.child_allocator.rawAlloc(adjusted_capacity, .of(Node), @returnAddress()) orelse { + const new_ptr = arena.child_allocator.rawAlloc(new_size, .of(Node), @returnAddress()) orelse { // we failed to preheat the arena properly, signal this to the user. ok = false; continue; @@ -193,7 +194,7 @@ pub fn reset(arena: *ArenaAllocator, mode: ResetMode) bool { arena.child_allocator.rawFree(allocated_slice, .of(Node), @returnAddress()); const new_first_node: *Node = @ptrCast(@alignCast(new_ptr)); new_first_node.* = .{ - .size = adjusted_capacity, + .size = .fromInt(new_size), .end_index = 0, .next = null, }; @@ -224,30 +225,49 @@ const Node = struct { /// accessing it. next: ?*Node, - const resize_bit: usize = 1; + const Size = packed struct(usize) { + resizing: bool, + _: @Int(.unsigned, @bitSizeOf(usize) - 1) = 0, + + fn fromInt(int: usize) Size { + assert(int >= @sizeOf(Node)); + const size: Size = @bitCast(int); + assert(!size.resizing); + return size; + } + + fn toInt(size: Size) usize { + var int = size; + int.resizing = false; + return @bitCast(int); + } + + comptime { + assert(Size{ .resizing = true } == @as(Size, @bitCast(@as(usize, 1)))); + } + }; fn loadBuf(node: *Node) []u8 { // monotonic is fine since `size` can only ever grow, so the buffer returned // by this function is always valid memory. - const size = @atomicLoad(usize, &node.size, .monotonic); - return @as([*]u8, @ptrCast(node))[0 .. size & ~resize_bit][@sizeOf(Node)..]; + const size = @atomicLoad(Size, &node.size, .monotonic); + return @as([*]u8, @ptrCast(node))[0..size.toInt()][@sizeOf(Node)..]; } /// Returns allocated slice or `null` if node is already (being) resized. fn beginResize(node: *Node) ?[]u8 { - const size = @atomicRmw(usize, &node.size, .Or, resize_bit, .acquire); // syncs with release in `endResize` - if (size & resize_bit != 0) return null; - return @as([*]u8, @ptrCast(node))[0..size]; + const size = @atomicRmw(Size, &node.size, .Or, .{ .resizing = true }, .acquire); // syncs with release in `endResize` + if (size.resizing) return null; + return @as([*]u8, @ptrCast(node))[0..size.toInt()]; } fn endResize(node: *Node, size: usize) void { - assert(size & resize_bit == 0); - return @atomicStore(usize, &node.size, size, .release); // syncs with acquire in `beginResize` + return @atomicStore(Size, &node.size, .fromInt(size), .release); // syncs with acquire in `beginResize` } /// Not threadsafe. fn allocatedSliceUnsafe(node: *Node) []u8 { - return @as([*]u8, @ptrCast(node))[0 .. node.size & ~resize_bit]; + return @as([*]u8, @ptrCast(node))[0..node.size.toInt()]; } }; @@ -408,7 +428,7 @@ fn alloc(ctx: *anyopaque, n: usize, alignment: Alignment, ret_addr: usize) ?[*]u it = node.next; }) { last_free = node; - assert(node.size & Node.resize_bit == 0); + assert(!node.size.resizing); const buf = node.allocatedSliceUnsafe()[@sizeOf(Node)..]; const aligned_index = alignedIndex(buf.ptr, 0, alignment); @@ -433,7 +453,7 @@ fn alloc(ctx: *anyopaque, n: usize, alignment: Alignment, ret_addr: usize) ?[*]u const new_size = mem.alignForward(usize, @sizeOf(Node) + aligned_index + n, 2); if (arena.child_allocator.rawResize(allocated_slice, .of(Node), new_size, @returnAddress())) { - node.size = new_size; + node.size = .fromInt(new_size); break :candidate .{ node, best_fit_prev }; } } @@ -486,12 +506,11 @@ fn alloc(ctx: *anyopaque, n: usize, alignment: Alignment, ret_addr: usize) ?[*]u const big_enough_size = prev_size + min_size + 16; break :size mem.alignForward(usize, big_enough_size + big_enough_size / 2, 2); }; - assert(size & Node.resize_bit == 0); const ptr = arena.child_allocator.rawAlloc(size, .of(Node), @returnAddress()) orelse return null; const new_node: *Node = @ptrCast(@alignCast(ptr)); new_node.* = .{ - .size = size, + .size = .fromInt(size), .end_index = undefined, // set below .next = undefined, // set below }; @@ -501,7 +520,7 @@ fn alloc(ctx: *anyopaque, n: usize, alignment: Alignment, ret_addr: usize) ?[*]u const buf = new_node.allocatedSliceUnsafe()[@sizeOf(Node)..]; const aligned_index = alignedIndex(buf.ptr, 0, alignment); - assert(new_node.size >= @sizeOf(Node) + aligned_index + n); + assert(new_node.size.toInt() >= @sizeOf(Node) + aligned_index + n); new_node.end_index = aligned_index + n; new_node.next = first_node; @@ -649,6 +668,7 @@ test "reset while retaining a buffer" { try std.testing.expect(arena_allocator.state.used_list.?.next != null); // This retains the first allocated buffer - try std.testing.expect(arena_allocator.reset(.{ .retain_with_limit = 1 })); + try std.testing.expect(arena_allocator.reset(.{ .retain_with_limit = 2 })); try std.testing.expect(arena_allocator.state.used_list.?.next == null); + try std.testing.expectEqual(2, arena_allocator.queryCapacity()); }