Skip to main content

tg_kernel_alloc/
lib.rs

1//! 内存分配。
2
3#![no_std]
4#![deny(missing_docs)]
5
6extern crate alloc;
7
8use alloc::alloc::handle_alloc_error;
9use core::{
10    alloc::{GlobalAlloc, Layout},
11    cell::UnsafeCell,
12    ptr::NonNull,
13};
14use customizable_buddy::{BuddyAllocator, LinkedListBuddy, UsizeBuddy};
15
16/// 静态单元格,用于安全地包装需要内部可变性的静态变量。
17///
18/// 这是对 `static mut` 的替代方案,通过 `UnsafeCell` 提供内部可变性,
19/// 同时避免了 `static mut` 带来的编译警告和潜在的未定义行为。
20///
21/// # Safety
22///
23/// 此类型实现了 `Sync`,但调用者必须确保:
24/// - 在单处理器环境下使用,或
25/// - 通过外部同步机制保证线程安全
26struct StaticCell<T> {
27    inner: UnsafeCell<T>,
28}
29
30// SAFETY: StaticCell 仅在单处理器环境下使用,不存在并发访问。
31unsafe impl<T> Sync for StaticCell<T> {}
32
33impl<T> StaticCell<T> {
34    /// 创建一个包含给定值的新 `StaticCell`。
35    const fn new(value: T) -> Self {
36        Self {
37            inner: UnsafeCell::new(value),
38        }
39    }
40
41    /// 获取内部值的可变指针。
42    ///
43    /// # Safety
44    ///
45    /// 调用者必须确保不存在对内部值的并发访问。
46    #[inline]
47    fn get(&self) -> *mut T {
48        self.inner.get()
49    }
50}
51
52/// 初始化内存分配。
53///
54/// 参数 `base_address` 表示动态内存区域的起始位置。
55///
56/// # 注意
57///
58/// 此函数必须在使用任何堆分配之前调用,且只能调用一次。
59#[inline]
60pub fn init(base_address: usize) {
61    // SAFETY: 此函数只在内核初始化时调用一次,此时没有其他代码会访问 HEAP。
62    // base_address 由调用者保证是有效的堆起始地址。
63    heap_mut().init(
64        core::mem::size_of::<usize>().trailing_zeros() as _,
65        NonNull::new(base_address as *mut u8).unwrap(),
66    );
67}
68
69/// 将一个内存块托管到内存分配器。
70///
71/// # Safety
72///
73/// 调用者必须确保:
74/// - `region` 内存块与已经转移到分配器的内存块都不重叠
75/// - `region` 未被其他对象引用
76/// - `region` 必须位于初始化时传入的起始位置之后
77/// - 内存块的所有权将转移到分配器
78#[inline]
79pub unsafe fn transfer(region: &'static mut [u8]) {
80    let ptr = NonNull::new(region.as_mut_ptr()).unwrap();
81    // SAFETY: 由调用者保证内存块有效且不重叠
82    heap_mut().transfer(ptr, region.len());
83}
84
85/// 堆分配器。
86///
87/// 最大容量:6 + 21 + 3 = 30 -> 1 GiB。
88/// 不考虑并发使用,因此没有加锁。
89///
90/// 使用 `StaticCell` 包装以避免 `static mut` 的使用,
91/// 通过 `heap_mut()` 函数获取可变引用。
92static HEAP: StaticCell<BuddyAllocator<21, UsizeBuddy, LinkedListBuddy>> =
93    StaticCell::new(BuddyAllocator::new());
94
95/// 获取堆分配器的可变引用。
96///
97/// # Safety
98///
99/// 此函数内部使用 unsafe 获取可变引用,仅在单处理器环境下安全。
100/// 调用者必须确保不存在对 HEAP 的并发访问。
101#[inline]
102fn heap_mut() -> &'static mut BuddyAllocator<21, UsizeBuddy, LinkedListBuddy> {
103    // SAFETY: 仅在单处理器环境下使用,不存在并发访问。
104    unsafe { &mut *HEAP.get() }
105}
106
107struct Global;
108
109#[global_allocator]
110static GLOBAL: Global = Global;
111
112// SAFETY: GlobalAlloc 的实现必须是 unsafe 的。
113// 此实现仅用于单处理器环境,不支持并发访问。
114unsafe impl GlobalAlloc for Global {
115    #[inline]
116    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
117        // SAFETY: 通过 heap_mut() 访问分配器,在单处理器环境下不会有并发的分配请求。
118        // layout 的有效性由调用者(Rust 的 alloc 机制)保证。
119        if let Ok((ptr, _)) = heap_mut().allocate_layout::<u8>(layout) {
120            ptr.as_ptr()
121        } else {
122            handle_alloc_error(layout)
123        }
124    }
125
126    #[inline]
127    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
128        // SAFETY: 通过 heap_mut() 访问分配器,在单处理器环境下不会有并发的释放请求。
129        // ptr 和 layout 的有效性由调用者保证(必须是之前 alloc 返回的)。
130        heap_mut().deallocate_layout(NonNull::new(ptr).unwrap(), layout)
131    }
132}