UEFI Allocator: add PAGE_SIZE shortcut

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

@@ -1,74 +1,101 @@
 // SPDX-License-Identifier: MIT OR Apache-2.0
 
 use alloc::vec::Vec;
-use uefi::boot::{self, AllocateType};
-use uefi::mem::memory_map::{MemoryMap, MemoryMapMut, MemoryType};
+use uefi::boot;
+use uefi::mem::memory_map::{MemoryMap, MemoryMapMut};
+use uefi_raw::table::boot::MemoryType;
 
 pub fn test() {
     info!("Testing memory functions");
 
-    test_allocate_pages();
-    test_allocate_pool();
+    bootservices::allocate_pages();
+    bootservices::allocate_pool();
 
-    vec_alloc();
-    alloc_alignment();
+    global::alloc_vec();
+    global::alloc_alignment();
 
     test_memory_map();
 }
 
-fn test_allocate_pages() {
-    let num_pages = 1;
-    let ptr =
-        boot::allocate_pages(AllocateType::AnyPages, MemoryType::LOADER_DATA, num_pages).unwrap();
-    let addr = ptr.as_ptr() as usize;
-    assert_eq!(addr % 4096, 0, "Page pointer is not page-aligned");
+/// Tests that directly use UEFI boot services to allocate memory.
+mod bootservices {
+    use uefi::boot;
+    use uefi::boot::AllocateType;
+    use uefi_raw::table::boot::MemoryType;
+
+    /// Tests the `allocate_pages` boot service.
+    pub fn allocate_pages() {
+        let num_pages = 1;
+        let ptr = boot::allocate_pages(AllocateType::AnyPages, MemoryType::LOADER_DATA, num_pages)
+            .unwrap();
+        let addr = ptr.as_ptr() as usize;
+        assert_eq!(addr % 4096, 0, "Page pointer is not page-aligned");
+
+        // Verify the page can be written to.
+        {
+            let ptr = ptr.as_ptr();
+            unsafe { ptr.write_volatile(0xff) };
+            unsafe { ptr.add(4095).write_volatile(0xff) };
+        }
 
-    // Verify the page can be written to.
-    {
-        let ptr = ptr.as_ptr();
-        unsafe { ptr.write_volatile(0xff) };
-        unsafe { ptr.add(4095).write_volatile(0xff) };
+        unsafe { boot::free_pages(ptr, num_pages) }.unwrap();
     }
 
-    unsafe { boot::free_pages(ptr, num_pages) }.unwrap();
-}
-
-fn test_allocate_pool() {
-    let ptr = boot::allocate_pool(MemoryType::LOADER_DATA, 10).unwrap();
+    /// Tests the `allocate_pool` boot service.
+    pub fn allocate_pool() {
+        let ptr = boot::allocate_pool(MemoryType::LOADER_DATA, 10).unwrap();
 
-    // Verify the allocation can be written to.
-    {
-        let ptr = ptr.as_ptr();
-        unsafe { ptr.write_volatile(0xff) };
-        unsafe { ptr.add(9).write_volatile(0xff) };
+        // Verify the allocation can be written to.
+        {
+            let ptr = ptr.as_ptr();
+            unsafe { ptr.write_volatile(0xff) };
+            unsafe { ptr.add(9).write_volatile(0xff) };
+        }
+        unsafe { boot::free_pool(ptr) }.unwrap();
     }
-    unsafe { boot::free_pool(ptr) }.unwrap();
 }
 
-// Simple test to ensure our custom allocator works with the `alloc` crate.
-fn vec_alloc() {
-    info!("Allocating a vector through the `alloc` crate");
+/// 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;
 
-    #[allow(clippy::useless_vec)]
-    let mut values = vec![-5, 16, 23, 4, 0];
+    /// Simple test to ensure our custom allocator works with the `alloc` crate.
+    pub fn alloc_vec() {
+        info!("Allocating a vector using the global allocator");
 
-    values.sort_unstable();
+        #[allow(clippy::useless_vec)]
+        let mut values = vec![-5, 16, 23, 4, 0];
 
-    assert_eq!(values[..], [-5, 0, 4, 16, 23], "Failed to sort vector");
-}
+        values.sort_unstable();
 
-// Simple test to ensure our custom allocator works with correct alignment.
-fn alloc_alignment() {
-    info!("Allocating a structure with alignment to 0x100");
+        assert_eq!(values[..], [-5, 0, 4, 16, 23], "Failed to sort vector");
+    }
 
-    #[repr(align(0x100))]
-    struct Block(
-        // Ignore warning due to field not being read.
-        #[allow(dead_code)] [u8; 0x100],
-    );
+    /// 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 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"
+            );
+        }
+    }
 }
 
 fn test_memory_map() {

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)

uefi/src/allocator.rs

