UEFI Allocator: Implement cores allocator_api

uefi-test-runner/src/boot/memory.rs

@@ -1,26 +1,28 @@
 // SPDX-License-Identifier: MIT OR Apache-2.0
 
-use alloc::vec::Vec;
-use uefi::boot;
-use uefi::mem::memory_map::{MemoryMap, MemoryMapMut};
-use uefi_raw::table::boot::MemoryType;
-
 pub fn test() {
     info!("Testing memory functions");
 
     bootservices::allocate_pages();
     bootservices::allocate_pool();
+    bootservices::test_memory_map();
 
     global::alloc_vec();
     global::alloc_alignment();
 
-    test_memory_map();
+    #[cfg(feature = "unstable")]
+    {
+        allocator_api::alloc_zst();
+        allocator_api::alloc_normal();
+    }
 }
 
 /// Tests that directly use UEFI boot services to allocate memory.
 mod bootservices {
+    use alloc::vec::Vec;
     use uefi::boot;
     use uefi::boot::AllocateType;
+    use uefi::mem::memory_map::{MemoryMap, MemoryMapMut};
     use uefi_raw::table::boot::MemoryType;
 
     /// Tests the `allocate_pages` boot service.
@@ -53,6 +55,44 @@ mod bootservices {
         }
         unsafe { boot::free_pool(ptr) }.unwrap();
     }
+
+    /// Tests getting the memory map.
+    pub fn test_memory_map() {
+        info!("Testing memory map functions");
+
+        let mut memory_map =
+            boot::memory_map(MemoryType::LOADER_DATA).expect("Failed to retrieve UEFI memory map");
+
+        memory_map.sort();
+
+        // Collect the descriptors into a vector
+        let descriptors = memory_map.entries().copied().collect::<Vec<_>>();
+
+        // Ensured we have at least one entry.
+        // Real memory maps usually have dozens of entries.
+        assert!(!descriptors.is_empty(), "Memory map is empty");
+
+        let mut curr_value = descriptors[0];
+
+        for value in descriptors.iter().skip(1) {
+            if value.phys_start <= curr_value.phys_start {
+                panic!("memory map sorting failed");
+            }
+            curr_value = *value;
+        }
+
+        // This is pretty much a basic sanity test to ensure returned memory
+        // isn't filled with random values.
+        let first_desc = descriptors[0];
+
+        #[cfg(target_arch = "x86_64")]
+        {
+            let phys_start = first_desc.phys_start;
+            assert_eq!(phys_start, 0, "Memory does not start at address 0");
+        }
+        let page_count = first_desc.page_count;
+        assert!(page_count != 0, "Memory map entry has size zero");
+    }
 }
 
 /// Tests that use [`uefi::allocator::Allocator`], which is configured as the
@@ -98,39 +138,37 @@ mod global {
     }
 }
 
