std: rework ArenaAllocator.State.promote to mutate the original state pointer

[mlugg]

Note: I think I'm okay to make this proposal based on what Vexu said in https://github.com/ziglang/zig/issues/13244#issuecomment-1287103168, since this is a small std change which doesn't impact compiler implementation and removes a potential footgun.

Motivation

std.heap.ArenaAllocator rightly suggests that when the child allocator is known through some other means, it can be more efficient to store a value of type ArenaAllocator.State, and promote it when it needs to be used. In theory, this is an elegant and simple pattern: however, the way it is currently designed is a bit unintuitive.

Suppose we've created an arena state and stored it, and now want to do some allocation. The correct code for this is as follows:

var arena = state.promote(allo);
try doSomethingWith(arena.allocator());
state = arena.state; // This is important!

The last line here is what I'm concerned with. promote simply creates an actual ArenaAllocator with the state that was passed into it; that means if you still want the state to be valid (as you normally would in this case) you have to copy it back afterwards. This is easy to forget or not realise, and missing this would lead to confusing state corruption on future allocations which could be quite tricky to debug.

This is made worse by the fact that state isn't updated until after doSomethingWith returns. That means if anything called within doSomethingWith itself makes use of state, all bets are off; we'll again get a corrupted state with overlapping allocations. This is an obvious footgun.

Proposal

Replace std.heap.ArenaAllocator.State.promote with the following:

pub fn promote(state: *State, child_allocator: Allocator) Promoted {
    return .{
        .child_allocator = child_allocator,
        .state = state,
    };
}

pub const Promoted = struct {
    child_allocator: Allocator,
    state: *State,

    pub fn allocator(self: Promoted) Allocator {
        return Allocator.init(self, alloc, resize, free);
    }

    pub fn deinit(self: Promoted) void {
        self.arena().deinit();
    }

    fn arena(self: Promoted) ArenaAllocator {
        return .{
            .child_allocator = self.child_allocator,
            .state = self.state.*,
        };
    }

    fn alloc(self: Promoted, n: usize, ptr_align: u29, len_align: u29, ra: usize) ![]u8 {
        var a = self.arena();
        defer self.state.* = a.state;
        return a.alloc(n, ptr_align, len_align, ra);
    }

    fn resize(self: Promoted, buf: []u8, buf_align: u29, new_len: usize, len_align: u29, ra: usize) ?usize {
        var a = self.arena();
        defer self.state.* = a.state;
        return a.resize(buf, buf_align, new_len, len_align, ra);
    }

    fn free(self: Promoted, buf: []u8, buf_align: u29, ra: usize) void {
        var a = self.arena();
        defer self.state.* = a.state;
        a.free(buf, buf_align, ra);
    }
};

Then, the interface can be used as e.g. doSomethingWith(state.promote(allo).allocator()) entirely safely, since every use of the allocator from Promoted.allocator automatically copies to the underlying state.

This makes the interface technically more restrictive, since promote now needs a mutable pointer to the underlying state. However, there are only two cases where you wish to use promote: to use the allocator, or to deinit the arena. In the former case, you almost certainly have a mutable state, since that's necessary for the State to really be useful: if for some reason you have a const state which you're going to use locally and then deinit, it's only one extra line to put it into a var or to just init an actual arena wrapping it. When deiniting, the logic is much the same: you almost certainly have a mutable pointer anyway. (Plus, it could be argued that ArenaAllocator.deinit should take a mutable pointer like ArenaAllocator.allocator - related #9814).

Performance Implications

The only real disadvantage I see to this proposal is that it almost certainly has slightly worse performance characteristics compared to status quo. This solution adds a level of indirection between the allocator and the actual arena state, and copies to the original state on every allocation (status quo has the advantage that the programmer chooses when to copy the state back, and so can do it only when absolutely necessary).

I haven't done any benchmarking on this code. However, its overhead should be pretty negligible, and I believe is an acceptable loss in exchange for this improvement in safety and ease of use.