@@ -52,38 +52,60 @@ pub fn exit_boot_services() {
52
52
pub struct Allocator ;
53
53
54
54
unsafe impl GlobalAlloc for Allocator {
55
+ /// Allocate memory using [`BootServices::allocate_pool`]. The allocation is
56
+ /// of type [`MemoryType::LOADER_DATA`].
55
57
unsafe fn alloc ( & self , layout : Layout ) -> * mut u8 {
56
58
let mem_ty = MemoryType :: LOADER_DATA ;
57
59
let size = layout. size ( ) ;
58
60
let align = layout. align ( ) ;
59
61
60
62
if align > 8 {
61
- // allocate more space for alignment
62
- let ptr = if let Ok ( ptr) = boot_services ( ) . as_ref ( ) . allocate_pool ( mem_ty, size + align)
63
- {
64
- ptr
65
- } else {
66
- return ptr:: null_mut ( ) ;
67
- } ;
68
- // calculate align offset
69
- let mut offset = ptr. align_offset ( align) ;
63
+ // The requested alignment is greater than 8, but `allocate_pool` is
64
+ // only guaranteed to provide eight-byte alignment. Allocate extra
65
+ // space so that we can return an appropriately-aligned pointer
66
+ // within the allocation.
67
+ let full_alloc_ptr =
68
+ if let Ok ( ptr) = boot_services ( ) . as_ref ( ) . allocate_pool ( mem_ty, size + align) {
69
+ ptr
70
+ } else {
71
+ return ptr:: null_mut ( ) ;
72
+ } ;
73
+
74
+ // Calculate the offset needed to get an aligned pointer within the
75
+ // full allocation. If that offset is zero, increase it to `align`
76
+ // so that we still have space to store the extra pointer described
77
+ // below.
78
+ let mut offset = full_alloc_ptr. align_offset ( align) ;
70
79
if offset == 0 {
71
80
offset = align;
72
81
}
73
- let return_ptr = ptr. add ( offset) ;
74
- // store allocated pointer before the struct
75
- ( return_ptr. cast :: < * mut u8 > ( ) ) . sub ( 1 ) . write ( ptr) ;
76
- return_ptr
82
+
83
+ // Before returning the aligned allocation, store a pointer to the
84
+ // full unaligned allocation in the bytes just before the aligned
85
+ // allocation. We know we have at least eight bytes there due to
86
+ // adding `align` to the memory allocation size. We also know the
87
+ // write is appropriately aligned for a `*mut u8` pointer because
88
+ // `align_ptr` is aligned, and alignments are always powers of two
89
+ // (as enforced by the `Layout` type).
90
+ let aligned_ptr = full_alloc_ptr. add ( offset) ;
91
+ ( aligned_ptr. cast :: < * mut u8 > ( ) ) . sub ( 1 ) . write ( full_alloc_ptr) ;
92
+ aligned_ptr
77
93
} else {
94
+ // The requested alignment is less than or equal to eight, and
95
+ // `allocate_pool` always provides eight-byte alignment, so we can
96
+ // use `allocate_pool` directly.
78
97
boot_services ( )
79
98
. as_ref ( )
80
99
. allocate_pool ( mem_ty, size)
81
100
. unwrap_or ( ptr:: null_mut ( ) )
82
101
}
83
102
}
84
103
104
+ /// Deallocate memory using [`BootServices::free_pool`].
85
105
unsafe fn dealloc ( & self , mut ptr : * mut u8 , layout : Layout ) {
86
106
if layout. align ( ) > 8 {
107
+ // Retrieve the pointer to the full allocation that was packed right
108
+ // before the aligned allocation in `alloc`.
87
109
ptr = ( ptr as * const * mut u8 ) . sub ( 1 ) . read ( ) ;
88
110
}
89
111
boot_services ( ) . as_ref ( ) . free_pool ( ptr) . unwrap ( ) ;
0 commit comments