1//! An allocator that attempts to allocate from the given buffer, falling back to
2//! `fallback_allocator` if this fails.
3
4const std = @import("../std.zig");
5const heap = std.heap;
6const testing = std.testing;
7
8const Alignment = std.mem.Alignment;
9const Allocator = std.mem.Allocator;
10const FixedBufferAllocator = std.heap.FixedBufferAllocator;
11
12const BufferFirstAllocator = @This();
13
14fallback_allocator: Allocator,
15fixed_buffer_allocator: FixedBufferAllocator,
16
17pub fn init(buffer: []u8, fallback_allocator: Allocator) BufferFirstAllocator {
18 return .{
19 .fallback_allocator = fallback_allocator,
20 .fixed_buffer_allocator = .init(buffer),
21 };
22}
23
24pub fn allocator(self: *BufferFirstAllocator) Allocator {
25 return .{
26 .ptr = self,
27 .vtable = &.{
28 .alloc = alloc,
29 .resize = resize,
30 .remap = remap,
31 .free = free,
32 },
33 };
34}
35
36fn alloc(
37 ctx: *anyopaque,
38 len: usize,
39 alignment: Alignment,
40 ra: usize,
41) ?[*]u8 {
42 const self: *BufferFirstAllocator = @ptrCast(@alignCast(ctx));
43 return FixedBufferAllocator.alloc(&self.fixed_buffer_allocator, len, alignment, ra) orelse
44 return self.fallback_allocator.rawAlloc(len, alignment, ra);
45}
46
47fn resize(
48 ctx: *anyopaque,
49 buf: []u8,
50 alignment: Alignment,
51 new_len: usize,
52 ra: usize,
53) bool {
54 const self: *BufferFirstAllocator = @ptrCast(@alignCast(ctx));
55 if (self.fixed_buffer_allocator.ownsPtr(buf.ptr)) {
56 return FixedBufferAllocator.resize(&self.fixed_buffer_allocator, buf, alignment, new_len, ra);
57 } else {
58 return self.fallback_allocator.rawResize(buf, alignment, new_len, ra);
59 }
60}
61
62fn remap(
63 context: *anyopaque,
64 memory: []u8,
65 alignment: Alignment,
66 new_len: usize,
67 return_address: usize,
68) ?[*]u8 {
69 const self: *BufferFirstAllocator = @ptrCast(@alignCast(context));
70 if (self.fixed_buffer_allocator.ownsPtr(memory.ptr)) {
71 return FixedBufferAllocator.remap(&self.fixed_buffer_allocator, memory, alignment, new_len, return_address);
72 } else {
73 return self.fallback_allocator.rawRemap(memory, alignment, new_len, return_address);
74 }
75}
76
77fn free(
78 ctx: *anyopaque,
79 buf: []u8,
80 alignment: Alignment,
81 ra: usize,
82) void {
83 const self: *BufferFirstAllocator = @ptrCast(@alignCast(ctx));
84 if (self.fixed_buffer_allocator.ownsPtr(buf.ptr)) {
85 return FixedBufferAllocator.free(&self.fixed_buffer_allocator, buf, alignment, ra);
86 } else {
87 return self.fallback_allocator.rawFree(buf, alignment, ra);
88 }
89}
90
91test "BufferFirstAllocator" {
92 // Buffer first specific tests
93 {
94 var buffer: [10]u8 = undefined;
95 var bfa_state: BufferFirstAllocator = .init(&buffer, std.testing.allocator);
96 const bfa = bfa_state.allocator();
97
98 // We're under the limit, so we should be allocated in the buffer
99 const txt0 = "hellowrld";
100 const buf0 = try bfa.create(@TypeOf(txt0.*));
101 buf0.* = txt0.*;
102 try testing.expect(bfa_state.fixed_buffer_allocator.ownsPtr(buf0.ptr));
103
104 // We're now over the limit, so we should be allocated from the fallback
105 const txt1 = "test!";
106 const buf1 = try bfa.create(@TypeOf(txt1.*));
107 buf1.* = txt1.*;
108 try testing.expect(!bfa_state.fixed_buffer_allocator.ownsPtr(buf1.ptr));
109
110 // Free the allocation that took up space in the buffer
111 try testing.expectEqualStrings(txt0, buf0);
112 bfa.destroy(buf0);
113
114 // The next allocation would go in the buffer, but it's too big so it doesn't
115 const txt2 = "qwertyqwerty";
116 const buf2 = try bfa.create(@TypeOf(txt2.*));
117 buf2.* = txt2.*;
118 try testing.expect(!bfa_state.fixed_buffer_allocator.ownsPtr(buf2.ptr));
119
120 // The next allocation is smaller and fits in the buffer
121 const txt3 = "dvorak";
122 const buf3 = try bfa.create(@TypeOf(txt3.*));
123 buf3.* = txt3.*;
124 try testing.expect(bfa_state.fixed_buffer_allocator.ownsPtr(buf3.ptr));
125
126 // The remainder in the buffer is too small for the following allocation so it falls back
127 const txt4 = "moretext";
128 const buf4 = try bfa.create(@TypeOf(txt4.*));
129 buf4.* = txt4.*;
130 try testing.expect(!bfa_state.fixed_buffer_allocator.ownsPtr(buf4.ptr));
131
132 // Check equality on the remaining buffers and free them
133 try testing.expectEqualStrings(txt1, buf1);
134 bfa.destroy(buf1);
135 try testing.expectEqualStrings(txt2, buf2);
136 bfa.destroy(buf2);
137 try testing.expectEqualStrings(txt3, buf3);
138 bfa.destroy(buf3);
139 try testing.expectEqualStrings(txt4, buf4);
140 bfa.destroy(buf4);
141
142 try testing.expectEqual(0, bfa_state.fixed_buffer_allocator.end_index);
143 }
144
145 // Standard allocator tests
146 {
147 var buf: [4096]u8 = undefined;
148 {
149 var bfa: BufferFirstAllocator = .init(&buf, std.testing.allocator);
150 try heap.testAllocator(bfa.allocator());
151 }
152 {
153 var bfa: BufferFirstAllocator = .init(&buf, std.testing.allocator);
154 try heap.testAllocatorAligned(bfa.allocator());
155 }
156 {
157 var bfa: BufferFirstAllocator = .init(&buf, std.testing.allocator);
158 try heap.testAllocatorLargeAlignment(bfa.allocator());
159 }
160 {
161 var bfa: BufferFirstAllocator = .init(&buf, std.testing.allocator);
162 try heap.testAllocatorAlignedShrink(bfa.allocator());
163 }
164 }
165}