Pool with generational references in Zig

Here are some resource management concepts I've found useful when working with Zig, trying to avoid heap memory allocations. The first is relatively straight forward: a pool. It's here mostly to demonstrate the second concept: generational references.

• Pools allow you to allocate and free single items from a fixed buffer, in O(1) time.
• Generational references enable a referee to invalidate all references to it without touching the referents, in O(1) time.

Pool

Let's create a pool type, starting with the type declarations and fields:

pub fn Pool(comptime T: type) type {
    return struct {
        const Self = @This();

        pub const Item = struct {
            next: ?*Item,
            value: T,
        };

        avail: ?*Item = null,
    };
}

avail is a list of free Item. Item encapsulate the given type T so that T values can be linked into the free list.

We now need a way to allocate T from the pool:

pub fn alloc(self: *Self) !*T {
    if (self.avail) |item| {
        self.avail = item.next;
        return &item.value;
    } else return error.PoolEmpty;
}

This simply pulls the first Item off avail and returns a pointer to its value. If there are no free items available, it returns error.PoolEmpty; we've exhausted the pool and no more items can be allocated.

We also need to be able to free the allocated items, linking them back into avail. Because free is given a pointer to T, but the Item which contains the list link is its parent, we need to get a pointer to the parent struct with @fieldParentPtr:

pub fn free(self: *Self, p: *T) void {
    const container = @fieldParentPtr(Item, "value", p);

    container.next = self.avail;
    self.avail = container;
}

Some sanity checking could be implemented here. Because the pool knows what memory the underlying slice occupies, free could make sure that the T being freed actually belongs to the pool before it blindly assumes that its parent is an Item, but I'll leave that out. For now, you'll just have to make sure to aim the footgun between your toes.

Finally, we implement init. It initializes by taking a slice of Item and adding each to the avail list:

pub fn init(storage: []Item) Self {
    var p = Self{};

    for (storage) |*container| p.free(&container.value);

    return p;
}

We can now use our pool as follows:

var storage: [100]Pool(f32).Item = undefined;
var pool = Pool(f32).init(&storage);
var value_p = try pool.alloc();
// ...
pool.free(value_p)

Generational references

A problem when freeing a resource is that of invalidating any references to it. Anyone referring to a resource will have to know not to use it when it's no longer valid.

Generational references are a great solution to this problem for our pool. Instead of having alloc return plain pointer, it returns a fat pointer, containing both the actual pointer and a number indicating the generation of the item it's referring to. The referee also carries a generation number.

• When we refer to an object, we use a reference with the current generation number of the referee.
• When we want to invalidate any reference to a particular referee, e.g. when freeing it, we increment the generation number of the referee without touching the references.
• Referents detect that their references are invalid by comparing the generation number of the reference to the current generation of the referee.

This way there is no need to track individual references. The owner of the reference simply detects that the referee is of a different generation before resolving the reference.

This concept is useful in combination with the pool, where the underlying memory itself will never be freed, so we know that the generation counters will remain intact in memory.

Let's construct a wrapper on top of the pool we implemented above, instead returning generational reference when allocating. Again, let's start with the type declarations and fields:

pub fn GenerationalPool(comptime T: type) type {
    return struct {
        const Self = @This();

        pub const Item = struct {
            value: T,
            gen: u64 = 0,
        };

        pub const Reference = struct {
            item: *Item,
            gen: u64,

            pub fn resolve(self: Reference) ?*T {
                return if (self.gen == self.item.gen)
                    &self.item.value
                else
                    null;
            }
        };

        pool: Pool(Item),
    };
}

We have an Item. This is what will be managed by the Pool. It contains the T we passed and the generation counter. Then we have the Reference. This refers to an Item via a pointer and a generation counter. We can resolve a reference which result in null if the generation of the reference differs from the generation of the item it refers to.

Let's add methods corresponding to the ones in the Pool. First, to allocate an item and return a Reference to it:

pub fn alloc(self: *Self) !Reference {
    const item = try self.pool.alloc();
    return .{ .item = item, .gen = item.gen };
}

Callers will hold this Reference and access it via resolve.

We need to be able to free items. We free by passing *T, to make sure that the reference is valid by the time it's freed. As with Pool this means we need a pointer to its parent Item, to be able to increment its generation.

pub fn free(self: *Self, p: *T) void {
    const item = @fieldParentPtr(Item, "value", p);
    item.gen +%= 1;
    self.pool.free(item);
}

Finally, init the pool by passing a slice of []Pool(Item).Item. We set the generation counters to 0 before we initialize the pool with the same slice.

pub fn init(storage: []Pool(Item).Item) Self {
    for (storage) |*item| item.value.gen = 0;
    return .{ .pool = Pool(Item).init(storage) };
}

Now we can allocate items from the pool, get a generational reference to them, resolve and free them.

For example, consider this:

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

var storage: [100]Pool(GenerationalPool(f32).Item).Item = undefined;
var pool = GenerationalPool(f32).init(&storage);

const ref1 = try pool.alloc();
const ref2 = ref1;

// Both references are valid
expect(ref1.resolve() != null);
expect(ref2.resolve() != null);

// We free via ref1. Both references should be invalidated
if(ref1.resolve()) |p| pool.free(p);
expect(ref1.resolve() == null);
expect(ref2.resolve() == null);

How safe is it?

The only problem, in theory, is that resolve can't tell the difference between generation 3, generation 2^64+3 and 2*2^64+3 and so on. If our application creates a reference and goes through exactly 2^64 generations of the same Item before resolving it, it'll get get a false match. Due to the astronomical number of allocations the application would have to make for this to happen, this is unlikely to be a problem in our case.

For a frame of reference, freeing an item a billion times per second, our application would still only reach that point in 584 years, and it would only be a problem if we had a reference stick around for all that time only to be resolved exactly at the right moment, the next such moment occurring in another 584 years.

If you can't rest easy without being perfectly safe even in theory, instead of wrapping the counter, you can return an error when allocating an item where the current generation of the item it allocated is 2^64-1. Then, alloc will fail after ~18 quintillion allocations to the same pool slot, and your distant descendants can deal with the error by e.g. restarting the software.

---