From b287b347278a81397ec2a28006ea216c9ae931fb Mon Sep 17 00:00:00 2001 From: Tristan Ross Date: Fri, 7 Feb 2025 18:46:25 -0800 Subject: [PATCH 1/4] std.os.uefi.tables.boot_services: adjust allocate page --- lib/std/os/uefi/tables/boot_services.zig | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/lib/std/os/uefi/tables/boot_services.zig b/lib/std/os/uefi/tables/boot_services.zig index f2a8b73b832a..3b1017268884 100644 --- a/lib/std/os/uefi/tables/boot_services.zig +++ b/lib/std/os/uefi/tables/boot_services.zig @@ -24,6 +24,8 @@ const OpenProtocolArgs = uefi.tables.OpenProtocolArgs; const OpenProtocolAttributes = uefi.tables.OpenProtocolAttributes; const ProtocolInformationEntry = uefi.tables.ProtocolInformationEntry; const EventNotify = uefi.tables.EventNotify; +const EfiEventNotify = uefi.tables.EfiEventNotify; +const EfiPhysicalAddress = uefi.tables.EfiPhysicalAddress; const cc = uefi.cc; const Error = Status.Error; @@ -48,7 +50,7 @@ pub const BootServices = extern struct { restoreTpl: *const fn (old_tpl: TaskPriorityLevel) callconv(cc) void, /// Allocates memory pages from the system. - _allocatePages: *const fn (alloc_type: AllocateType, mem_type: MemoryType, pages: usize, memory: *[*]align(4096) Page) callconv(cc) Status, + _allocatePages: *const fn (alloc_type: AllocateType, mem_type: MemoryType, pages: usize, memory: EfiPhysicalAddress) callconv(cc) Status, /// Frees memory pages. _freePages: *const fn (memory: [*]align(4096) Page, pages: usize) callconv(cc) Status, @@ -377,7 +379,7 @@ pub const BootServices = extern struct { mem_type: MemoryType, pages: usize, ) AllocatePagesError![]align(4096) Page { - var ptr: [*]align(4096) Page = switch (location) { + var ptr: EfiPhysicalAddress = switch (location) { .any => undefined, .address, .max_address => |ptr| ptr, }; From 789a93285b62dfe5f4df7d580e7b62e004f476c9 Mon Sep 17 00:00:00 2001 From: Tristan Ross Date: Fri, 7 Feb 2025 20:22:10 -0800 Subject: [PATCH 2/4] std.os.uefi: refactor allocators --- lib/std/heap.zig | 11 +- lib/std/os/uefi.zig | 10 +- lib/std/os/uefi/allocator.zig | 299 +++++++++++++++++++++++++++++ lib/std/os/uefi/pool_allocator.zig | 169 ---------------- 4 files changed, 314 insertions(+), 175 deletions(-) create mode 100644 lib/std/os/uefi/allocator.zig delete mode 100644 lib/std/os/uefi/pool_allocator.zig diff --git a/lib/std/heap.zig b/lib/std/heap.zig index 51e4fe44a2ca..2166aaee10ba 100644 --- a/lib/std/heap.zig +++ b/lib/std/heap.zig @@ -353,10 +353,13 @@ else if (builtin.target.cpu.arch.isWasm()) .{ } else if (builtin.target.os.tag == .plan9) .{ .ptr = undefined, .vtable = &SbrkAllocator(std.os.plan9.sbrk).vtable, -} else .{ - .ptr = undefined, - .vtable = &PageAllocator.vtable, -}; +} else if (builtin.target.os.tag == .uefi) + std.os.uefi.global_page_allocator.allocator() +else + .{ + .ptr = undefined, + .vtable = &PageAllocator.vtable, + }; pub const smp_allocator: Allocator = .{ .ptr = undefined, diff --git a/lib/std/os/uefi.zig b/lib/std/os/uefi.zig index 53aa884d4a90..30539ac29999 100644 --- a/lib/std/os/uefi.zig +++ b/lib/std/os/uefi.zig @@ -15,8 +15,14 @@ pub const tables = @import("uefi/tables.zig"); /// Defaults to `.loader_data`, the default data allocation type /// used by UEFI applications to allocate pool memory. pub var efi_pool_memory_type: tables.MemoryType = .loader_data; -pub const pool_allocator = @import("uefi/pool_allocator.zig").pool_allocator; -pub const raw_pool_allocator = @import("uefi/pool_allocator.zig").raw_pool_allocator; + +const allocator = @import("uefi/allocator.zig"); +pub const PageAllocator = allocator.Page; +pub const PoolAllocator = allocator.Pool; +pub const RawPoolAllocator = allocator.RawPool; + +pub var global_page_allocator = PageAllocator{}; +pub var global_pool_allocator = PoolAllocator{}; /// The EFI image's handle that is passed to its entry point. pub var handle: Handle = undefined; diff --git a/lib/std/os/uefi/allocator.zig b/lib/std/os/uefi/allocator.zig new file mode 100644 index 000000000000..6bf813d085fa --- /dev/null +++ b/lib/std/os/uefi/allocator.zig @@ -0,0 +1,299 @@ +const std = @import("../../std.zig"); + +const mem = std.mem; +const uefi = std.os.uefi; + +const Allocator = mem.Allocator; + +const assert = std.debug.assert; + +/// Allocates memory in pages. +/// +/// This allocator is backed by `allocatePages` and is therefore only suitable for usage when Boot Services are available. +pub const Page = struct { + memory_type: uefi.tables.MemoryType = .LoaderData, + + pub fn allocator(self: *Page) Allocator { + return Allocator{ + .ptr = self, + .vtable = &vtable, + }; + } + + const vtable = Allocator.VTable{ + .alloc = alloc, + .resize = resize, + .remap = remap, + .free = free, + }; + + fn alloc( + ctx: *anyopaque, + len: usize, + alignment: mem.Alignment, + ret_addr: usize, + ) ?[*]u8 { + _ = alignment; + _ = ret_addr; + + const self: *Page = @ptrCast(@alignCast(ctx)); + + assert(len > 0); + const pages = mem.alignForward(usize, len, 4096) / 4096; + + const buf = uefi.system_table.boot_services.?.allocatePages(.any, self.memory_type, pages) catch return null; + return buf.ptr; + } + + fn resize( + ctx: *anyopaque, + buf: []u8, + alignment: mem.Alignment, + new_len: usize, + ret_addr: usize, + ) bool { + _ = alignment; + _ = ret_addr; + + const self: *Page = @ptrCast(@alignCast(ctx)); + + // If the buffer was originally larger than the new length, we can grow or shrink it in place. + const original_len = mem.alignForward(usize, buf.len, 4096); + const new_aligned_len = mem.alignForward(usize, new_len, 4096); + + if (original_len >= new_aligned_len) return true; + + const new_pages_required = (new_aligned_len - original_len) / 4096; + const start_of_new_pages = @intFromPtr(buf.ptr) + original_len; + + // Try to allocate the necessary pages at the end of the buffer. + const new_pages = uefi.system_table.boot_services.?.allocatePages(.{ .at_address = start_of_new_pages }, self.memory_type, new_pages_required) catch return false; + _ = new_pages; + + // If the above function succeeds, then the new pages were successfully allocated. + return true; + } + + fn remap( + ctx: *anyopaque, + buf: []u8, + alignment: mem.Alignment, + new_len: usize, + ret_addr: usize, + ) ?[*]u8 { + _ = ctx; + _ = buf; + _ = alignment; + _ = new_len; + _ = ret_addr; + return null; + } + + fn free( + ctx: *anyopaque, + buf: []u8, + alignment: mem.Alignment, + ret_addr: usize, + ) void { + _ = ctx; + _ = alignment; + _ = ret_addr; + + const aligned_len = mem.alignForward(usize, buf.len, 4096); + const ptr: [*]align(4096) u8 = @alignCast(buf.ptr); + + uefi.system_table.boot_services.?.freePages(ptr[0..aligned_len]); + } +}; + +/// Supports the full std.mem.Allocator interface, including up to page alignment. +/// +/// This allocator is backed by `allocatePool` and is therefore only suitable for usage when Boot Services are available. +pub const Pool = struct { + memory_type: uefi.tables.MemoryType = .LoaderData, + + pub fn allocator(self: *Pool) Allocator { + return Allocator{ + .ptr = self, + .vtable = &vtable, + }; + } + + const vtable = Allocator.VTable{ + .alloc = alloc, + .resize = resize, + .remap = remap, + .free = free, + }; + + const Header = struct { + ptr: [*]align(8) u8, + len: usize, + }; + + fn getHeader(ptr: [*]u8) *align(1) Header { + return @ptrCast(ptr - @sizeOf(Header)); + } + + fn alloc( + ctx: *anyopaque, + len: usize, + alignment: mem.Alignment, + ret_addr: usize, + ) ?[*]u8 { + _ = ret_addr; + const self: *Pool = @ptrCast(@alignCast(ctx)); + + assert(len > 0); + + const ptr_align = alignment.toByteUnits(); + + // The maximum size of the metadata and any alignment padding. + const metadata_len = mem.alignForward(usize, @sizeOf(Header), ptr_align); + + const full_len = metadata_len + len; + + const buf = uefi.system_table.boot_services.?.allocatePool(self.memory_type, full_len) catch return null; + const unaligned_ptr = buf.ptr; + + const unaligned_addr = @intFromPtr(unaligned_ptr); + const aligned_addr = mem.alignForward(usize, unaligned_addr + @sizeOf(Header), ptr_align); + + const aligned_ptr: [*]u8 = @ptrFromInt(aligned_addr); + getHeader(aligned_ptr).ptr = unaligned_ptr; + getHeader(aligned_ptr).len = unaligned_addr + full_len - aligned_addr; + + return aligned_ptr; + } + + fn resize( + ctx: *anyopaque, + buf: []u8, + alignment: mem.Alignment, + new_len: usize, + ret_addr: usize, + ) bool { + _ = ctx; + _ = alignment; + _ = ret_addr; + + // If the buffer was originally larger than the new length, we can grow or shrink it in place. + if (getHeader(buf.ptr).len >= new_len) return true; + + // Otherwise, we cannot grow the buffer. + return false; + } + + fn remap( + ctx: *anyopaque, + buf: []u8, + alignment: mem.Alignment, + new_len: usize, + ret_addr: usize, + ) ?[*]u8 { + _ = ctx; + _ = buf; + _ = alignment; + _ = new_len; + _ = ret_addr; + return null; + } + + fn free( + ctx: *anyopaque, + buf: []u8, + alignment: mem.Alignment, + ret_addr: usize, + ) void { + _ = ctx; + _ = alignment; + _ = ret_addr; + + const header = getHeader(buf.ptr); + + uefi.system_table.boot_services.?.freePool(header.ptr[0..header.len]); + } +}; + +/// Asserts all allocations are at most 8 byte aligned. This is the highest alignment UEFI will give us directly. +/// +/// This allocator is backed by `allocatePool` and is therefore only suitable for usage when Boot Services are available. +pub const RawPool = struct { + memory_type: uefi.tables.MemoryType = .LoaderData, + + pub fn allocator(self: *RawPool) Allocator { + return Allocator{ + .ptr = self, + .vtable = &vtable, + }; + } + + pub const vtable = Allocator.VTable{ + .alloc = alloc, + .resize = resize, + .remap = remap, + .free = free, + }; + + fn alloc( + ctx: *anyopaque, + len: usize, + alignment: mem.Alignment, + ret_addr: usize, + ) ?[*]u8 { + _ = ret_addr; + const self: *RawPool = @ptrCast(@alignCast(ctx)); + + // UEFI pool allocations are 8 byte aligned, so we can't do better than that. + std.debug.assert(@intFromEnum(alignment) <= 3); + + const buf = uefi.system_table.boot_services.?.allocatePool(self.memory_type, len) catch return null; + return buf.ptr; + } + + fn resize( + ctx: *anyopaque, + buf: []u8, + alignment: mem.Alignment, + new_len: usize, + ret_addr: usize, + ) bool { + _ = ctx; + _ = alignment; + _ = ret_addr; + + // The original capacity is not known, so we can't ever grow the buffer. + if (new_len > buf.len) return false; + + // If this is a shrink, it will happen in place. + return true; + } + + fn remap( + ctx: *anyopaque, + buf: []u8, + alignment: mem.Alignment, + new_len: usize, + ret_addr: usize, + ) ?[*]u8 { + _ = ctx; + _ = buf; + _ = alignment; + _ = new_len; + _ = ret_addr; + return null; + } + + fn free( + ctx: *anyopaque, + buf: []u8, + alignment: mem.Alignment, + ret_addr: usize, + ) void { + _ = ctx; + _ = alignment; + _ = ret_addr; + + uefi.system_table.boot_services.?.freePool(@alignCast(buf)); + } +}; diff --git a/lib/std/os/uefi/pool_allocator.zig b/lib/std/os/uefi/pool_allocator.zig deleted file mode 100644 index 5acdcf6af4fe..000000000000 --- a/lib/std/os/uefi/pool_allocator.zig +++ /dev/null @@ -1,169 +0,0 @@ -const std = @import("std"); - -const mem = std.mem; -const uefi = std.os.uefi; - -const assert = std.debug.assert; - -const Allocator = mem.Allocator; - -const UefiPoolAllocator = struct { - fn getHeader(ptr: [*]u8) *[*]align(8) u8 { - return @as(*[*]align(8) u8, @ptrFromInt(@intFromPtr(ptr) - @sizeOf(usize))); - } - - fn alloc( - _: *anyopaque, - len: usize, - alignment: mem.Alignment, - ret_addr: usize, - ) ?[*]u8 { - _ = ret_addr; - - assert(len > 0); - - const ptr_align = alignment.toByteUnits(); - - const metadata_len = mem.alignForward(usize, @sizeOf(usize), ptr_align); - - const full_len = metadata_len + len; - - const unaligned_slice = uefi.system_table.boot_services.?.allocatePool( - uefi.efi_pool_memory_type, - full_len, - ) catch return null; - - const unaligned_addr = @intFromPtr(unaligned_slice.ptr); - const aligned_addr = mem.alignForward(usize, unaligned_addr + @sizeOf(usize), ptr_align); - - const aligned_ptr = unaligned_slice.ptr + (aligned_addr - unaligned_addr); - getHeader(aligned_ptr).* = unaligned_slice.ptr; - - return aligned_ptr; - } - - fn resize( - _: *anyopaque, - buf: []u8, - alignment: mem.Alignment, - new_len: usize, - ret_addr: usize, - ) bool { - _ = ret_addr; - _ = alignment; - - if (new_len > buf.len) return false; - return true; - } - - fn remap( - _: *anyopaque, - buf: []u8, - alignment: mem.Alignment, - new_len: usize, - ret_addr: usize, - ) ?[*]u8 { - _ = alignment; - _ = ret_addr; - - if (new_len > buf.len) return null; - return buf.ptr; - } - - fn free( - _: *anyopaque, - buf: []u8, - alignment: mem.Alignment, - ret_addr: usize, - ) void { - _ = alignment; - _ = ret_addr; - uefi.system_table.boot_services.?.freePool(getHeader(buf.ptr).*) catch unreachable; - } -}; - -/// Supports the full Allocator interface, including alignment. -/// For a direct call of `allocatePool`, see `raw_pool_allocator`. -pub const pool_allocator = Allocator{ - .ptr = undefined, - .vtable = &pool_allocator_vtable, -}; - -const pool_allocator_vtable = Allocator.VTable{ - .alloc = UefiPoolAllocator.alloc, - .resize = UefiPoolAllocator.resize, - .remap = UefiPoolAllocator.remap, - .free = UefiPoolAllocator.free, -}; - -/// Asserts allocations are 8 byte aligned and calls `boot_services.allocatePool`. -pub const raw_pool_allocator = Allocator{ - .ptr = undefined, - .vtable = &raw_pool_allocator_table, -}; - -const raw_pool_allocator_table = Allocator.VTable{ - .alloc = uefi_alloc, - .resize = uefi_resize, - .remap = uefi_remap, - .free = uefi_free, -}; - -fn uefi_alloc( - _: *anyopaque, - len: usize, - alignment: mem.Alignment, - ret_addr: usize, -) ?[*]u8 { - _ = ret_addr; - - std.debug.assert(@intFromEnum(alignment) <= 3); - - const slice = uefi.system_table.boot_services.?.allocatePool( - uefi.efi_pool_memory_type, - len, - ) catch return null; - - return slice.ptr; -} - -fn uefi_resize( - _: *anyopaque, - buf: []u8, - alignment: mem.Alignment, - new_len: usize, - ret_addr: usize, -) bool { - _ = ret_addr; - - std.debug.assert(@intFromEnum(alignment) <= 3); - - if (new_len > buf.len) return false; - return true; -} - -fn uefi_remap( - _: *anyopaque, - buf: []u8, - alignment: mem.Alignment, - new_len: usize, - ret_addr: usize, -) ?[*]u8 { - _ = ret_addr; - - std.debug.assert(@intFromEnum(alignment) <= 3); - - if (new_len > buf.len) return null; - return buf.ptr; -} - -fn uefi_free( - _: *anyopaque, - buf: []u8, - alignment: mem.Alignment, - ret_addr: usize, -) void { - _ = alignment; - _ = ret_addr; - uefi.system_table.boot_services.?.freePool(@alignCast(buf.ptr)) catch unreachable; -} From 7952530f36e18f86ddb284e56bd3d23058bf62d3 Mon Sep 17 00:00:00 2001 From: Tristan Ross Date: Sat, 8 Feb 2025 17:17:30 -0800 Subject: [PATCH 3/4] std.os.uefi: move allocators to heap --- lib/std/heap.zig | 4 ++- .../uefi_allocators.zig} | 31 ++++++++++--------- lib/std/os/uefi.zig | 8 ----- lib/std/os/uefi/tables/boot_services.zig | 6 ++-- 4 files changed, 23 insertions(+), 26 deletions(-) rename lib/std/{os/uefi/allocator.zig => heap/uefi_allocators.zig} (89%) diff --git a/lib/std/heap.zig b/lib/std/heap.zig index 2166aaee10ba..8a0faf9e18eb 100644 --- a/lib/std/heap.zig +++ b/lib/std/heap.zig @@ -25,6 +25,8 @@ pub const GeneralPurposeAllocatorConfig = DebugAllocatorConfig; /// Deprecated; to be removed after 0.14.0 is tagged. pub const GeneralPurposeAllocator = DebugAllocator; +pub const uefi = @import("heap/uefi_allocators.zig"); + const memory_pool = @import("heap/memory_pool.zig"); pub const MemoryPool = memory_pool.MemoryPool; pub const MemoryPoolAligned = memory_pool.MemoryPoolAligned; @@ -354,7 +356,7 @@ else if (builtin.target.cpu.arch.isWasm()) .{ .ptr = undefined, .vtable = &SbrkAllocator(std.os.plan9.sbrk).vtable, } else if (builtin.target.os.tag == .uefi) - std.os.uefi.global_page_allocator.allocator() + uefi.global_page_allocator.allocator() else .{ .ptr = undefined, diff --git a/lib/std/os/uefi/allocator.zig b/lib/std/heap/uefi_allocators.zig similarity index 89% rename from lib/std/os/uefi/allocator.zig rename to lib/std/heap/uefi_allocators.zig index 6bf813d085fa..5ef102846679 100644 --- a/lib/std/os/uefi/allocator.zig +++ b/lib/std/heap/uefi_allocators.zig @@ -1,4 +1,4 @@ -const std = @import("../../std.zig"); +const std = @import("../std.zig"); const mem = std.mem; const uefi = std.os.uefi; @@ -7,13 +7,16 @@ const Allocator = mem.Allocator; const assert = std.debug.assert; +pub var global_page_allocator = PageAllocator{}; +pub var global_pool_allocator = PoolAllocator{}; + /// Allocates memory in pages. /// /// This allocator is backed by `allocatePages` and is therefore only suitable for usage when Boot Services are available. -pub const Page = struct { - memory_type: uefi.tables.MemoryType = .LoaderData, +pub const PageAllocator = struct { + memory_type: uefi.tables.MemoryType = .loader_data, - pub fn allocator(self: *Page) Allocator { + pub fn allocator(self: *PageAllocator) Allocator { return Allocator{ .ptr = self, .vtable = &vtable, @@ -36,7 +39,7 @@ pub const Page = struct { _ = alignment; _ = ret_addr; - const self: *Page = @ptrCast(@alignCast(ctx)); + const self: *PageAllocator = @ptrCast(@alignCast(ctx)); assert(len > 0); const pages = mem.alignForward(usize, len, 4096) / 4096; @@ -55,7 +58,7 @@ pub const Page = struct { _ = alignment; _ = ret_addr; - const self: *Page = @ptrCast(@alignCast(ctx)); + const self: *PageAllocator = @ptrCast(@alignCast(ctx)); // If the buffer was originally larger than the new length, we can grow or shrink it in place. const original_len = mem.alignForward(usize, buf.len, 4096); @@ -109,10 +112,10 @@ pub const Page = struct { /// Supports the full std.mem.Allocator interface, including up to page alignment. /// /// This allocator is backed by `allocatePool` and is therefore only suitable for usage when Boot Services are available. -pub const Pool = struct { - memory_type: uefi.tables.MemoryType = .LoaderData, +pub const PoolAllocator = struct { + memory_type: uefi.tables.MemoryType = .loader_data, - pub fn allocator(self: *Pool) Allocator { + pub fn allocator(self: *PoolAllocator) Allocator { return Allocator{ .ptr = self, .vtable = &vtable, @@ -142,7 +145,7 @@ pub const Pool = struct { ret_addr: usize, ) ?[*]u8 { _ = ret_addr; - const self: *Pool = @ptrCast(@alignCast(ctx)); + const self: *PoolAllocator = @ptrCast(@alignCast(ctx)); assert(len > 0); @@ -218,10 +221,10 @@ pub const Pool = struct { /// Asserts all allocations are at most 8 byte aligned. This is the highest alignment UEFI will give us directly. /// /// This allocator is backed by `allocatePool` and is therefore only suitable for usage when Boot Services are available. -pub const RawPool = struct { - memory_type: uefi.tables.MemoryType = .LoaderData, +pub const RawPoolAllocator = struct { + memory_type: uefi.tables.MemoryType = .loader_data, - pub fn allocator(self: *RawPool) Allocator { + pub fn allocator(self: *RawPoolAllocator) Allocator { return Allocator{ .ptr = self, .vtable = &vtable, @@ -242,7 +245,7 @@ pub const RawPool = struct { ret_addr: usize, ) ?[*]u8 { _ = ret_addr; - const self: *RawPool = @ptrCast(@alignCast(ctx)); + const self: *RawPoolAllocator = @ptrCast(@alignCast(ctx)); // UEFI pool allocations are 8 byte aligned, so we can't do better than that. std.debug.assert(@intFromEnum(alignment) <= 3); diff --git a/lib/std/os/uefi.zig b/lib/std/os/uefi.zig index 30539ac29999..58452bfddb66 100644 --- a/lib/std/os/uefi.zig +++ b/lib/std/os/uefi.zig @@ -16,14 +16,6 @@ pub const tables = @import("uefi/tables.zig"); /// used by UEFI applications to allocate pool memory. pub var efi_pool_memory_type: tables.MemoryType = .loader_data; -const allocator = @import("uefi/allocator.zig"); -pub const PageAllocator = allocator.Page; -pub const PoolAllocator = allocator.Pool; -pub const RawPoolAllocator = allocator.RawPool; - -pub var global_page_allocator = PageAllocator{}; -pub var global_pool_allocator = PoolAllocator{}; - /// The EFI image's handle that is passed to its entry point. pub var handle: Handle = undefined; diff --git a/lib/std/os/uefi/tables/boot_services.zig b/lib/std/os/uefi/tables/boot_services.zig index 3b1017268884..a1be5cc8ce80 100644 --- a/lib/std/os/uefi/tables/boot_services.zig +++ b/lib/std/os/uefi/tables/boot_services.zig @@ -25,7 +25,7 @@ const OpenProtocolAttributes = uefi.tables.OpenProtocolAttributes; const ProtocolInformationEntry = uefi.tables.ProtocolInformationEntry; const EventNotify = uefi.tables.EventNotify; const EfiEventNotify = uefi.tables.EfiEventNotify; -const EfiPhysicalAddress = uefi.tables.EfiPhysicalAddress; +const PhysicalAddress = uefi.tables.PhysicalAddress; const cc = uefi.cc; const Error = Status.Error; @@ -50,7 +50,7 @@ pub const BootServices = extern struct { restoreTpl: *const fn (old_tpl: TaskPriorityLevel) callconv(cc) void, /// Allocates memory pages from the system. - _allocatePages: *const fn (alloc_type: AllocateType, mem_type: MemoryType, pages: usize, memory: EfiPhysicalAddress) callconv(cc) Status, + _allocatePages: *const fn (alloc_type: AllocateType, mem_type: MemoryType, pages: usize, memory: PhysicalAddress) callconv(cc) Status, /// Frees memory pages. _freePages: *const fn (memory: [*]align(4096) Page, pages: usize) callconv(cc) Status, @@ -379,7 +379,7 @@ pub const BootServices = extern struct { mem_type: MemoryType, pages: usize, ) AllocatePagesError![]align(4096) Page { - var ptr: EfiPhysicalAddress = switch (location) { + var ptr: PhysicalAddress = switch (location) { .any => undefined, .address, .max_address => |ptr| ptr, }; From d5966c85474ff0af8f63c101cfcb7d96a5c145f7 Mon Sep 17 00:00:00 2001 From: Tristan Ross Date: Thu, 10 Apr 2025 00:24:18 -0700 Subject: [PATCH 4/4] std.os.heap.uefi_allocators: split each allocator into a separate file --- lib/std/heap/uefi/PageAllocator.zig | 106 +++++++++ lib/std/heap/uefi/PoolAllocator.zig | 116 ++++++++++ lib/std/heap/uefi/RawPoolAllocator.zig | 90 ++++++++ lib/std/heap/uefi_allocators.zig | 302 +------------------------ 4 files changed, 315 insertions(+), 299 deletions(-) create mode 100644 lib/std/heap/uefi/PageAllocator.zig create mode 100644 lib/std/heap/uefi/PoolAllocator.zig create mode 100644 lib/std/heap/uefi/RawPoolAllocator.zig diff --git a/lib/std/heap/uefi/PageAllocator.zig b/lib/std/heap/uefi/PageAllocator.zig new file mode 100644 index 000000000000..ddbd3430eb53 --- /dev/null +++ b/lib/std/heap/uefi/PageAllocator.zig @@ -0,0 +1,106 @@ +/// Allocates memory in pages. +/// +/// This allocator is backed by `allocatePages` and is therefore only suitable for usage when Boot Services are available. +const std = @import("../../std.zig"); +const PageAllocator = @This(); + +const mem = std.mem; +const uefi = std.os.uefi; + +const Allocator = mem.Allocator; + +const assert = std.debug.assert; + +memory_type: uefi.tables.MemoryType = .loader_data, + +pub fn allocator(self: *PageAllocator) Allocator { + return Allocator{ + .ptr = self, + .vtable = &vtable, + }; +} + +const vtable = Allocator.VTable{ + .alloc = alloc, + .resize = resize, + .remap = remap, + .free = free, +}; + +fn alloc( + ctx: *anyopaque, + len: usize, + alignment: mem.Alignment, + ret_addr: usize, +) ?[*]u8 { + _ = alignment; + _ = ret_addr; + + const self: *PageAllocator = @ptrCast(@alignCast(ctx)); + + assert(len > 0); + const pages = mem.alignForward(usize, len, 4096) / 4096; + + const buf = uefi.system_table.boot_services.?.allocatePages(.any, self.memory_type, pages) catch return null; + return buf.ptr; +} + +fn resize( + ctx: *anyopaque, + buf: []u8, + alignment: mem.Alignment, + new_len: usize, + ret_addr: usize, +) bool { + _ = alignment; + _ = ret_addr; + + const self: *PageAllocator = @ptrCast(@alignCast(ctx)); + + // If the buffer was originally larger than the new length, we can grow or shrink it in place. + const original_len = mem.alignForward(usize, buf.len, 4096); + const new_aligned_len = mem.alignForward(usize, new_len, 4096); + + if (original_len >= new_aligned_len) return true; + + const new_pages_required = (new_aligned_len - original_len) / 4096; + const start_of_new_pages = @intFromPtr(buf.ptr) + original_len; + + // Try to allocate the necessary pages at the end of the buffer. + const new_pages = uefi.system_table.boot_services.?.allocatePages(.{ .at_address = start_of_new_pages }, self.memory_type, new_pages_required) catch return false; + _ = new_pages; + + // If the above function succeeds, then the new pages were successfully allocated. + return true; +} + +fn remap( + ctx: *anyopaque, + buf: []u8, + alignment: mem.Alignment, + new_len: usize, + ret_addr: usize, +) ?[*]u8 { + _ = ctx; + _ = buf; + _ = alignment; + _ = new_len; + _ = ret_addr; + return null; +} + +fn free( + ctx: *anyopaque, + buf: []u8, + alignment: mem.Alignment, + ret_addr: usize, +) void { + _ = ctx; + _ = alignment; + _ = ret_addr; + + const aligned_len = mem.alignForward(usize, buf.len, 4096); + const ptr: [*]align(4096) u8 = @alignCast(buf.ptr); + + uefi.system_table.boot_services.?.freePages(ptr[0..aligned_len]); +} diff --git a/lib/std/heap/uefi/PoolAllocator.zig b/lib/std/heap/uefi/PoolAllocator.zig new file mode 100644 index 000000000000..0a155f22e8cf --- /dev/null +++ b/lib/std/heap/uefi/PoolAllocator.zig @@ -0,0 +1,116 @@ +/// Supports the full std.mem.Allocator interface, including up to page alignment. +/// +/// This allocator is backed by `allocatePool` and is therefore only suitable for usage when Boot Services are available. +const std = @import("../../std.zig"); +const PoolAllocator = @This(); + +const mem = std.mem; +const uefi = std.os.uefi; + +const Allocator = mem.Allocator; + +const assert = std.debug.assert; + +memory_type: uefi.tables.MemoryType = .loader_data, + +pub fn allocator(self: *PoolAllocator) Allocator { + return Allocator{ + .ptr = self, + .vtable = &vtable, + }; +} + +const vtable = Allocator.VTable{ + .alloc = alloc, + .resize = resize, + .remap = remap, + .free = free, +}; + +const Header = struct { + ptr: [*]align(8) u8, + len: usize, +}; + +fn getHeader(ptr: [*]u8) *align(1) Header { + return @ptrCast(ptr - @sizeOf(Header)); +} + +fn alloc( + ctx: *anyopaque, + len: usize, + alignment: mem.Alignment, + ret_addr: usize, +) ?[*]u8 { + _ = ret_addr; + const self: *PoolAllocator = @ptrCast(@alignCast(ctx)); + + assert(len > 0); + + const ptr_align = alignment.toByteUnits(); + + // The maximum size of the metadata and any alignment padding. + const metadata_len = mem.alignForward(usize, @sizeOf(Header), ptr_align); + + const full_len = metadata_len + len; + + const buf = uefi.system_table.boot_services.?.allocatePool(self.memory_type, full_len) catch return null; + const unaligned_ptr = buf.ptr; + + const unaligned_addr = @intFromPtr(unaligned_ptr); + const aligned_addr = mem.alignForward(usize, unaligned_addr + @sizeOf(Header), ptr_align); + + const aligned_ptr: [*]u8 = @ptrFromInt(aligned_addr); + getHeader(aligned_ptr).ptr = unaligned_ptr; + getHeader(aligned_ptr).len = unaligned_addr + full_len - aligned_addr; + + return aligned_ptr; +} + +fn resize( + ctx: *anyopaque, + buf: []u8, + alignment: mem.Alignment, + new_len: usize, + ret_addr: usize, +) bool { + _ = ctx; + _ = alignment; + _ = ret_addr; + + // If the buffer was originally larger than the new length, we can grow or shrink it in place. + if (getHeader(buf.ptr).len >= new_len) return true; + + // Otherwise, we cannot grow the buffer. + return false; +} + +fn remap( + ctx: *anyopaque, + buf: []u8, + alignment: mem.Alignment, + new_len: usize, + ret_addr: usize, +) ?[*]u8 { + _ = ctx; + _ = buf; + _ = alignment; + _ = new_len; + _ = ret_addr; + return null; +} + +fn free( + ctx: *anyopaque, + buf: []u8, + alignment: mem.Alignment, + ret_addr: usize, +) void { + _ = ctx; + _ = alignment; + _ = ret_addr; + + const header = getHeader(buf.ptr); + + uefi.system_table.boot_services.?.freePool(header.ptr[0..header.len]); +} diff --git a/lib/std/heap/uefi/RawPoolAllocator.zig b/lib/std/heap/uefi/RawPoolAllocator.zig new file mode 100644 index 000000000000..c3636b2b3dad --- /dev/null +++ b/lib/std/heap/uefi/RawPoolAllocator.zig @@ -0,0 +1,90 @@ +/// Asserts all allocations are at most 8 byte aligned. This is the highest alignment UEFI will give us directly. +/// +/// This allocator is backed by `allocatePool` and is therefore only suitable for usage when Boot Services are available. +const std = @import("../../std.zig"); +const RawPoolAllocator = @This(); + +const mem = std.mem; +const uefi = std.os.uefi; + +const Allocator = mem.Allocator; + +const assert = std.debug.assert; + +memory_type: uefi.tables.MemoryType = .loader_data, + +pub fn allocator(self: *RawPoolAllocator) Allocator { + return Allocator{ + .ptr = self, + .vtable = &vtable, + }; +} + +pub const vtable = Allocator.VTable{ + .alloc = alloc, + .resize = resize, + .remap = remap, + .free = free, +}; + +fn alloc( + ctx: *anyopaque, + len: usize, + alignment: mem.Alignment, + ret_addr: usize, +) ?[*]u8 { + _ = ret_addr; + const self: *RawPoolAllocator = @ptrCast(@alignCast(ctx)); + + // UEFI pool allocations are 8 byte aligned, so we can't do better than that. + std.debug.assert(@intFromEnum(alignment) <= 3); + + const buf = uefi.system_table.boot_services.?.allocatePool(self.memory_type, len) catch return null; + return buf.ptr; +} + +fn resize( + ctx: *anyopaque, + buf: []u8, + alignment: mem.Alignment, + new_len: usize, + ret_addr: usize, +) bool { + _ = ctx; + _ = alignment; + _ = ret_addr; + + // The original capacity is not known, so we can't ever grow the buffer. + if (new_len > buf.len) return false; + + // If this is a shrink, it will happen in place. + return true; +} + +fn remap( + ctx: *anyopaque, + buf: []u8, + alignment: mem.Alignment, + new_len: usize, + ret_addr: usize, +) ?[*]u8 { + _ = ctx; + _ = buf; + _ = alignment; + _ = new_len; + _ = ret_addr; + return null; +} + +fn free( + ctx: *anyopaque, + buf: []u8, + alignment: mem.Alignment, + ret_addr: usize, +) void { + _ = ctx; + _ = alignment; + _ = ret_addr; + + uefi.system_table.boot_services.?.freePool(@alignCast(buf)); +} diff --git a/lib/std/heap/uefi_allocators.zig b/lib/std/heap/uefi_allocators.zig index 5ef102846679..68e92638cda6 100644 --- a/lib/std/heap/uefi_allocators.zig +++ b/lib/std/heap/uefi_allocators.zig @@ -1,302 +1,6 @@ -const std = @import("../std.zig"); - -const mem = std.mem; -const uefi = std.os.uefi; - -const Allocator = mem.Allocator; - -const assert = std.debug.assert; - pub var global_page_allocator = PageAllocator{}; pub var global_pool_allocator = PoolAllocator{}; -/// Allocates memory in pages. -/// -/// This allocator is backed by `allocatePages` and is therefore only suitable for usage when Boot Services are available. -pub const PageAllocator = struct { - memory_type: uefi.tables.MemoryType = .loader_data, - - pub fn allocator(self: *PageAllocator) Allocator { - return Allocator{ - .ptr = self, - .vtable = &vtable, - }; - } - - const vtable = Allocator.VTable{ - .alloc = alloc, - .resize = resize, - .remap = remap, - .free = free, - }; - - fn alloc( - ctx: *anyopaque, - len: usize, - alignment: mem.Alignment, - ret_addr: usize, - ) ?[*]u8 { - _ = alignment; - _ = ret_addr; - - const self: *PageAllocator = @ptrCast(@alignCast(ctx)); - - assert(len > 0); - const pages = mem.alignForward(usize, len, 4096) / 4096; - - const buf = uefi.system_table.boot_services.?.allocatePages(.any, self.memory_type, pages) catch return null; - return buf.ptr; - } - - fn resize( - ctx: *anyopaque, - buf: []u8, - alignment: mem.Alignment, - new_len: usize, - ret_addr: usize, - ) bool { - _ = alignment; - _ = ret_addr; - - const self: *PageAllocator = @ptrCast(@alignCast(ctx)); - - // If the buffer was originally larger than the new length, we can grow or shrink it in place. - const original_len = mem.alignForward(usize, buf.len, 4096); - const new_aligned_len = mem.alignForward(usize, new_len, 4096); - - if (original_len >= new_aligned_len) return true; - - const new_pages_required = (new_aligned_len - original_len) / 4096; - const start_of_new_pages = @intFromPtr(buf.ptr) + original_len; - - // Try to allocate the necessary pages at the end of the buffer. - const new_pages = uefi.system_table.boot_services.?.allocatePages(.{ .at_address = start_of_new_pages }, self.memory_type, new_pages_required) catch return false; - _ = new_pages; - - // If the above function succeeds, then the new pages were successfully allocated. - return true; - } - - fn remap( - ctx: *anyopaque, - buf: []u8, - alignment: mem.Alignment, - new_len: usize, - ret_addr: usize, - ) ?[*]u8 { - _ = ctx; - _ = buf; - _ = alignment; - _ = new_len; - _ = ret_addr; - return null; - } - - fn free( - ctx: *anyopaque, - buf: []u8, - alignment: mem.Alignment, - ret_addr: usize, - ) void { - _ = ctx; - _ = alignment; - _ = ret_addr; - - const aligned_len = mem.alignForward(usize, buf.len, 4096); - const ptr: [*]align(4096) u8 = @alignCast(buf.ptr); - - uefi.system_table.boot_services.?.freePages(ptr[0..aligned_len]); - } -}; - -/// Supports the full std.mem.Allocator interface, including up to page alignment. -/// -/// This allocator is backed by `allocatePool` and is therefore only suitable for usage when Boot Services are available. -pub const PoolAllocator = struct { - memory_type: uefi.tables.MemoryType = .loader_data, - - pub fn allocator(self: *PoolAllocator) Allocator { - return Allocator{ - .ptr = self, - .vtable = &vtable, - }; - } - - const vtable = Allocator.VTable{ - .alloc = alloc, - .resize = resize, - .remap = remap, - .free = free, - }; - - const Header = struct { - ptr: [*]align(8) u8, - len: usize, - }; - - fn getHeader(ptr: [*]u8) *align(1) Header { - return @ptrCast(ptr - @sizeOf(Header)); - } - - fn alloc( - ctx: *anyopaque, - len: usize, - alignment: mem.Alignment, - ret_addr: usize, - ) ?[*]u8 { - _ = ret_addr; - const self: *PoolAllocator = @ptrCast(@alignCast(ctx)); - - assert(len > 0); - - const ptr_align = alignment.toByteUnits(); - - // The maximum size of the metadata and any alignment padding. - const metadata_len = mem.alignForward(usize, @sizeOf(Header), ptr_align); - - const full_len = metadata_len + len; - - const buf = uefi.system_table.boot_services.?.allocatePool(self.memory_type, full_len) catch return null; - const unaligned_ptr = buf.ptr; - - const unaligned_addr = @intFromPtr(unaligned_ptr); - const aligned_addr = mem.alignForward(usize, unaligned_addr + @sizeOf(Header), ptr_align); - - const aligned_ptr: [*]u8 = @ptrFromInt(aligned_addr); - getHeader(aligned_ptr).ptr = unaligned_ptr; - getHeader(aligned_ptr).len = unaligned_addr + full_len - aligned_addr; - - return aligned_ptr; - } - - fn resize( - ctx: *anyopaque, - buf: []u8, - alignment: mem.Alignment, - new_len: usize, - ret_addr: usize, - ) bool { - _ = ctx; - _ = alignment; - _ = ret_addr; - - // If the buffer was originally larger than the new length, we can grow or shrink it in place. - if (getHeader(buf.ptr).len >= new_len) return true; - - // Otherwise, we cannot grow the buffer. - return false; - } - - fn remap( - ctx: *anyopaque, - buf: []u8, - alignment: mem.Alignment, - new_len: usize, - ret_addr: usize, - ) ?[*]u8 { - _ = ctx; - _ = buf; - _ = alignment; - _ = new_len; - _ = ret_addr; - return null; - } - - fn free( - ctx: *anyopaque, - buf: []u8, - alignment: mem.Alignment, - ret_addr: usize, - ) void { - _ = ctx; - _ = alignment; - _ = ret_addr; - - const header = getHeader(buf.ptr); - - uefi.system_table.boot_services.?.freePool(header.ptr[0..header.len]); - } -}; - -/// Asserts all allocations are at most 8 byte aligned. This is the highest alignment UEFI will give us directly. -/// -/// This allocator is backed by `allocatePool` and is therefore only suitable for usage when Boot Services are available. -pub const RawPoolAllocator = struct { - memory_type: uefi.tables.MemoryType = .loader_data, - - pub fn allocator(self: *RawPoolAllocator) Allocator { - return Allocator{ - .ptr = self, - .vtable = &vtable, - }; - } - - pub const vtable = Allocator.VTable{ - .alloc = alloc, - .resize = resize, - .remap = remap, - .free = free, - }; - - fn alloc( - ctx: *anyopaque, - len: usize, - alignment: mem.Alignment, - ret_addr: usize, - ) ?[*]u8 { - _ = ret_addr; - const self: *RawPoolAllocator = @ptrCast(@alignCast(ctx)); - - // UEFI pool allocations are 8 byte aligned, so we can't do better than that. - std.debug.assert(@intFromEnum(alignment) <= 3); - - const buf = uefi.system_table.boot_services.?.allocatePool(self.memory_type, len) catch return null; - return buf.ptr; - } - - fn resize( - ctx: *anyopaque, - buf: []u8, - alignment: mem.Alignment, - new_len: usize, - ret_addr: usize, - ) bool { - _ = ctx; - _ = alignment; - _ = ret_addr; - - // The original capacity is not known, so we can't ever grow the buffer. - if (new_len > buf.len) return false; - - // If this is a shrink, it will happen in place. - return true; - } - - fn remap( - ctx: *anyopaque, - buf: []u8, - alignment: mem.Alignment, - new_len: usize, - ret_addr: usize, - ) ?[*]u8 { - _ = ctx; - _ = buf; - _ = alignment; - _ = new_len; - _ = ret_addr; - return null; - } - - fn free( - ctx: *anyopaque, - buf: []u8, - alignment: mem.Alignment, - ret_addr: usize, - ) void { - _ = ctx; - _ = alignment; - _ = ret_addr; - - uefi.system_table.boot_services.?.freePool(@alignCast(buf)); - } -}; +pub const PageAllocator = @import("uefi/PageAllocator.zig"); +pub const PoolAllocator = @import("uefi/PoolAllocator.zig"); +pub const RawPoolAllocator = @import("uefi/RawPoolAllocator.zig");