uefi: allocator: use shortcut for PAGE_SIZE + add test

Allocating page-aligned memory via the global allocator is not
uncommon for UEFI OS loaders. Therefore, it is feasible to use a
shortcut in the allocator, and directly use boot::allocate_pages()
rather than boot::allocate_pool().

We can look at the TRACE messages of `cargo xtask run` to
verify that the shortcut is taken.

Author: phip1611

Diff for: uefi-test-runner/src/boot/memory.rs

@@ -58,9 +58,12 @@ mod bootservices {
 /// Tests that use [`uefi::allocator::Allocator`], which is configured as the
 /// global allocator.
 mod global {
+    use alloc::boxed::Box;
+    use uefi_raw::table::boot::PAGE_SIZE;
+
     /// Simple test to ensure our custom allocator works with the `alloc` crate.
     pub fn alloc_vec() {
-        info!("Allocating a vector through the `alloc` crate");
+        info!("Allocating a vector using the global allocator");
 
         #[allow(clippy::useless_vec)]
         let mut values = vec![-5, 16, 23, 4, 0];
@@ -71,17 +74,27 @@ mod global {
     }
 
     /// Simple test to ensure our custom allocator works with correct alignment.
+    #[allow(dead_code)] // Ignore warning due to field not being read.
     pub fn alloc_alignment() {
-        info!("Allocating a structure with alignment to 0x100");
-
-        #[repr(align(0x100))]
-        struct Block(
-            // Ignore warning due to field not being read.
-            #[allow(dead_code)] [u8; 0x100],
-        );
+        {
+            info!("Allocating a structure with alignment of 0x100 using the global allocator");
+            #[repr(align(0x100))]
+            struct Block([u8; 0x100]);
 
-        let value = vec![Block([1; 0x100])];
-        assert_eq!(value.as_ptr() as usize % 0x100, 0, "Wrong alignment");
+            let value = vec![Block([1; 0x100])];
+            assert_eq!(value.as_ptr() as usize % 0x100, 0, "Wrong alignment");
+        }
+        {
+            info!("Allocating a memory page ({PAGE_SIZE}) using the global allocator");
+            #[repr(align(4096))]
+            struct Page([u8; PAGE_SIZE]);
+            let value = Box::new(Page([0; PAGE_SIZE]));
+            assert_eq!(
+                value.0.as_ptr().align_offset(PAGE_SIZE),
+                0,
+                "Wrong alignment"
+            );
+        }
     }
 }
 

Diff for: uefi/CHANGELOG.md

@@ -37,6 +37,8 @@
 - 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 
+  allocations.
 
 
 # uefi - 0.34.1 (2025-02-07)

Diff for: uefi/src/allocator.rs

@@ -7,12 +7,13 @@
 //!
 //! [`helpers`]: uefi::helpers
 
-use crate::boot;
+use crate::boot::{self, AllocateType};
 use crate::mem::memory_map::MemoryType;
 use crate::proto::loaded_image::LoadedImage;
 use core::alloc::{GlobalAlloc, Layout};
 use core::ptr::{self, NonNull};
 use core::sync::atomic::{AtomicU32, Ordering};
+use uefi_raw::table::boot::PAGE_SIZE;
 
 /// Get the memory type to use for allocation.
 ///
@@ -75,6 +76,20 @@ fn alloc_pool_aligned(memory_type: MemoryType, size: usize, align: usize) -> *mu
     }
 }
 
+/// Returns whether the allocation is a multiple of a [`PAGE_SIZE`] and is
+/// aligned to [`PAGE_SIZE`].
+///
+/// This does not only check the alignment but also the size. For types
+/// allocated by Rust itself (e.g., `Box<T>`), the size is always at least the
+/// alignment, as specified in the [Rust type layout]. However, to be also safe
+/// when it comes to manual invocations, we additionally check if the size is
+/// a multiple of [`PAGE_SIZE`].
+///
+/// [Rust type layout]: https://doc.rust-lang.org/reference/type-layout.html
+const fn layout_allows_page_alloc_shortcut(layout: &Layout) -> bool {
+    layout.size() % PAGE_SIZE == 0 && layout.align() == PAGE_SIZE
+}
+
 /// Allocator using UEFI boot services.
 ///
 /// This type implements [`GlobalAlloc`] and can be marked with the
