uefi: allocator: use shortcut for PAGE_SIZE

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().

## Testing

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

Author: phip1611

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"
+            );
+        }
     }
 }
 

uefi/CHANGELOG.md

@@ -28,6 +28,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)

uefi/src/allocator.rs

@@ -12,11 +12,13 @@
 //! will panic.
 
 use crate::boot;
+use crate::boot::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.
 ///
@@ -79,7 +81,7 @@ fn alloc_pool_aligned(memory_type: MemoryType, size: usize, align: usize) -> *mu
     }
 }
 
-/// Allocator which uses the UEFI pool allocation functions.
+/// Allocator which uses the UEFI allocation functions.
 ///
 /// The allocator can only be used as long as the UEFI boot services are
 /// available and have not been exited.
@@ -110,24 +112,51 @@ unsafe impl GlobalAlloc for Allocator {
                     .map(|ptr| ptr.as_ptr())
                     .unwrap_or(ptr::null_mut())
             }
+            // Allocating pages is actually very expected in UEFI OS loaders, so
+            // it makes sense to provide this optimization.
+            // In Rust, each allocation's size must be at least the alignment,
+            // as specified in the Rust type layout [0]. Therefore, we don't
+            // have a risk of wasting memory. Further, using page-alignment for
+            // small allocations is quiet unusual.
+            // [0]: https://doc.rust-lang.org/reference/type-layout.html
+            PAGE_SIZE => {
+                log::trace!("Taking PAGE_SIZE shortcut");
+                let count = size.div_ceil(PAGE_SIZE);
+                boot::allocate_pages(AllocateType::AnyPages, memory_type, count)
+                    .map(|ptr| ptr.as_ptr())
+                    .unwrap_or(ptr::null_mut())
+            }
             9.. => {
                 alloc_pool_aligned(memory_type, size, align)
             }
         }
     }
 
     /// Deallocate memory using [`boot::free_pool`].
-    unsafe fn dealloc(&self, mut ptr: *mut u8, layout: Layout) {
-        if layout.align() > 8 {
-            // Retrieve the pointer to the full allocation that was packed right
-            // before the aligned allocation in `alloc`.
-            ptr = unsafe { (ptr as *const *mut u8).sub(1).read() };
-        }
-
-        // OK to unwrap: `ptr` is required to be a valid allocation by the trait API.
+    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
         let ptr = NonNull::new(ptr).unwrap();
-
-        // Warning: this will panic after exiting boot services.
-        unsafe { boot::free_pool(ptr) }.unwrap();
+        match layout.align() {
+            0..=8 /* UEFI default alignment */ => {
+                // Warning: this will panic after exiting boot services.
+                unsafe { boot::free_pool(ptr) }.unwrap();
+            }
+            /* Corresponds the the alloc() match arm. */
+            PAGE_SIZE => {
+                log::trace!("Taking PAGE_SIZE shortcut");
+                let count = layout.size().div_ceil(PAGE_SIZE);
+                unsafe {
+                    boot::free_pages(ptr, count).unwrap()
+                }
+            }
+            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 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();
+            }
+        }
     }
 }