@@ -63,60 +63,70 @@ unsafe impl GlobalAlloc for Allocator {
63
63
let align = layout. align ( ) ;
64
64
let memory_type = get_memory_type ( ) ;
65
65
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 ( ) )
84
74
}
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
+ }
85
96
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
+ }
97
109
}
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 ( ) )
105
110
}
106
111
}
107
112
108
113
/// 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
+ }
114
130
}
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 ( ) ;
121
131
}
122
132
}
0 commit comments