1const std = @import("std");
  2
  3const mem = std.mem;
  4const uefi = std.os.uefi;
  5
  6const assert = std.debug.assert;
  7
  8const Allocator = mem.Allocator;
  9
 10const UefiPoolAllocator = struct {
 11    fn getHeader(ptr: [*]u8) *[*]align(8) u8 {
 12        return @as(*[*]align(8) u8, @ptrFromInt(@intFromPtr(ptr) - @sizeOf(usize)));
 13    }
 14
 15    fn alloc(
 16        _: *anyopaque,
 17        len: usize,
 18        alignment: mem.Alignment,
 19        ret_addr: usize,
 20    ) ?[*]u8 {
 21        _ = ret_addr;
 22
 23        assert(len > 0);
 24
 25        const ptr_align = alignment.toByteUnits();
 26
 27        const metadata_len = mem.alignForward(usize, @sizeOf(usize), ptr_align);
 28
 29        const full_len = metadata_len + len;
 30
 31        const unaligned_slice = uefi.system_table.boot_services.?.allocatePool(
 32            uefi.efi_pool_memory_type,
 33            full_len,
 34        ) catch return null;
 35
 36        const unaligned_addr = @intFromPtr(unaligned_slice.ptr);
 37        const aligned_addr = mem.alignForward(usize, unaligned_addr + @sizeOf(usize), ptr_align);
 38
 39        const aligned_ptr = unaligned_slice.ptr + (aligned_addr - unaligned_addr);
 40        getHeader(aligned_ptr).* = unaligned_slice.ptr;
 41
 42        return aligned_ptr;
 43    }
 44
 45    fn resize(
 46        _: *anyopaque,
 47        buf: []u8,
 48        alignment: mem.Alignment,
 49        new_len: usize,
 50        ret_addr: usize,
 51    ) bool {
 52        _ = ret_addr;
 53        _ = alignment;
 54
 55        if (new_len > buf.len) return false;
 56        return true;
 57    }
 58
 59    fn remap(
 60        _: *anyopaque,
 61        buf: []u8,
 62        alignment: mem.Alignment,
 63        new_len: usize,
 64        ret_addr: usize,
 65    ) ?[*]u8 {
 66        _ = alignment;
 67        _ = ret_addr;
 68
 69        if (new_len > buf.len) return null;
 70        return buf.ptr;
 71    }
 72
 73    fn free(
 74        _: *anyopaque,
 75        buf: []u8,
 76        alignment: mem.Alignment,
 77        ret_addr: usize,
 78    ) void {
 79        _ = alignment;
 80        _ = ret_addr;
 81        uefi.system_table.boot_services.?.freePool(getHeader(buf.ptr).*) catch unreachable;
 82    }
 83};
 84
 85/// Supports the full Allocator interface, including alignment.
 86/// For a direct call of `allocatePool`, see `raw_pool_allocator`.
 87pub const pool_allocator = Allocator{
 88    .ptr = undefined,
 89    .vtable = &pool_allocator_vtable,
 90};
 91
 92const pool_allocator_vtable = Allocator.VTable{
 93    .alloc = UefiPoolAllocator.alloc,
 94    .resize = UefiPoolAllocator.resize,
 95    .remap = UefiPoolAllocator.remap,
 96    .free = UefiPoolAllocator.free,
 97};
 98
 99/// Asserts allocations are 8 byte aligned and calls `boot_services.allocatePool`.
100pub const raw_pool_allocator = Allocator{
101    .ptr = undefined,
102    .vtable = &raw_pool_allocator_table,
103};
104
105const raw_pool_allocator_table = Allocator.VTable{
106    .alloc = uefi_alloc,
107    .resize = uefi_resize,
108    .remap = uefi_remap,
109    .free = uefi_free,
110};
111
112fn uefi_alloc(
113    _: *anyopaque,
114    len: usize,
115    alignment: mem.Alignment,
116    ret_addr: usize,
117) ?[*]u8 {
118    _ = ret_addr;
119
120    std.debug.assert(@intFromEnum(alignment) <= 3);
121
122    const slice = uefi.system_table.boot_services.?.allocatePool(
123        uefi.efi_pool_memory_type,
124        len,
125    ) catch return null;
126
127    return slice.ptr;
128}
129
130fn uefi_resize(
131    _: *anyopaque,
132    buf: []u8,
133    alignment: mem.Alignment,
134    new_len: usize,
135    ret_addr: usize,
136) bool {
137    _ = ret_addr;
138
139    std.debug.assert(@intFromEnum(alignment) <= 3);
140
141    if (new_len > buf.len) return false;
142    return true;
143}
144
145fn uefi_remap(
146    _: *anyopaque,
147    buf: []u8,
148    alignment: mem.Alignment,
149    new_len: usize,
150    ret_addr: usize,
151) ?[*]u8 {
152    _ = ret_addr;
153
154    std.debug.assert(@intFromEnum(alignment) <= 3);
155
156    if (new_len > buf.len) return null;
157    return buf.ptr;
158}
159
160fn uefi_free(
161    _: *anyopaque,
162    buf: []u8,
163    alignment: mem.Alignment,
164    ret_addr: usize,
165) void {
166    _ = alignment;
167    _ = ret_addr;
168    uefi.system_table.boot_services.?.freePool(@alignCast(buf.ptr)) catch unreachable;
169}