1use super::global::SYSTEM_TABLE;
5use super::MemoryType;
6use core::alloc::{GlobalAlloc, Layout};
7use core::ffi::c_void;
8use core::ptr;
9
10#[allow(clippy::module_name_repetitions)]
11pub struct BootAllocator;
12
13impl BootAllocator {
14 const UEFI_ALIGNMENT: usize = 8;
16
17 const VOID_PTR_LAYOUT: Layout = Layout::new::<*const c_void>();
18
19 fn extend_layout_for_alignment(orig_layout: Layout) -> (Layout, usize) {
20 let (saved_ptr_layout, saved_ptr_offset) = orig_layout
26 .extend(Self::VOID_PTR_LAYOUT)
27 .expect("Could not construct extended layout for alignment structure");
28 let adjusted_size = saved_ptr_layout.size() + saved_ptr_layout.align();
29 let adjusted_layout =
30 Layout::from_size_align(adjusted_size, saved_ptr_layout.align())
31 .expect("Could not construct extended layout for alignable allocation");
32 (adjusted_layout, saved_ptr_offset)
33 }
34
35 unsafe fn alloc(layout: Layout) -> *mut u8 {
36 let system_table = SYSTEM_TABLE.expect("System table not initialized");
37 let boot_services =
38 (*system_table).boot_services.expect("Boot services unavailable");
39
40 let (adjusted_layout, saved_ptr_offset) =
41 if layout.align() <= Self::UEFI_ALIGNMENT {
42 (layout, 0)
44 } else {
45 Self::extend_layout_for_alignment(layout)
46 };
47
48 let mut buffer: *mut c_void = ptr::null_mut();
49 (boot_services.allocate_pool)(
50 MemoryType::LoaderData,
51 adjusted_layout.size(),
52 &raw mut buffer,
53 )
54 .into_result()
55 .expect("allocate_pool() failed");
56
57 if adjusted_layout.align() > Self::UEFI_ALIGNMENT {
58 let orig_address = buffer as usize;
59 let align_offset = orig_address % adjusted_layout.align();
60 if align_offset != 0 {
61 let shifted_ptr = orig_address + (adjusted_layout.align() - align_offset);
63 buffer = shifted_ptr as *mut c_void;
64 }
65 let saved_ptr_ptr = (buffer as usize + saved_ptr_offset) as *mut *mut c_void;
69 *saved_ptr_ptr = orig_address as *mut c_void;
70 }
71
72 buffer.cast()
73 }
74
75 unsafe fn dealloc(buffer: *mut u8, layout: Layout) {
76 let system_table = SYSTEM_TABLE.expect("System table not initialized");
77 let boot_services =
78 (*system_table).boot_services.expect("Boot services unavailable");
79
80 let original_ptr: *mut c_void = if layout.align() <= Self::UEFI_ALIGNMENT {
81 buffer.cast()
83 } else {
84 let (_, saved_ptr_offset) = Self::extend_layout_for_alignment(layout);
88 let saved_ptr_ptr = (buffer as usize + saved_ptr_offset) as *mut *mut c_void;
89 *saved_ptr_ptr
90 };
91
92 (boot_services.free_pool)(original_ptr)
93 .into_result()
94 .expect("Could not free memory");
95 }
96}
97
98unsafe impl GlobalAlloc for BootAllocator {
99 unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
100 BootAllocator::alloc(layout)
101 }
102
103 unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
104 BootAllocator::dealloc(ptr, layout);
105 }
106}