Use r-efi-alloc for allocator

Switch to using r-efi-alloc for allocation. This is useful since a lot
of UEFI projects will probably end up using this allocator anyway to
have custom allocation types or to simply not use std.

Signed-off-by: Ayush Singh <[email protected]>

Cargo.lock

@@ -2953,6 +2953,17 @@ dependencies = [
  "rustc-std-workspace-core",
 ]
 
+[[package]]
+name = "r-efi-alloc"
+version = "1.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "31d6f09fe2b6ad044bc3d2c34ce4979796581afd2f1ebc185837e02421e02fd7"
+dependencies = [
+ "compiler_builtins",
+ "r-efi",
+ "rustc-std-workspace-core",
+]
+
 [[package]]
 name = "rand"
 version = "0.7.3"
@@ -4909,6 +4920,7 @@ dependencies = [
  "panic_unwind",
  "profiler_builtins",
  "r-efi",
+ "r-efi-alloc",
  "rand 0.7.3",
  "rustc-demangle",
  "std_detect",

library/std/Cargo.toml

@@ -49,6 +49,7 @@ wasi = { version = "0.11.0", features = ['rustc-dep-of-std'], default-features =
 
 [target.'cfg(target_os = "uefi")'.dependencies]
 r-efi = { version = "4.1.0", features = ['rustc-dep-of-std', 'efiapi']}
+r-efi-alloc = { version = "1.0.0", features = ['rustc-dep-of-std']}
 
 [features]
 backtrace = [

library/std/src/sys/uefi/alloc.rs

@@ -1,10 +1,7 @@
 //! Global Allocator for UEFI.
-//! Uses `EFI_BOOT_SERVICES.AllocatePool()` and `EFI_BOOT_SERVICES.FreePool()`.
-//! Takes a lot of inspiration from Windows allocator for Alignment > 8.
+//! Uses [r-efi-alloc](https://crates.io/crates/r-efi-alloc)
 
-use crate::alloc::{GlobalAlloc, Layout, System};
-use crate::ptr;
-use crate::sys::uefi::common;
+use crate::alloc::{handle_alloc_error, GlobalAlloc, Layout, System};
 
 pub(crate) const POOL_ALIGNMENT: usize = 8;
 
@@ -13,87 +10,25 @@ const MEMORY_TYPE: u32 = r_efi::efi::LOADER_DATA;
 #[stable(feature = "alloc_system_type", since = "1.28.0")]
 unsafe impl GlobalAlloc for System {
     unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
-        let align = layout.align();
-        let size = layout.size();
-
-        // Return NULL pointer if boot_services pointer cannot be obtained. The only time this
-        // should happen is if SystemTable has not been initialized
-        let boot_services = match common::try_boot_services() {
-            Some(x) => x,
-            None => return ptr::null_mut(),
+        let system_table = match crate::os::uefi::env::try_system_table() {
+            None => return crate::ptr::null_mut(),
+            Some(x) => x.as_ptr() as *mut _,
         };
 
-        let allocate_pool_ptr = unsafe { (*boot_services.as_ptr()).allocate_pool };
-
-        let mut ptr: *mut crate::ffi::c_void = ptr::null_mut();
-        let aligned_size = align_size(size, align);
-
-        let r = (allocate_pool_ptr)(MEMORY_TYPE, aligned_size, &mut ptr);
-
-        if r.is_error() || ptr.is_null() {
-            return ptr::null_mut();
+        if layout.size() > 0 {
+            unsafe { r_efi_alloc::raw::alloc(system_table, layout, MEMORY_TYPE) }
+        } else {
+            layout.dangling().as_ptr()
         }
-
-        unsafe { align_ptr(ptr.cast(), align) }
     }
 
     unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
-        let boot_services = match common::try_boot_services() {
-            Some(x) => x,
-            None => return,
+        let system_table = match crate::os::uefi::env::try_system_table() {
+            None => handle_alloc_error(layout),
+            Some(x) => x.as_ptr() as *mut _,
         };
-
-        let free_pool_ptr = unsafe { (*boot_services.as_ptr()).free_pool };
-
-        let ptr = unsafe { unalign_ptr(ptr, layout.align()) };
-        let r = (free_pool_ptr)(ptr.cast());
-
-        assert!(!r.is_error());
-    }
-}
-
-#[inline]
-fn align_size(size: usize, align: usize) -> usize {
-    if align > POOL_ALIGNMENT {
-        // Allocate extra padding in order to be able to satisfy the alignment.
-        size + align
-    } else {
-        size
-    }
-}
-
-#[repr(C)]
-struct Header(*mut u8);
-
-#[inline]
-unsafe fn align_ptr(ptr: *mut u8, align: usize) -> *mut u8 {
-    if align > POOL_ALIGNMENT {
-        let offset = ptr.align_offset(align);
-
-        // SAFETY: `MIN_ALIGN` <= `offset` <= `layout.align()` and the size of the allocated
-        // block is `layout.align() + layout.size()`. `aligned` will thus be a correctly aligned
-        // pointer inside the allocated block with at least `layout.size()` bytes after it and at
-        // least `MIN_ALIGN` bytes of padding before it.
-        let aligned = unsafe { ptr.add(offset) };
-
-        // SAFETY: Because the size and alignment of a header is <= `MIN_ALIGN` and `aligned`
-        // is aligned to at least `MIN_ALIGN` and has at least `MIN_ALIGN` bytes of padding before
-        // it, it is safe to write a header directly before it.
-        unsafe { crate::ptr::write((aligned as *mut Header).offset(-1), Header(ptr)) };
-
-        aligned
-    } else {
-        ptr
-    }
-}
-
-#[inline]
-unsafe fn unalign_ptr(ptr: *mut u8, align: usize) -> *mut u8 {
-    if align > POOL_ALIGNMENT {
-        // SAFETY: Because of the contract of `System`, `ptr` is guaranteed to be non-null
-        // and have a header readable directly before it.
-        unsafe { crate::ptr::read((ptr as *mut Header).offset(-1)).0 }
-    } else {
-        ptr
+        if layout.size() > 0 {
+            unsafe { r_efi_alloc::raw::dealloc(system_table, ptr, layout) }
+        }
     }
 }

src/tools/tidy/src/deps.rs

@@ -273,6 +273,7 @@ const PERMITTED_RUSTC_DEPENDENCIES: &[&str] = &[
     "zerovec",
     "zerovec-derive",
     "r-efi",
+    "r-efi-alloc",
 ];
 
 const PERMITTED_CRANELIFT_DEPENDENCIES: &[&str] = &[