fn alloc(ctx: *anyopaque, n: usize, alignment: Alignment, ret_addr: usize) ?[*]u8 {
const arena: *ArenaAllocator = @ptrCast(@alignCast(ctx));
_ = ret_addr;
assert(n > 0);
var cur_first_node = arena.loadFirstNode();
var cur_new_node: ?*Node = null;
defer if (cur_new_node) |node| {
node.next = null; // optimize for empty free list
arena.pushFreeList(node, node);
};
retry: while (true) {
const first_node: ?*Node, const prev_size: usize = first_node: {
const node = cur_first_node orelse break :first_node .{ null, 0 };
const buf = node.loadBuf();
// To avoid using a CAS loop in the hot path we atomically increase
// `end_index` by a large enough amount to be able to always provide
// the required alignment within the reserved memory. To recover the
// space this potentially wastes we try to subtract the 'overshoot'
// with a single cmpxchg afterwards, which may fail.
const alignable = n + alignment.toByteUnits() - 1;
const end_index = @atomicRmw(usize, &node.end_index, .Add, alignable, .acquire); // acquire any memory that may have been freed
const aligned_index = alignedIndex(buf.ptr, end_index, alignment);
assert(end_index + alignable >= aligned_index + n);
if (end_index + alignable != aligned_index + n) {
_ = @cmpxchgStrong(
usize,
&node.end_index,
end_index + alignable,
aligned_index + n,
.monotonic, // no need to release alignment padding; there's no one accessing it!
.monotonic,
);
}
if (aligned_index + n > buf.len) break :first_node .{ node, buf.len };
return buf[aligned_index..][0..n].ptr;
};
resize: {
// Before attempting to get our hands on a new node, we try to resize
// the one we're currently holding. This is an exclusive operation;
// if another thread is already in this section we can never resize.
const node = first_node orelse break :resize;
const allocated_slice = node.beginResize() orelse break :resize;
var size = allocated_slice.len;
defer node.endResize(size, allocated_slice.len);
const buf = allocated_slice[@sizeOf(Node)..];
const end_index = @atomicLoad(usize, &node.end_index, .monotonic);
const aligned_index = alignedIndex(buf.ptr, end_index, alignment);
const new_size = mem.alignForward(usize, @sizeOf(Node) + aligned_index + n, 2);
if (new_size <= allocated_slice.len) {
// A `resize` or `free` call managed to sneak in and we need to
// guarantee that `size` is only ever increased; retry!
continue :retry;
}
if (arena.child_allocator.rawResize(allocated_slice, .of(Node), new_size, @returnAddress())) {
size = new_size;
// strong because a spurious failure could result in suboptimal
// usage of this node
if (null == @cmpxchgStrong(
usize,
&node.end_index,
end_index,
aligned_index + n,
.acquire, // acquire any memory that may have been freed
.monotonic,
)) {
const new_buf = allocated_slice.ptr[0..new_size][@sizeOf(Node)..];
return new_buf[aligned_index..][0..n].ptr;
}
}
}
// We need a new node! First, we search `free_list` for one that's big
// enough, if we don't find one there we fall back to allocating a new
// node with `child_allocator` (if we haven't already done that!).
from_free_list: {
// We 'steal' the entire free list to operate on it without other
// threads getting up into our business.
// This is a rather pragmatic approach, but since the free list isn't
// used very frequently it's fine performance-wise, even under load.
// Also this avoids the ABA problem; stealing the list with an atomic
// swap doesn't introduce any potentially stale `next` pointers.
const free_list = arena.stealFreeList() orelse break :from_free_list;
const first_free: *Node, const last_free: *Node, const node: *Node, const prev: ?*Node = find: {
var best_fit_prev: ?*Node = null;
var best_fit: ?*Node = null;
var best_fit_diff: usize = std.math.maxInt(usize);
var it_prev: ?*Node = null;
var it: ?*Node = free_list;
while (it) |node| : ({
it_prev = node;
it = node.next;
}) {
assert(!node.size.resizing);
const buf = node.allocatedSliceUnsafe()[@sizeOf(Node)..];
const aligned_index = alignedIndex(buf.ptr, 0, alignment);
const diff = aligned_index + n -| buf.len;
if (diff < best_fit_diff) {
best_fit_prev = it_prev;
best_fit = node;
best_fit_diff = diff;
}
}
break :find .{ free_list, it_prev.?, best_fit.?, best_fit_prev };
};
const aligned_index, const need_resize = aligned_index: {
const buf = node.allocatedSliceUnsafe()[@sizeOf(Node)..];
const aligned_index = alignedIndex(buf.ptr, 0, alignment);
break :aligned_index .{ aligned_index, aligned_index + n > buf.len };
};
if (need_resize) {
// Ideally we want to use all nodes in `free_list` eventually,
// so even if none fit we'll try to resize the one that was the
// closest to being large enough.
const new_size = mem.alignForward(usize, @sizeOf(Node) + aligned_index + n, 2);
if (arena.child_allocator.rawResize(node.allocatedSliceUnsafe(), .of(Node), new_size, @returnAddress())) {
node.size = .fromInt(new_size);
} else {
arena.pushFreeList(first_free, last_free);
break :from_free_list; // we couldn't find a fitting free node
}
}
const buf = node.allocatedSliceUnsafe()[@sizeOf(Node)..];
const old_next = node.next;
node.end_index = aligned_index + n;
node.next = first_node;
switch (arena.tryPushNode(node)) {
.success => {
// Finish removing node from free list.
if (prev) |p| p.next = old_next;
// Push remaining stolen free list back onto `arena.state.free_list`.
const new_first_free = if (node == first_free) old_next else first_free;
const new_last_free = if (node == last_free) prev else last_free;
if (new_first_free) |first| {
const last = new_last_free.?;
arena.pushFreeList(first, last);
}
return buf[aligned_index..][0..n].ptr;
},
.failure => |old_first_node| {
// restore free list to as we found it
node.next = old_next;
arena.pushFreeList(first_free, last_free);
cur_first_node = old_first_node;
continue :retry; // there's a new first node; retry!
},
}
}
const new_node: *Node = new_node: {
if (cur_new_node) |new_node| {
break :new_node new_node;
} else {
@branchHint(.cold);
}
const size: Node.Size = size: {
const min_size = @sizeOf(Node) + alignment.toByteUnits() + n;
const big_enough_size = prev_size + min_size + 16;
const size = mem.alignForward(usize, big_enough_size + big_enough_size / 2, 2);
break :size .fromInt(size);
};
const ptr = arena.child_allocator.rawAlloc(size.toInt(), .of(Node), @returnAddress()) orelse
return null;
const new_node: *Node = @ptrCast(@alignCast(ptr));
new_node.* = .{
.size = size,
.end_index = undefined, // set below
.next = undefined, // set below
};
cur_new_node = new_node;
break :new_node new_node;
};
const buf = new_node.allocatedSliceUnsafe()[@sizeOf(Node)..];
const aligned_index = alignedIndex(buf.ptr, 0, alignment);
assert(new_node.size.toInt() >= @sizeOf(Node) + aligned_index + n);
new_node.end_index = aligned_index + n;
new_node.next = first_node;
switch (arena.tryPushNode(new_node)) {
.success => {
cur_new_node = null;
return buf[aligned_index..][0..n].ptr;
},
.failure => |old_first_node| {
cur_first_node = old_first_node;
},
}
}
}
