Skip to content

Commit abcabc8

Browse files
committed
uefi: allocator: improve documentation and switch to match {}
The refactoring to a match {} block simplifies the following improvements.
1 parent a0a76eb commit abcabc8

File tree

1 file changed

+47
-39
lines changed

1 file changed

+47
-39
lines changed

uefi/src/allocator.rs

Lines changed: 47 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,10 @@
11
// SPDX-License-Identifier: MIT OR Apache-2.0
22

3-
//! This module implements Rust's global allocator interface using UEFI's memory allocation functions.
3+
//! This module [`Allocator`], which uses UEFI boot services to allocate memory.
44
//!
5-
//! If the `global_allocator` feature is enabled, the [`Allocator`] will be used
5+
//! By implementing the [`GlobalAlloc`] trait, this type can be designated as the
6+
//! global allocator using the `#[global_allocator]` attribute. If the
7+
//! `global_allocator` crate feature is enabled, the [`Allocator`] will be used
68
//! as the global Rust allocator.
79
//!
810
//! This allocator can only be used while boot services are active. If boot
@@ -44,7 +46,8 @@ fn get_memory_type() -> MemoryType {
4446

4547
/// Allocator which uses the UEFI pool allocation functions.
4648
///
47-
/// Only valid for as long as the UEFI boot services are available.
49+
/// The allocator can only be used as long as the UEFI boot services are
50+
/// available and have not been exited.
4851
#[derive(Debug)]
4952
pub struct Allocator;
5053

@@ -63,45 +66,50 @@ unsafe impl GlobalAlloc for Allocator {
6366
let align = layout.align();
6467
let memory_type = get_memory_type();
6568

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;
69+
match align {
70+
0..=8 /* UEFI default alignment */ => {
71+
// The requested alignment is less than or equal to eight, and
72+
// `allocate_pool` always provides eight-byte alignment, so we can
73+
// use `allocate_pool` directly.
74+
boot::allocate_pool(memory_type, size)
75+
.map(|ptr| ptr.as_ptr())
76+
.unwrap_or(ptr::null_mut())
8477
}
78+
9.. => {
79+
// The requested alignment is greater than 8, but `allocate_pool` is
80+
// only guaranteed to provide eight-byte alignment. Allocate extra
81+
// space so that we can return an appropriately-aligned pointer
82+
// within the allocation.
83+
let full_alloc_ptr = boot::allocate_pool(memory_type, size + align);
84+
let full_alloc_ptr = if let Ok(ptr) = full_alloc_ptr
85+
{
86+
ptr.as_ptr()
87+
} else {
88+
return ptr::null_mut();
89+
};
90+
91+
// Calculate the offset needed to get an aligned pointer within the
92+
// full allocation. If that offset is zero, increase it to `align`
93+
// so that we still have space to store the extra pointer described
94+
// below.
95+
let mut offset = full_alloc_ptr.align_offset(align);
96+
if offset == 0 {
97+
offset = align;
98+
}
8599

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
100+
// Before returning the aligned allocation, store a pointer to the
101+
// full unaligned allocation in the bytes just before the aligned
102+
// allocation. We know we have at least eight bytes there due to
103+
// adding `align` to the memory allocation size. We also know the
104+
// write is appropriately aligned for a `*mut u8` pointer because
105+
// `align_ptr` is aligned, and alignments are always powers of two
106+
// (as enforced by the `Layout` type).
107+
unsafe {
108+
let aligned_ptr = full_alloc_ptr.add(offset);
109+
(aligned_ptr.cast::<*mut u8>()).sub(1).write(full_alloc_ptr);
110+
aligned_ptr
111+
}
97112
}
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())
105113
}
106114
}
107115

0 commit comments

Comments
 (0)