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.