| ... | @@ -0,0 +1,123 @@ |
| 1 | //! A "composite" allocator that attempts to allocate first on the stack, using |
| 2 | //! the provided FixedBufferAllocator; upon failure, the provided secondary |
| 3 | //! allocator is used. reset() is NOT provided, even though available for the |
| 4 | //! (primary) FixedBufferAllocator, because it may not be available for the |
| 5 | //! provided secondary allocator (so callers must call reset() on underlying |
| 6 | //! allocators, themselves, when desirable). |
| 7 | |
| 8 | const std = @import("../std.zig"); |
| 9 | const mem = std.mem; |
| 10 | const Allocator = mem.Allocator; |
| 11 | const FixedBufferAllocator = std.heap.FixedBufferAllocator; |
| 12 | const assert = std.debug.assert; |
| 13 | |
| 14 | const StackFirstAllocator = @This(); |
| 15 | |
| 16 | primary: FixedBufferAllocator, |
| 17 | secondary: Allocator, |
| 18 | |
| 19 | pub fn init(buffer: []u8, secondary_allocator: Allocator) StackFirstAllocator { |
| 20 | return .{ |
| 21 | .primary = .init(buffer), |
| 22 | .secondary = secondary_allocator, |
| 23 | }; |
| 24 | } |
| 25 | |
| 26 | pub fn allocator(self: *StackFirstAllocator) Allocator { |
| 27 | return .{ |
| 28 | .ptr = self, |
| 29 | .vtable = &.{ |
| 30 | .alloc = alloc, |
| 31 | .resize = resize, |
| 32 | .remap = remap, |
| 33 | .free = free, |
| 34 | }, |
| 35 | }; |
| 36 | } |
| 37 | |
| 38 | pub fn alloc(ctx: *anyopaque, len: usize, alignment: mem.Alignment, ret_addr: usize) ?[*]u8 { |
| 39 | const self: *StackFirstAllocator = @ptrCast(@alignCast(ctx)); |
| 40 | return FixedBufferAllocator.alloc(&self.primary, len, alignment, ret_addr) orelse |
| 41 | self.secondary.rawAlloc(len, alignment, ret_addr); |
| 42 | } |
| 43 | |
| 44 | pub fn resize(ctx: *anyopaque, memory: []u8, alignment: mem.Alignment, new_len: usize, ret_addr: usize) bool { |
| 45 | const self: *StackFirstAllocator = @ptrCast(@alignCast(ctx)); |
| 46 | return if (self.primary.ownsPtr(memory.ptr)) |
| 47 | FixedBufferAllocator.resize(&self.primary, memory, alignment, new_len, ret_addr) |
| 48 | else |
| 49 | self.secondary.rawResize(memory, alignment, new_len, ret_addr); |
| 50 | } |
| 51 | |
| 52 | pub fn remap(ctx: *anyopaque, memory: []u8, alignment: mem.Alignment, new_len: usize, ret_addr: usize) ?[*]u8 { |
| 53 | const self: *StackFirstAllocator = @ptrCast(@alignCast(ctx)); |
| 54 | return if (self.primary.ownsPtr(memory.ptr)) |
| 55 | FixedBufferAllocator.remap(&self.primary, memory, alignment, new_len, ret_addr) |
| 56 | else |
| 57 | self.secondary.rawRemap(memory, alignment, new_len, ret_addr); |
| 58 | } |
| 59 | |
| 60 | pub fn free(ctx: *anyopaque, memory: []u8, alignment: mem.Alignment, ret_addr: usize) void { |
| 61 | const self: *StackFirstAllocator = @ptrCast(@alignCast(ctx)); |
| 62 | return if (self.primary.ownsPtr(memory.ptr)) |
| 63 | FixedBufferAllocator.free(&self.primary, memory, alignment, ret_addr) |
| 64 | else |
| 65 | self.secondary.rawFree(memory, alignment, ret_addr); |
| 66 | } |
| 67 | |
| 68 | test StackFirstAllocator { |
| 69 | var arena = std.heap.ArenaAllocator.init(std.testing.allocator); |
| 70 | defer arena.deinit(); |
| 71 | var buffer: [10]u8 = undefined; |
| 72 | var sfa = StackFirstAllocator.init(&buffer, arena.allocator()); |
| 73 | |
| 74 | const expect = std.testing.expect; |
| 75 | const expectEqualStrings = std.testing.expectEqualStrings; |
| 76 | |
| 77 | const al = sfa.allocator(); |
| 78 | const txt = "0123456789"; |
| 79 | |
| 80 | const dest = try al.alloc(u8, txt.len); |
| 81 | @memcpy(dest, txt); |
| 82 | try expectEqualStrings(txt, dest); |
| 83 | try expect(sfa.primary.ownsPtr(dest.ptr)); |
| 84 | |
| 85 | const txt2 = "abcde"; |
| 86 | const dest2 = try al.alloc(u8, txt2.len); |
| 87 | @memcpy(dest2, txt2); |
| 88 | try expectEqualStrings(txt2, dest2); |
| 89 | try expect(!sfa.primary.ownsPtr(dest2.ptr)); |
| 90 | |
| 91 | sfa.primary.reset(); |
| 92 | |
| 93 | const txt3 = "0123"; |
| 94 | const dest3 = try al.alloc(u8, txt3.len); |
| 95 | @memcpy(dest3, txt3); |
| 96 | try expectEqualStrings(txt3, dest3); |
| 97 | try expect(sfa.primary.ownsPtr(dest3.ptr)); |
| 98 | |
| 99 | sfa.primary.reset(); |
| 100 | //arena.reset(); // unnecessary, but allowed (note `defer arena.deinit()` above) |
| 101 | |
| 102 | // stock tests: |
| 103 | { |
| 104 | var buf: [16]u8 = undefined; |
| 105 | var a = StackFirstAllocator.init(&buf, std.testing.allocator); |
| 106 | try std.heap.testAllocator(a.allocator()); |
| 107 | } |
| 108 | { |
| 109 | var buf: [16]u8 = undefined; |
| 110 | var a = StackFirstAllocator.init(&buf, std.testing.allocator); |
| 111 | try std.heap.testAllocatorAligned(a.allocator()); |
| 112 | } |
| 113 | { |
| 114 | var buf: [16]u8 = undefined; |
| 115 | var a = StackFirstAllocator.init(&buf, std.testing.allocator); |
| 116 | try std.heap.testAllocatorLargeAlignment(a.allocator()); |
| 117 | } |
| 118 | { |
| 119 | var buf: [16]u8 = undefined; |
| 120 | var a = StackFirstAllocator.init(&buf, std.testing.allocator); |
| 121 | try std.heap.testAllocatorAlignedShrink(a.allocator()); |
| 122 | } |
| 123 | } |