Arena comptime lock

Help
,
1

I’m using several arenas and I lock them when I use them, so that I don’t accidentally reset them further down the stack, and can only get/reset them in one scope:

code 
//! Arenas used by the editor.

/// Temporary allocations that don't survive function scope.
temp: Arena = .{},

/// Used when drawing rows.
render: Arena = .{},

/// Used when building edit targets.
targets: Arena = .{},

///////////////////////////////////////////////////////////////////////////////
//
//                              Types
//
///////////////////////////////////////////////////////////////////////////////

const Kind = enum { temp, targets, render };

const Arena = struct {
    /// Underlying arena instance.
    instance: std.heap.ArenaAllocator = undefined,

    /// True when an arena is already being used and shouldn't be reset.
    ///
    /// For example, if an outer scope starts using an arena, resetting it on
    /// scope exit, but it calls another function that does the same, if the
    /// second function resets the arena it invalidates all memory.
    ///
    /// To avoid this, arenas can only be used in one scope, and if they are
    /// needed further down their allocator must be passed explicitly.
    busy: bool = false,
};

///////////////////////////////////////////////////////////////////////////////
//
//                              Init/deinit
//
///////////////////////////////////////////////////////////////////////////////

pub fn init(gpa: mem.Allocator) Arenas {
    return .{
        .temp = .{ .instance = .init(gpa) },
        .render = .{ .instance = .init(gpa) },
        .targets = .{ .instance = .init(gpa) },
    };
}

pub fn deinit(self: *const Arenas) void {
    self.temp.instance.deinit();
    self.render.instance.deinit();
    self.targets.instance.deinit();
}

///////////////////////////////////////////////////////////////////////////////
//
//                              Methods
//
///////////////////////////////////////////////////////////////////////////////

/// Resets arenas with different strategies.
pub fn reset(self: *Arenas, comptime which: Kind) void {
    switch (which) {
        .temp => {
            self.temp.busy = false;
            _ = self.temp.instance.reset(.{ .retain_with_limit = 1024 * 1024 });
        },
        .targets => {
            self.targets.busy = false;
            _ = self.targets.instance.reset(.{ .retain_with_limit = 1024 * 1024 });
        },
        .render => {
            self.render.busy = false;
            _ = self.render.instance.reset(.retain_capacity);
        },
    }
}

/// Returns an arena allocator.
pub fn get(self: *Arenas, comptime which: Kind) !mem.Allocator {
    switch (which) {
        .temp => {
            if (self.temp.busy) {
                return error.ArenaIsBusy;
            }
            else {
                self.temp.busy = true;
                return self.temp.instance.allocator();
            }
        },

        .targets => {
            if (self.targets.busy) {
                return error.ArenaIsBusy;
            }
            else {
                self.targets.busy = true;
                return self.targets.instance.allocator();
            }
        },

        .render => {
            if (self.render.busy) {
                return error.ArenaIsBusy;
            }
            else {
                self.render.busy = true;
                return self.render.instance.allocator();
            }
        },
    }
}

///////////////////////////////////////////////////////////////////////////////
//
//                              Tests
//
///////////////////////////////////////////////////////////////////////////////

test "Arenas.get rejects a busy arena until reset" {
    inline for ([_]Kind{ .temp, .targets, .render }) |kind|
    {
        var arenas = Arenas.init(std.testing.allocator);
        defer arenas.deinit();

        _ = try arenas.get(kind);
        try testing.expectError(error.ArenaIsBusy, arenas.get(kind));

        arenas.reset(kind);
        _ = try arenas.get(kind);
    }
}

test "Arenas.reset only clears the requested arena" {
    var arenas = Arenas.init(std.testing.allocator);
    defer arenas.deinit();

    _ = try arenas.get(.temp);
    arenas.reset(.targets);

    try testing.expectError(error.ArenaIsBusy, arenas.get(.temp));
    _ = try arenas.get(.targets);
}

///////////////////////////////////////////////////////////////////////////////
//
//                              Constants
//
///////////////////////////////////////////////////////////////////////////////

const Arenas = @This();

const std = @import("std");
const mem = std.mem;
const testing = std.testing;

But I’m wondering if these busy checks can be moved to comptime, because from the functions flow it should be evident already at comptime if an arena is being used where it shouldn’t (because it’s being used somewhere up in the stack). Is this somehow possible?

2

This sort of comptime property enforcement / tracking seems difficult and not very ergonomic in Zig, you are basically adhoc re-inventing rust style borrow mechanics[1] where you try to statically ensure that some resource is used at most once.

I think it is better to stay in Zig style and just write normal code being aware of where something is allowed or not. If you really want safety checks you could for example use std.debug.SafetyLocks similar to std.ArrayHashMapUnmanaged’s (un-/lockPointers).

You would use one safety lock per arena, then you could replace the get and reset functions with two functions acquire and release the former one just locks the lock and returns the allocator, while the latter unlocks the lock and then calls reset. (You also could put the call to reset in acquire but then you have one unnecessary reset of the initially fresh arena)

That way the two functions create a span of time between which any nested calls to acquire or release cause an assertion failure in safe modes.

In unsafe modes all the safety locks change to zero sized types and get optimized out to noops.
This also means that you don’t need any booleans to track the state.

  1. or trying to use linear types, which isn’t a good fit for what comptime supports (it would be possible to model but tends to be very verbose because you are basically emulating a more complex type system with a simpler one) ↩︎