@@ -1,20 +1,19 @@
 // SPDX-License-Identifier: MIT OR Apache-2.0
 
-//! This module implements Rust's global allocator interface using UEFI's memory allocation functions.
+//! This module exports [`Allocator`].
 //!
-//! If the `global_allocator` feature is enabled, the [`Allocator`] will be used
-//! as the global Rust allocator.
+//! The allocator can be used as global Rust allocator using the
+//! `global_allocator` crate feature. See [`helpers`] for more info.
 //!
-//! This allocator can only be used while boot services are active. If boot
-//! services are not active, `alloc` will return a null pointer, and `dealloc`
-//! will panic.
+//! [`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.
 ///
@@ -42,15 +41,69 @@ fn get_memory_type() -> MemoryType {
     }
 }
 
-/// Allocator which uses the UEFI pool allocation functions.
+/// Helper to get a custom alignment out of an allocation with an alignment of
+/// eight (UEFI default alignment). This works by allocating extra space and
+/// storing a pointer to the actual allocation right above the allocation
+/// handed out via the public API.
+fn alloc_pool_aligned(memory_type: MemoryType, size: usize, align: usize) -> *mut u8 {
+    let full_alloc_ptr = boot::allocate_pool(memory_type, size + align);
+    let full_alloc_ptr = if let Ok(ptr) = full_alloc_ptr {
+        ptr.as_ptr()
+    } else {
+        return ptr::null_mut();
+    };
+
+    // Calculate the offset needed to get an aligned pointer within the
+    // full allocation. If that offset is zero, increase it to `align`
+    // so that we still have space to store the extra pointer described
+    // below.
+    let mut offset = full_alloc_ptr.align_offset(align);
+    if offset == 0 {
+        offset = align;
+    }
+
+    // Before returning the aligned allocation, store a pointer to the
+    // full unaligned allocation in the bytes just before the aligned
+    // allocation. We know we have at least eight bytes there due to
+    // adding `align` to the memory allocation size. We also know the
+    // write is appropriately aligned for a `*mut u8` pointer because
+    // `align_ptr` is aligned, and alignments are always powers of two
+    // (as enforced by the `Layout` type).
+    unsafe {
+        let aligned_ptr = full_alloc_ptr.add(offset);
+        (aligned_ptr.cast::<*mut u8>()).sub(1).write(full_alloc_ptr);
+        aligned_ptr
+    }
+}
+
+/// 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
+/// `#[global_allocator]` attribute to be used as global Rust allocator.
 ///
-/// Only valid for as long as the UEFI boot services are available.
+/// Note that if boot services are not active (anymore), [`Allocator::alloc`]
+/// will return a null pointer and [`Allocator::dealloc`] will panic.
 #[derive(Debug)]
 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
@@ -59,64 +112,60 @@ 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);
 
-        if align > 8 {
-            // The requested alignment is greater than 8, but `allocate_pool` is
-            // only guaranteed to provide eight-byte alignment. Allocate extra
-            // space so that we can return an appropriately-aligned pointer
-            // within the allocation.
-            let full_alloc_ptr = if let Ok(ptr) = boot::allocate_pool(memory_type, size + align) {
-                ptr.as_ptr()
-            } else {
-                return ptr::null_mut();
-            };
-
-            // Calculate the offset needed to get an aligned pointer within the
-            // full allocation. If that offset is zero, increase it to `align`
-            // so that we still have space to store the extra pointer described
-            // below.
-            let mut offset = full_alloc_ptr.align_offset(align);
-            if offset == 0 {
-                offset = align;
+        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())
             }
-
-            // Before returning the aligned allocation, store a pointer to the
-            // full unaligned allocation in the bytes just before the aligned
-            // allocation. We know we have at least eight bytes there due to
-            // adding `align` to the memory allocation size. We also know the
-            // write is appropriately aligned for a `*mut u8` pointer because
-            // `align_ptr` is aligned, and alignments are always powers of two
-            // (as enforced by the `Layout` type).
-            unsafe {
-                let aligned_ptr = full_alloc_ptr.add(offset);
-                (aligned_ptr.cast::<*mut u8>()).sub(1).write(full_alloc_ptr);
-                aligned_ptr
+            (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, layout.size())
+                    .map(|ptr| ptr.as_ptr())
+                    .unwrap_or(ptr::null_mut())
             }
-        } else {
-            // 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)
-                .map(|ptr| ptr.as_ptr())
-                .unwrap_or(ptr::null_mut())
+            (false, 9..) => alloc_pool_aligned(memory_type, layout.size(), layout.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.
+    /// Deallocate memory using the UEFI boot services.
+    ///
+    /// This will panic after exiting boot services.
+    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();
+        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();
+            }
+            (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 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();
+            }
+        }
     }
 }