@@ -63,60 +63,70 @@ unsafe impl GlobalAlloc for Allocator {
6363 let align = layout. align ( ) ;
6464 let memory_type = get_memory_type ( ) ;
6565
66- if align > 8 {
67- // The requested alignment is greater than 8, but `allocate_pool` is
68- // only guaranteed to provide eight-byte alignment. Allocate extra
69- // space so that we can return an appropriately-aligned pointer
70- // within the allocation.
71- let full_alloc_ptr = if let Ok ( ptr) = boot:: allocate_pool ( memory_type, size + align) {
72- ptr. as_ptr ( )
73- } else {
74- return ptr:: null_mut ( ) ;
75- } ;
76- 77- // Calculate the offset needed to get an aligned pointer within the
78- // full allocation. If that offset is zero, increase it to `align`
79- // so that we still have space to store the extra pointer described
80- // below.
81- let mut offset = full_alloc_ptr. align_offset ( align) ;
82- if offset == 0 {
83- offset = align;
66+ match align {
67+ 0 ..=8 /* UEFI default alignment */ => {
68+ // The requested alignment is less than or equal to eight, and
69+ // `allocate_pool` always provides eight-byte alignment, so we can
70+ // use `allocate_pool` directly.
71+ boot:: allocate_pool ( memory_type, size)
72+ . map ( |ptr| ptr. as_ptr ( ) )
73+ . unwrap_or ( ptr:: null_mut ( ) )
8474 }
75+ 9 .. => {
76+ // The requested alignment is greater than 8, but `allocate_pool` is
77+ // only guaranteed to provide eight-byte alignment. Allocate extra
78+ // space so that we can return an appropriately-aligned pointer
79+ // within the allocation.
80+ let full_alloc_ptr = boot:: allocate_pool ( memory_type, size + align) ;
81+ let full_alloc_ptr = if let Ok ( ptr) = full_alloc_ptr
82+ {
83+ ptr. as_ptr ( )
84+ } else {
85+ return ptr:: null_mut ( ) ;
86+ } ;
87+ 88+ // Calculate the offset needed to get an aligned pointer within the
89+ // full allocation. If that offset is zero, increase it to `align`
90+ // so that we still have space to store the extra pointer described
91+ // below.
92+ let mut offset = full_alloc_ptr. align_offset ( align) ;
93+ if offset == 0 {
94+ offset = align;
95+ }
8596
86- // Before returning the aligned allocation, store a pointer to the
87- // full unaligned allocation in the bytes just before the aligned
88- // allocation. We know we have at least eight bytes there due to
89- // adding `align` to the memory allocation size. We also know the
90- // write is appropriately aligned for a `*mut u8` pointer because
91- // `align_ptr` is aligned, and alignments are always powers of two
92- // (as enforced by the `Layout` type).
93- unsafe {
94- let aligned_ptr = full_alloc_ptr. add ( offset) ;
95- ( aligned_ptr. cast :: < * mut u8 > ( ) ) . sub ( 1 ) . write ( full_alloc_ptr) ;
96- aligned_ptr
97+ // Before returning the aligned allocation, store a pointer to the
98+ // full unaligned allocation in the bytes just before the aligned
99+ // allocation. We know we have at least eight bytes there due to
100+ // adding `align` to the memory allocation size. We also know the
101+ // write is appropriately aligned for a `*mut u8` pointer because
102+ // `align_ptr` is aligned, and alignments are always powers of two
103+ // (as enforced by the `Layout` type).
104+ unsafe {
105+ let aligned_ptr = full_alloc_ptr. add ( offset) ;
106+ ( aligned_ptr. cast :: < * mut u8 > ( ) ) . sub ( 1 ) . write ( full_alloc_ptr) ;
107+ aligned_ptr
108+ }
97109 }
98- } else {
99- // The requested alignment is less than or equal to eight, and
100- // `allocate_pool` always provides eight-byte alignment, so we can
101- // use `allocate_pool` directly.
102- boot:: allocate_pool ( memory_type, size)
103- . map ( |ptr| ptr. as_ptr ( ) )
104- . unwrap_or ( ptr:: null_mut ( ) )
105110 }
106111 }
107112
108113 /// Deallocate memory using [`boot::free_pool`].
109- unsafe fn dealloc ( & self , mut ptr : * mut u8 , layout : Layout ) {
110- if layout. align ( ) > 8 {
111- // Retrieve the pointer to the full allocation that was packed right
112- // before the aligned allocation in `alloc`.
113- ptr = unsafe { ( ptr as * const * mut u8 ) . sub ( 1 ) . read ( ) } ;
114+ ///
115+ /// This will panic after exiting boot services.
116+ unsafe fn dealloc ( & self , ptr : * mut u8 , layout : Layout ) {
117+ match layout. align ( ) {
118+ 0 ..=8 => {
119+ // OK to unwrap: `ptr` is required to be a valid allocation by the trait API.
120+ let ptr = NonNull :: new ( ptr) . unwrap ( ) ;
121+ unsafe { boot:: free_pool ( ptr) } . unwrap ( ) ;
122+ }
123+ 9 .. => {
124+ // Retrieve the pointer to the full allocation that was packed right
125+ // before the aligned allocation in `alloc`.
126+ let ptr = unsafe { ( ptr as * const * mut u8 ) . sub ( 1 ) . read ( ) } ;
127+ let ptr = NonNull :: new ( ptr) . unwrap ( ) ;
128+ unsafe { boot:: free_pool ( ptr) } . unwrap ( ) ;
129+ }
114130 }
115- 116- // OK to unwrap: `ptr` is required to be a valid allocation by the trait API.
117- let ptr = NonNull :: new ( ptr) . unwrap ( ) ;
118- 119- // Warning: this will panic after exiting boot services.
120- unsafe { boot:: free_pool ( ptr) } . unwrap ( ) ;
121131 }
122132}
0 commit comments