@@ -86,8 +101,9 @@ fn alloc_pool_aligned(memory_type: MemoryType, size: usize, align: usize) -> *mu
 pub struct Allocator;
 
 unsafe impl GlobalAlloc for Allocator {
-    /// Allocate memory using [`boot::allocate_pool`]. The allocation's [memory
-    /// type] matches the current image's [data type].
+    /// Allocate memory using the UEFI boot services.
+    ///
+    /// The allocation's [memory type] matches the current image's [data type].
     ///
     /// [memory type]: MemoryType
     /// [data type]: LoadedImage::data_type
@@ -96,40 +112,58 @@ unsafe impl GlobalAlloc for Allocator {
             return ptr::null_mut();
         }
 
-        let size = layout.size();
-        let align = layout.align();
         let memory_type = get_memory_type();
+        let use_page_shortcut = layout_allows_page_alloc_shortcut(&layout);
 
-        match align {
-            0..=8 /* UEFI default alignment */ => {
+        match (use_page_shortcut, layout.align()) {
+            // Allocating pages is actually very expected in UEFI OS loaders, so
+            // it makes sense to provide this optimization.
+            (true, _) => {
+                // To spammy, but useful for manual testing.
+                // log::trace!("Taking PAGE_SIZE shortcut for layout={layout:?}");
+                let count = layout.size().div_ceil(PAGE_SIZE);
+                boot::allocate_pages(AllocateType::AnyPages, memory_type, count)
+                    .map(|ptr| ptr.as_ptr())
+                    .unwrap_or(ptr::null_mut())
+            }
+            (false, 0..=8 /* UEFI default alignment */) => {
                 // The requested alignment is less than or equal to eight, and
                 // `allocate_pool` always provides eight-byte alignment, so we can
                 // use `allocate_pool` directly.
-                boot::allocate_pool(memory_type, size)
+                boot::allocate_pool(memory_type, layout.size())
                     .map(|ptr| ptr.as_ptr())
                     .unwrap_or(ptr::null_mut())
             }
-            9.. => {
-                alloc_pool_aligned(memory_type, size, align)
-            }
+            (false, 9..) => alloc_pool_aligned(memory_type, layout.size(), layout.align()),
         }
     }
 
-    /// Deallocate memory using [`boot::free_pool`].
+    /// Deallocate memory using the UEFI boot services.
     ///
     /// This will panic after exiting boot services.
     unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
-        match layout.align() {
-            0..=8 => {
-                // OK to unwrap: `ptr` is required to be a valid allocation by the trait API.
-                let ptr = NonNull::new(ptr).unwrap();
+        let ptr = NonNull::new(ptr).unwrap();
+
+        let use_page_shortcut = layout_allows_page_alloc_shortcut(&layout);
+
+        match (use_page_shortcut, layout.align()) {
+            (true, _) => {
+                // To spammy, but useful for manual testing.
+                // log::trace!("Taking PAGE_SIZE shortcut for layout={layout:?}");
+                let count = layout.size().div_ceil(PAGE_SIZE);
+                unsafe { boot::free_pages(ptr, count).unwrap() }
+            }
+            (false, 0..=8 /* UEFI default alignment */) => {
+                // Warning: this will panic after exiting boot services.
                 unsafe { boot::free_pool(ptr) }.unwrap();
             }
-            9.. => {
+            (false, 9..) => {
+                let ptr = ptr.as_ptr().cast::<*mut u8>();
                 // Retrieve the pointer to the full allocation that was packed right
                 // before the aligned allocation in `alloc`.
-                let ptr = unsafe { (ptr as *const *mut u8).sub(1).read() };
-                let ptr = NonNull::new(ptr).unwrap();
+                let actual_alloc_ptr = unsafe { ptr.sub(1).read() };
+                let ptr = NonNull::new(actual_alloc_ptr).unwrap();
+                // Warning: this will panic after exiting boot services.
                 unsafe { boot::free_pool(ptr) }.unwrap();
             }
         }