-fn test_memory_map() {
-    info!("Testing memory map functions");
-
-    let mut memory_map =
-        boot::memory_map(MemoryType::LOADER_DATA).expect("Failed to retrieve UEFI memory map");
-
-    memory_map.sort();
-
-    // Collect the descriptors into a vector
-    let descriptors = memory_map.entries().copied().collect::<Vec<_>>();
-
-    // Ensured we have at least one entry.
-    // Real memory maps usually have dozens of entries.
-    assert!(!descriptors.is_empty(), "Memory map is empty");
-
-    let mut curr_value = descriptors[0];
-
-    for value in descriptors.iter().skip(1) {
-        if value.phys_start <= curr_value.phys_start {
-            panic!("memory map sorting failed");
-        }
-        curr_value = *value;
+/// Tests the `allocator_api` on the UEFI allocator.
+#[cfg(feature = "unstable")]
+mod allocator_api {
+    use core::alloc::Layout;
+    use uefi::allocator::Allocator;
+
+    pub fn alloc_zst() {
+        let layout = Layout::from_size_align(0, 1024).unwrap();
+        let ptr = <Allocator as core::alloc::Allocator>::allocate(&Allocator, layout).unwrap();
+        assert_eq!(ptr.len(), 0);
+        assert_eq!(ptr.as_ptr().cast::<u8>().align_offset(layout.align()), 0);
     }
 
-    // This is pretty much a basic sanity test to ensure returned memory
-    // isn't filled with random values.
-    let first_desc = descriptors[0];
-
-    #[cfg(target_arch = "x86_64")]
-    {
-        let phys_start = first_desc.phys_start;
-        assert_eq!(phys_start, 0, "Memory does not start at address 0");
+    pub fn alloc_normal() {
+        let layout = Layout::from_size_align(64, 64).unwrap();
+        let allocation =
+            <Allocator as core::alloc::Allocator>::allocate(&Allocator, layout).unwrap();
+        assert_eq!(allocation.len(), 64);
+        assert_eq!(
+            allocation
+                .as_ptr()
+                .cast::<u8>()
+                .align_offset(layout.align()),
+            0
+        );
+
+        unsafe {
+            core::ptr::write_bytes(allocation.as_ptr().cast::<u8>(), 42, allocation.len());
+        }
+        unsafe {
+            assert_eq!(allocation.as_ref()[42], 42);
+        }
     }
-    let page_count = first_desc.page_count;
-    assert!(page_count != 0, "Memory map entry has size zero");
 }

uefi-test-runner/src/main.rs

@@ -2,6 +2,7 @@
 
 #![no_std]
 #![no_main]
+#![cfg_attr(feature = "unstable", feature(allocator_api))]
 
 #[macro_use]
 extern crate log;

uefi/CHANGELOG.md

@@ -41,8 +41,10 @@
 - The `Display` impl for `CStr8` now excludes the trailing null character.
 - `VariableKeys` initializes with a larger name buffer to work around firmware
   bugs on some devices.
-- The UEFI `allocator::Allocator` has been optimized for page-aligned 
+- The UEFI `allocator::Allocator` has been optimized for page-aligned
   allocations.
+- The UEFI `allocator::Allocator` now implements `core::alloc::Allocator`
+  (`allocator_api`), when the `unstable` feature is used.
 
 
 # uefi - 0.34.1 (2025-02-07)

uefi/src/allocator.rs

@@ -169,3 +169,38 @@ unsafe impl GlobalAlloc for Allocator {
         }
     }
 }
+
+#[cfg(feature = "unstable")]
+unsafe impl core::alloc::Allocator for Allocator {
+    fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, core::alloc::AllocError> {
+        // Stable alternative for Layout::dangling()
+        fn dangling_for_layout(layout: Layout) -> NonNull<u8> {
+            let align = layout.align();
+            // SAFETY: align is non-zero, so align as usize is a valid address for a NonNull.
+            unsafe {
+                let ptr = align as *mut u8;
+                NonNull::new_unchecked(ptr)
+            }
+        }
+
+        match layout.size() {
+            0 => Ok(NonNull::slice_from_raw_parts(
+                dangling_for_layout(layout),
+                0,
+            )),
+            // SAFETY: `layout` is non-zero in size,
+            size => {
+                let ptr = unsafe { <Allocator as GlobalAlloc>::alloc(self, layout) };
+                NonNull::new(ptr)
+                    .ok_or(core::alloc::AllocError)
+                    .map(|ptr| NonNull::slice_from_raw_parts(ptr, size))
+            }
+        }
+    }
+
+    unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
+        if layout.size() != 0 {
+            unsafe { <Allocator as GlobalAlloc>::dealloc(self, ptr.as_ptr(), layout) }
+        }
+    }
+}

uefi/src/lib.rs

@@ -212,7 +212,7 @@
 //! [uefi-std-tr-issue]: https://github.com/rust-lang/rust/issues/100499
 //! [unstable features]: https://doc.rust-lang.org/unstable-book/
 
-#![cfg_attr(all(feature = "unstable", feature = "alloc"), feature(allocator_api))]
+#![cfg_attr(feature = "unstable", feature(allocator_api))]
 #![cfg_attr(docsrs, feature(doc_auto_cfg))]
 #![no_std]
 #![deny(