feat(uefi): add post-boot-services allocator

Signed-off-by: Martin Kröning <martin.kroening@eonerc.rwth-aachen.de>

Cargo.toml

@@ -8,6 +8,7 @@ edition = "2021"
 
 [dependencies]
 align-address = "0.3"
+allocator-api2 = { version = "0.2", default-features = false }
 anstyle = { version = "1", default-features = false }
 cfg-if = "1"
 hermit-entry = { version = "0.10", features = ["loader"] }
@@ -33,11 +34,8 @@ aarch64-cpu = "9"
 hermit-dtb = { version = "0.1" }
 goblin = { version = "0.8", default-features = false, features = ["elf64"] }
 
-[target.'cfg(target_os = "none")'.dependencies]
-allocator-api2 = { version = "0.2", default-features = false }
-
 [target.'cfg(target_os = "uefi")'.dependencies]
-uefi = { version = "0.28", features = ["alloc", "global_allocator", "panic_handler", "qemu"] }
+uefi = { version = "0.28", features = ["alloc", "panic_handler", "qemu"] }
 qemu-exit = "3"
 
 [target.'cfg(target_arch = "riscv64")'.dependencies]

src/main.rs

@@ -14,6 +14,7 @@ use hermit_entry::boot_info::{BootInfo, RawBootInfo};
 mod macros;
 
 mod arch;
+mod bump_allocator;
 mod log;
 mod os;
 

src/os/uefi/allocator.rs

@@ -0,0 +1,59 @@
+use core::alloc::{GlobalAlloc, Layout};
+use core::mem::MaybeUninit;
+use core::ptr::{self, NonNull};
+
+use alloc::vec;
+use allocator_api2::alloc::Allocator;
+use one_shot_mutex::OneShotMutex;
+
+use crate::bump_allocator::BumpAllocator;
+
+pub enum GlobalAllocator {
+	Uefi,
+	Bump(BumpAllocator),
+}
+
+pub struct LockedAllocator(OneShotMutex<GlobalAllocator>);
+
+impl LockedAllocator {
+    pub const fn uefi() -> Self {
+        Self(OneShotMutex::new(GlobalAllocator::Uefi))
+    }
+}
+
+unsafe impl GlobalAlloc for LockedAllocator {
+	unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+		match &*self.0.lock() {
+			GlobalAllocator::Uefi => unsafe { uefi::allocator::Allocator.alloc(layout) },
+			GlobalAllocator::Bump(bump) => bump
+				.allocate(layout)
+				// FIXME: Use NonNull::as_mut_ptr once `slice_ptr_get` is stabilized
+				// https://github.com/rust-lang/rust/issues/74265
+				.map_or(ptr::null_mut(), |ptr| ptr.as_ptr().cast()),
+		}
+	}
+
+	unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
+		match &*self.0.lock() {
+			GlobalAllocator::Uefi => unsafe { uefi::allocator::Allocator.dealloc(ptr, layout) },
+			GlobalAllocator::Bump(bump) => unsafe {
+				bump.deallocate(NonNull::new(ptr).unwrap(), layout)
+			},
+		}
+	}
+}
+
+#[global_allocator]
+static ALLOCATOR: LockedAllocator = LockedAllocator::uefi();
+
+pub fn exit_boot_services() {
+    assert!(matches!(*ALLOCATOR.0.lock(), GlobalAllocator::Uefi));
+
+    let mem = vec![MaybeUninit::uninit(); 4096].leak();
+
+    let bump = BumpAllocator::from(mem);
+
+    *ALLOCATOR.0.lock() = GlobalAllocator::Bump(bump);
+
+	uefi::allocator::exit_boot_services();
+}

src/os/uefi/mod.rs

@@ -1,3 +1,4 @@
+mod allocator;
 mod console;
 
 use alloc::string::String;
@@ -14,6 +15,9 @@ pub use self::console::CONSOLE;
 #[entry]
 fn loader_main(_handle: Handle, mut system_table: SystemTable<Boot>) -> Status {
 	uefi::helpers::init(&mut system_table).unwrap();
+	unsafe {
+		uefi::allocator::init(&mut system_table);
+	}
 	crate::log::init();
 
 	let app = read_app(system_table.boot_services());