Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly
Sign up
Appearance settings

Commit 3b0c930

Browse files
committed
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(). We can look at the TRACE messages of `cargo xtask run` to verify that the shortcut is taken.
1 parent bd3e22a commit 3b0c930

File tree

3 files changed

+78
-29
lines changed

3 files changed

+78
-29
lines changed

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

Lines changed: 23 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -58,9 +58,12 @@ mod bootservices {
5858
/// Tests that use [`uefi::allocator::Allocator`], which is configured as the
5959
/// global allocator.
6060
mod global {
61+
use alloc::boxed::Box;
62+
use uefi_raw::table::boot::PAGE_SIZE;
63+
6164
/// Simple test to ensure our custom allocator works with the `alloc` crate.
6265
pub fn alloc_vec() {
63-
info!("Allocating a vector through the `alloc` crate");
66+
info!("Allocating a vector using the global allocator");
6467

6568
#[allow(clippy::useless_vec)]
6669
let mut values = vec![-5, 16, 23, 4, 0];
@@ -71,17 +74,27 @@ mod global {
7174
}
7275

7376
/// Simple test to ensure our custom allocator works with correct alignment.
77+
#[allow(dead_code)] // Ignore warning due to field not being read.
7478
pub fn alloc_alignment() {
75-
info!("Allocating a structure with alignment to 0x100");
76-
77-
#[repr(align(0x100))]
78-
struct Block(
79-
// Ignore warning due to field not being read.
80-
#[allow(dead_code)] [u8; 0x100],
81-
);
79+
{
80+
info!("Allocating a structure with alignment of 0x100 using the global allocator");
81+
#[repr(align(0x100))]
82+
struct Block([u8; 0x100]);
8283

83-
let value = vec![Block([1; 0x100])];
84-
assert_eq!(value.as_ptr() as usize % 0x100, 0, "Wrong alignment");
84+
let value = vec![Block([1; 0x100])];
85+
assert_eq!(value.as_ptr() as usize % 0x100, 0, "Wrong alignment");
86+
}
87+
{
88+
info!("Allocating a memory page ({PAGE_SIZE}) using the global allocator");
89+
#[repr(align(4096))]
90+
struct Page([u8; PAGE_SIZE]);
91+
let value = Box::new(Page([0; PAGE_SIZE]));
92+
assert_eq!(
93+
value.0.as_ptr().align_offset(PAGE_SIZE),
94+
0,
95+
"Wrong alignment"
96+
);
97+
}
8598
}
8699
}
87100

‎uefi/CHANGELOG.md

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,8 @@
3737
- The `Display` impl for `CStr8` now excludes the trailing null character.
3838
- `VariableKeys` initializes with a larger name buffer to work around firmware
3939
bugs on some devices.
40+
- The UEFI `allocator::Allocator` has been optimized for page-aligned
41+
allocations.
4042

4143

4244
# uefi - 0.34.1 (2025年02月07日)

‎uefi/src/allocator.rs

Lines changed: 53 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -7,12 +7,13 @@
77
//!
88
//! [`helpers`]: uefi::helpers
99
10-
use crate::boot;
10+
use crate::boot::{self,AllocateType};
1111
use crate::mem::memory_map::MemoryType;
1212
use crate::proto::loaded_image::LoadedImage;
1313
use core::alloc::{GlobalAlloc, Layout};
1414
use core::ptr::{self, NonNull};
1515
use core::sync::atomic::{AtomicU32, Ordering};
16+
use uefi_raw::table::boot::PAGE_SIZE;
1617

1718
/// Get the memory type to use for allocation.
1819
///
@@ -75,6 +76,20 @@ fn alloc_pool_aligned(memory_type: MemoryType, size: usize, align: usize) -> *mu
7576
}
7677
}
7778

79+
/// Returns whether the allocation is a multiple of a [`PAGE_SIZE`] and is
80+
/// aligned to [`PAGE_SIZE`].
81+
///
82+
/// This does not only check the alignment but also the size. For types
83+
/// allocated by Rust itself (e.g., `Box<T>`), the size is always at least the
84+
/// alignment, as specified in the [Rust type layout]. However, to be also safe
85+
/// when it comes to manual invocations, we additionally check if the size is
86+
/// a multiple of [`PAGE_SIZE`].
87+
///
88+
/// [Rust type layout]: https://doc.rust-lang.org/reference/type-layout.html
89+
fn layout_allows_page_alloc_shortcut(layout: &Layout) -> bool {
90+
layout.size() % PAGE_SIZE == 0 && layout.align() == PAGE_SIZE
91+
}
92+
7893
/// Allocator using UEFI boot services.
7994
///
8095
/// 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
86101
pub struct Allocator;
87102

