|
12 | 12 | //! will panic. |
13 | 13 | |
14 | 14 | use crate::boot; |
| 15 | +use crate::boot::AllocateType; |
15 | 16 | use crate::mem::memory_map::MemoryType; |
16 | 17 | use crate::proto::loaded_image::LoadedImage; |
17 | 18 | use core::alloc::{GlobalAlloc, Layout}; |
18 | 19 | use core::ptr::{self, NonNull}; |
19 | 20 | use core::sync::atomic::{AtomicU32, Ordering}; |
| 21 | +use uefi_raw::table::boot::PAGE_SIZE; |
20 | 22 |
|
21 | 23 | /// Get the memory type to use for allocation. |
22 | 24 | /// |
@@ -79,7 +81,7 @@ fn alloc_pool_aligned(memory_type: MemoryType, size: usize, align: usize) -> *mu |
79 | 81 | } |
80 | 82 | } |
81 | 83 |
|
82 | | -/// Allocator which uses the UEFI pool allocation functions. |
| 84 | +/// Allocator which uses the UEFI allocation functions. |
83 | 85 | /// |
84 | 86 | /// The allocator can only be used as long as the UEFI boot services are |
85 | 87 | /// available and have not been exited. |
@@ -110,24 +112,51 @@ unsafe impl GlobalAlloc for Allocator { |
110 | 112 | .map(|ptr| ptr.as_ptr()) |
111 | 113 | .unwrap_or(ptr::null_mut()) |
112 | 114 | } |
| 115 | + // Allocating pages is actually very expected in UEFI OS loaders, so |
| 116 | + // it makes sense to provide this optimization. |
| 117 | + // In Rust, each allocation's size must be at least the alignment, |
| 118 | + // as specified in the Rust type layout [0]. Therefore, we don't |
| 119 | + // have a risk of wasting memory. Further, using page-alignment for |
| 120 | + // small allocations is quiet unusual. |
| 121 | + // [0]: https://doc.rust-lang.org/reference/type-layout.html |
| 122 | + PAGE_SIZE => { |
| 123 | + log::trace!("Taking PAGE_SIZE shortcut"); |
| 124 | + let count = 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 | + } |
113 | 129 | 9.. => { |
114 | 130 | alloc_pool_aligned(memory_type, size, align) |
115 | 131 | } |
116 | 132 | } |
117 | 133 | } |
118 | 134 |
|
119 | 135 | /// Deallocate memory using [`boot::free_pool`]. |
120 | | - unsafe fn dealloc(&self, mut ptr: *mut u8, layout: Layout) { |
121 | | - if layout.align() > 8 { |
122 | | - // Retrieve the pointer to the full allocation that was packed right |
123 | | - // before the aligned allocation in `alloc`. |
124 | | - ptr = unsafe { (ptr as *const *mut u8).sub(1).read() }; |
125 | | - } |
126 | | - |
127 | | - // OK to unwrap: `ptr` is required to be a valid allocation by the trait API. |
| 136 | + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { |
128 | 137 | let ptr = NonNull::new(ptr).unwrap(); |
129 | | - |
130 | | - // Warning: this will panic after exiting boot services. |
131 | | - unsafe { boot::free_pool(ptr) }.unwrap(); |
| 138 | + match layout.align() { |
| 139 | + 0..=8 /* UEFI default alignment */ => { |
| 140 | + // Warning: this will panic after exiting boot services. |
| 141 | + unsafe { boot::free_pool(ptr) }.unwrap(); |
| 142 | + } |
| 143 | + /* Corresponds the the alloc() match arm. */ |
| 144 | + PAGE_SIZE => { |
| 145 | + log::trace!("Taking PAGE_SIZE shortcut"); |
| 146 | + let count = layout.size().div_ceil(PAGE_SIZE); |
| 147 | + unsafe { |
| 148 | + boot::free_pages(ptr, count).unwrap() |
| 149 | + } |
| 150 | + } |
| 151 | + 9.. => { |
| 152 | + let ptr = ptr.as_ptr().cast::<*mut u8>(); |
| 153 | + // Retrieve the pointer to the full allocation that was packed right |
| 154 | + // before the aligned allocation in `alloc`. |
| 155 | + let actual_alloc_ptr = unsafe { ptr.sub(1).read() }; |
| 156 | + let ptr = NonNull::new(actual_alloc_ptr).unwrap(); |
| 157 | + // Warning: this will panic after exiting boot services. |
| 158 | + unsafe { boot::free_pool(ptr) }.unwrap(); |
| 159 | + } |
| 160 | + } |
132 | 161 | } |
133 | 162 | } |
0 commit comments