Skip to content

Commit e0c9d78

Browse files
committed
uefi: allocator: use match{} arms to separate cases
This prepares changes in upcoming commits.
1 parent 4248ee8 commit e0c9d78

File tree

1 file changed

+57
-47
lines changed

1 file changed

+57
-47
lines changed

uefi/src/allocator.rs

Lines changed: 57 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -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

Comments
 (0)