1const std = @import("../std.zig");
2const Allocator = std.mem.Allocator;
3const assert = std.debug.assert;
4const mem = std.mem;
5
6const FixedBufferAllocator = @This();
7
8end_index: usize,
9buffer: []u8,
10
11pub fn init(buffer: []u8) FixedBufferAllocator {
12 return .{
13 .buffer = buffer,
14 .end_index = 0,
15 };
16}
17
18/// Using this at the same time as the interface returned by `threadSafeAllocator` is not thread safe.
19pub fn allocator(self: *FixedBufferAllocator) Allocator {
20 return .{
21 .ptr = self,
22 .vtable = &.{
23 .alloc = alloc,
24 .resize = resize,
25 .remap = remap,
26 .free = free,
27 },
28 };
29}
30
31/// Provides a lock free thread safe `Allocator` interface to the underlying `FixedBufferAllocator`
32///
33/// Using this at the same time as the interface returned by `allocator` is not thread safe.
34pub fn threadSafeAllocator(self: *FixedBufferAllocator) Allocator {
35 return .{
36 .ptr = self,
37 .vtable = &.{
38 .alloc = threadSafeAlloc,
39 .resize = threadSafeResize,
40 .remap = threadSafeRemap,
41 .free = threadSafeFree,
42 },
43 };
44}
45
46pub fn ownsPtr(self: *FixedBufferAllocator, ptr: [*]u8) bool {
47 return sliceContainsPtr(self.buffer, ptr);
48}
49
50pub fn ownsSlice(self: *FixedBufferAllocator, slice: []u8) bool {
51 return sliceContainsSlice(self.buffer, slice);
52}
53
54/// This has false negatives when the last allocation had an
55/// adjusted_index. In such case we won't be able to determine what the
56/// last allocation was because the alignForward operation done in alloc is
57/// not reversible.
58pub fn isLastAllocation(self: *FixedBufferAllocator, buf: []u8) bool {
59 return buf.ptr + buf.len == self.buffer.ptr + self.end_index;
60}
61
62pub fn alloc(ctx: *anyopaque, n: usize, alignment: mem.Alignment, ra: usize) ?[*]u8 {
63 const self: *FixedBufferAllocator = @ptrCast(@alignCast(ctx));
64 _ = ra;
65 const ptr_align = alignment.toByteUnits();
66 const adjust_off = mem.alignPointerOffset(self.buffer.ptr + self.end_index, ptr_align) orelse return null;
67 const adjusted_index = self.end_index + adjust_off;
68 const new_end_index = adjusted_index + n;
69 if (new_end_index > self.buffer.len) return null;
70 self.end_index = new_end_index;
71 return self.buffer.ptr + adjusted_index;
72}
73
74pub fn resize(
75 ctx: *anyopaque,
76 buf: []u8,
77 alignment: mem.Alignment,
78 new_size: usize,
79 return_address: usize,
80) bool {
81 const self: *FixedBufferAllocator = @ptrCast(@alignCast(ctx));
82 _ = alignment;
83 _ = return_address;
84 assert(@inComptime() or self.ownsSlice(buf));
85
86 if (!self.isLastAllocation(buf)) {
87 if (new_size > buf.len) return false;
88 return true;
89 }
90
91 if (new_size <= buf.len) {
92 const sub = buf.len - new_size;
93 self.end_index -= sub;
94 return true;
95 }
96
97 const add = new_size - buf.len;
98 if (add + self.end_index > self.buffer.len) return false;
99
100 self.end_index += add;
101 return true;
102}
103
104pub fn remap(
105 context: *anyopaque,
106 memory: []u8,
107 alignment: mem.Alignment,
108 new_len: usize,
109 return_address: usize,
110) ?[*]u8 {
111 return if (resize(context, memory, alignment, new_len, return_address)) memory.ptr else null;
112}
113
114pub fn free(
115 ctx: *anyopaque,
116 buf: []u8,
117 alignment: mem.Alignment,
118 return_address: usize,
119) void {
120 const self: *FixedBufferAllocator = @ptrCast(@alignCast(ctx));
121 _ = alignment;
122 _ = return_address;
123 assert(@inComptime() or self.ownsSlice(buf));
124
125 if (self.isLastAllocation(buf)) {
126 self.end_index -= buf.len;
127 }
128}
129
130fn threadSafeAlloc(ctx: *anyopaque, n: usize, alignment: mem.Alignment, ret_addr: usize) ?[*]u8 {
131 const self: *FixedBufferAllocator = @ptrCast(@alignCast(ctx));
132 _ = ret_addr;
133 const ptr_align = alignment.toByteUnits();
134 var cur_end_index = @atomicLoad(usize, &self.end_index, .monotonic);
135 while (true) {
136 const adjust_off = mem.alignPointerOffset(self.buffer.ptr + cur_end_index, ptr_align) orelse return null;
137 const adjusted_index = cur_end_index + adjust_off;
138 const new_end_index = adjusted_index + n;
139 if (new_end_index > self.buffer.len) return null;
140 cur_end_index = @cmpxchgWeak(
141 usize,
142 &self.end_index,
143 cur_end_index,
144 new_end_index,
145 .acquire, // acquire any memory that may have been freed
146 .monotonic,
147 ) orelse
148 return self.buffer[adjusted_index..new_end_index].ptr;
149 }
150}
151
152fn threadSafeResize(ctx: *anyopaque, memory: []u8, alignment: mem.Alignment, new_len: usize, ret_addr: usize) bool {
153 const fba: *FixedBufferAllocator = @ptrCast(@alignCast(ctx));
154 _ = alignment;
155 _ = ret_addr;
156
157 const cur_end_index = @atomicLoad(usize, &fba.end_index, .monotonic);
158 if (fba.buffer.ptr + cur_end_index != memory.ptr + memory.len) {
159 // It's not the most recent allocation, so it cannot be expanded,
160 // but it's fine if they want to make it smaller.
161 return new_len <= memory.len;
162 }
163
164 if (new_len <= memory.len) {
165 const new_end_index = cur_end_index - (memory.len - new_len);
166 assert(fba.buffer.ptr + new_end_index == memory.ptr + new_len);
167
168 _ = @cmpxchgStrong(
169 usize,
170 &fba.end_index,
171 cur_end_index,
172 new_end_index,
173 .release, // release freed memory
174 .monotonic,
175 );
176 return true; // Shrinking allocations should always succeed.
177 }
178
179 if (fba.buffer.len - cur_end_index >= new_len - memory.len) {
180 const new_end_index = cur_end_index + (new_len - memory.len);
181 assert(fba.buffer.ptr + new_end_index == memory.ptr + new_len);
182
183 return null == @cmpxchgStrong(
184 usize,
185 &fba.end_index,
186 cur_end_index,
187 new_end_index,
188 .acquire, // acquire any memory that may have been freed
189 .monotonic,
190 );
191 }
192
193 return false;
194}
195
196fn threadSafeRemap(ctx: *anyopaque, memory: []u8, alignment: mem.Alignment, new_len: usize, ret_addr: usize) ?[*]u8 {
197 return if (threadSafeResize(ctx, memory, alignment, new_len, ret_addr)) memory.ptr else null;
198}
199
200fn threadSafeFree(ctx: *anyopaque, memory: []u8, alignment: mem.Alignment, ret_addr: usize) void {
201 const fba: *FixedBufferAllocator = @ptrCast(@alignCast(ctx));
202 _ = alignment;
203 _ = ret_addr;
204
205 assert(memory.len > 0);
206
207 const cur_end_index = @atomicLoad(usize, &fba.end_index, .monotonic);
208 if (fba.buffer.ptr + cur_end_index != memory.ptr + memory.len) {
209 // Not the most recent allocation; we cannot free it.
210 return;
211 }
212
213 const new_end_index = cur_end_index - memory.len;
214 assert(fba.buffer.ptr + new_end_index == memory.ptr);
215
216 _ = @cmpxchgStrong(
217 usize,
218 &fba.end_index,
219 cur_end_index,
220 new_end_index,
221 .release, // release freed memory
222 .monotonic,
223 );
224}
225
226pub fn reset(self: *FixedBufferAllocator) void {
227 self.end_index = 0;
228}
229
230fn sliceContainsPtr(container: []u8, ptr: [*]u8) bool {
231 return @intFromPtr(ptr) >= @intFromPtr(container.ptr) and
232 @intFromPtr(ptr) < (@intFromPtr(container.ptr) + container.len);
233}
234
235fn sliceContainsSlice(container: []u8, slice: []u8) bool {
236 return @intFromPtr(slice.ptr) >= @intFromPtr(container.ptr) and
237 (@intFromPtr(slice.ptr) + slice.len) <= (@intFromPtr(container.ptr) + container.len);
238}
239
240var test_fixed_buffer_allocator_memory: [800000 * @sizeOf(u64)]u8 = undefined;
241
242test FixedBufferAllocator {
243 var fixed_buffer_allocator = mem.validationWrap(FixedBufferAllocator.init(test_fixed_buffer_allocator_memory[0..]));
244 const a = fixed_buffer_allocator.allocator();
245
246 try std.heap.testAllocator(a);
247 try std.heap.testAllocatorAligned(a);
248 try std.heap.testAllocatorLargeAlignment(a);
249 try std.heap.testAllocatorAlignedShrink(a);
250}
251
252test reset {
253 var buf: [8]u8 align(@alignOf(u64)) = undefined;
254 var fba = FixedBufferAllocator.init(buf[0..]);
255 const a = fba.allocator();
256
257 const X = 0xeeeeeeeeeeeeeeee;
258 const Y = 0xffffffffffffffff;
259
260 const x = try a.create(u64);
261 x.* = X;
262 try std.testing.expectError(error.OutOfMemory, a.create(u64));
263
264 fba.reset();
265 const y = try a.create(u64);
266 y.* = Y;
267
268 // we expect Y to have overwritten X.
269 try std.testing.expect(x.* == y.*);
270 try std.testing.expect(y.* == Y);
271}
272
273test "reuse memory on realloc" {
274 var small_fixed_buffer: [10]u8 = undefined;
275 // check if we re-use the memory
276 {
277 var fixed_buffer_allocator = FixedBufferAllocator.init(small_fixed_buffer[0..]);
278 const a = fixed_buffer_allocator.allocator();
279
280 const slice0 = try a.alloc(u8, 5);
281 try std.testing.expect(slice0.len == 5);
282 const slice1 = try a.realloc(slice0, 10);
283 try std.testing.expect(slice1.ptr == slice0.ptr);
284 try std.testing.expect(slice1.len == 10);
285 try std.testing.expectError(error.OutOfMemory, a.realloc(slice1, 11));
286 }
287 // check that we don't re-use the memory if it's not the most recent block
288 {
289 var fixed_buffer_allocator = FixedBufferAllocator.init(small_fixed_buffer[0..]);
290 const a = fixed_buffer_allocator.allocator();
291
292 var slice0 = try a.alloc(u8, 2);
293 slice0[0] = 1;
294 slice0[1] = 2;
295 const slice1 = try a.alloc(u8, 2);
296 const slice2 = try a.realloc(slice0, 4);
297 try std.testing.expect(slice0.ptr != slice2.ptr);
298 try std.testing.expect(slice1.ptr != slice2.ptr);
299 try std.testing.expect(slice2[0] == 1);
300 try std.testing.expect(slice2[1] == 2);
301 }
302}
303
304test "thread safe version" {
305 var fixed_buffer_allocator = FixedBufferAllocator.init(test_fixed_buffer_allocator_memory[0..]);
306
307 try std.heap.testAllocator(fixed_buffer_allocator.threadSafeAllocator());
308 try std.heap.testAllocatorAligned(fixed_buffer_allocator.threadSafeAllocator());
309 try std.heap.testAllocatorLargeAlignment(fixed_buffer_allocator.threadSafeAllocator());
310 try std.heap.testAllocatorAlignedShrink(fixed_buffer_allocator.threadSafeAllocator());
311}