88103
unsafe impl GlobalAlloc for Allocator {
89-
/// Allocate memory using [`boot::allocate_pool`]. The allocation's [memory
90-
/// type] matches the current image's [data type].
104+
/// Allocate memory using the UEFI boot services.
105+
///
106+
/// The allocation's [memory type] matches the current image's [data type].
91107
///
92108
/// [memory type]: MemoryType
93109
/// [data type]: LoadedImage::data_type
@@ -96,40 +112,58 @@ unsafe impl GlobalAlloc for Allocator {
96112
return ptr::null_mut();
97113
}
98114

99-
let size = layout.size();
100-
let align = layout.align();
101115
let memory_type = get_memory_type();
116+
let use_page_shortcut = layout_allows_page_alloc_shortcut(&layout);
102117

103-
match align {
104-
0..=8 /* UEFI default alignment */ => {
118+
match (use_page_shortcut, layout.align()) {
119+
// Allocating pages is actually very expected in UEFI OS loaders, so
120+
// it makes sense to provide this optimization.
121+
(true, _) => {
122+
// To spammy, but useful for manual testing.
123+
// log::trace!("Taking PAGE_SIZE shortcut for layout={layout:?}");
124+
let count = layout.size().div_ceil(PAGE_SIZE);
125+
boot::allocate_pages(AllocateType::AnyPages, memory_type, count)
126+
.map(|ptr| ptr.as_ptr())
127+
.unwrap_or(ptr::null_mut())
128+
}
129+
(false, 0..=8 /* UEFI default alignment */) => {
105130
// The requested alignment is less than or equal to eight, and
106131
// `allocate_pool` always provides eight-byte alignment, so we can
107132
// use `allocate_pool` directly.
108-
boot::allocate_pool(memory_type, size)
133+
boot::allocate_pool(memory_type, layout.size())
109134
.map(|ptr| ptr.as_ptr())
110135
.unwrap_or(ptr::null_mut())
111136
}
112-
9.. => {
113-
alloc_pool_aligned(memory_type, size, align)
114-
}
137+
(false, 9..) => alloc_pool_aligned(memory_type, layout.size(), layout.align()),
115138
}
116139
}
117140

118-
/// Deallocate memory using [`boot::free_pool`].
141+
/// Deallocate memory using the UEFI boot services.
119142
///
120143
/// This will panic after exiting boot services.
121144
unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
122-
match layout.align() {
123-
0..=8 => {
124-
// OK to unwrap: `ptr` is required to be a valid allocation by the trait API.
125-
let ptr = NonNull::new(ptr).unwrap();
145+
let ptr = NonNull::new(ptr).unwrap();
146+
147+
let use_page_shortcut = layout_allows_page_alloc_shortcut(&layout);
148+
149+
match (use_page_shortcut, layout.align()) {
150+
(true, _) => {
151+
// To spammy, but useful for manual testing.
152+
// log::trace!("Taking PAGE_SIZE shortcut for layout={layout:?}");
153+
let count = layout.size().div_ceil(PAGE_SIZE);
154+
unsafe { boot::free_pages(ptr, count).unwrap() }
155+
}
156+
(false, 0..=8 /* UEFI default alignment */) => {
157+
// Warning: this will panic after exiting boot services.
126158
unsafe { boot::free_pool(ptr) }.unwrap();
127159
}
128-
9.. => {
160+
(false, 9..) => {
161+
let ptr = ptr.as_ptr().cast::<*mut u8>();
129162
// Retrieve the pointer to the full allocation that was packed right
130163
// before the aligned allocation in `alloc`.
131-
let ptr = unsafe { (ptr as *const *mut u8).sub(1).read() };
132-
let ptr = NonNull::new(ptr).unwrap();
164+
let actual_alloc_ptr = unsafe { ptr.sub(1).read() };
165+
let ptr = NonNull::new(actual_alloc_ptr).unwrap();
166+
// Warning: this will panic after exiting boot services.
133167
unsafe { boot::free_pool(ptr) }.unwrap();
134168
}
135169
}

0 commit comments

Comments
(0)

AltStyle によって変換されたページ (->オリジナル) /