1 use crate::{
2     arch::mm::LockedFrameAllocator,
3     libs::align::page_align_up,
4     mm::{MMArch, MemoryManagementArch, VirtAddr},
5 };
6 
7 use core::{
8     alloc::{AllocError, GlobalAlloc, Layout},
9     intrinsics::unlikely,
10     ptr::NonNull,
11 };
12 
13 use super::page_frame::{FrameAllocator, PageFrameCount};
14 
15 /// 类kmalloc的分配器应当实现的trait
16 pub trait LocalAlloc {
local_alloc(&self, layout: Layout) -> *mut u817     unsafe fn local_alloc(&self, layout: Layout) -> *mut u8;
local_alloc_zeroed(&self, layout: Layout) -> *mut u818     unsafe fn local_alloc_zeroed(&self, layout: Layout) -> *mut u8;
local_dealloc(&self, ptr: *mut u8, layout: Layout)19     unsafe fn local_dealloc(&self, ptr: *mut u8, layout: Layout);
20 }
21 
22 pub struct KernelAllocator;
23 
24 impl KernelAllocator {
alloc_in_buddy(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError>25     unsafe fn alloc_in_buddy(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
26         // 计算需要申请的页数,向上取整
27         let count = (page_align_up(layout.size()) / MMArch::PAGE_SIZE).next_power_of_two();
28         let page_frame_count = PageFrameCount::new(count);
29         let (phy_addr, allocated_frame_count) = LockedFrameAllocator
30             .allocate(page_frame_count)
31             .ok_or(AllocError)?;
32 
33         let virt_addr = unsafe { MMArch::phys_2_virt(phy_addr).ok_or(AllocError)? };
34         if unlikely(virt_addr.is_null()) {
35             return Err(AllocError);
36         }
37 
38         let slice = unsafe {
39             core::slice::from_raw_parts_mut(
40                 virt_addr.data() as *mut u8,
41                 allocated_frame_count.data() * MMArch::PAGE_SIZE,
42             )
43         };
44         return Ok(NonNull::from(slice));
45     }
46 
free_in_buddy(&self, ptr: *mut u8, layout: Layout)47     unsafe fn free_in_buddy(&self, ptr: *mut u8, layout: Layout) {
48         // 由于buddy分配的页数量是2的幂,因此释放的时候也需要按照2的幂向上取整。
49         let count = (page_align_up(layout.size()) / MMArch::PAGE_SIZE).next_power_of_two();
50         let page_frame_count = PageFrameCount::new(count);
51         let phy_addr = MMArch::virt_2_phys(VirtAddr::new(ptr as usize)).unwrap();
52         LockedFrameAllocator.free(phy_addr, page_frame_count);
53     }
54 }
55 
56 /// 为内核SLAB分配器实现LocalAlloc的trait
57 impl LocalAlloc for KernelAllocator {
local_alloc(&self, layout: Layout) -> *mut u858     unsafe fn local_alloc(&self, layout: Layout) -> *mut u8 {
59         return self
60             .alloc_in_buddy(layout)
61             .map(|x| x.as_mut_ptr() as *mut u8)
62             .unwrap_or(core::ptr::null_mut() as *mut u8);
63     }
64 
local_alloc_zeroed(&self, layout: Layout) -> *mut u865     unsafe fn local_alloc_zeroed(&self, layout: Layout) -> *mut u8 {
66         return self
67             .alloc_in_buddy(layout)
68             .map(|x| {
69                 let ptr: *mut u8 = x.as_mut_ptr();
70                 core::ptr::write_bytes(ptr, 0, x.len());
71                 ptr
72             })
73             .unwrap_or(core::ptr::null_mut() as *mut u8);
74     }
75 
local_dealloc(&self, ptr: *mut u8, layout: Layout)76     unsafe fn local_dealloc(&self, ptr: *mut u8, layout: Layout) {
77         self.free_in_buddy(ptr, layout);
78     }
79 }
80 
81 /// 为内核slab分配器实现GlobalAlloc特性
82 unsafe impl GlobalAlloc for KernelAllocator {
alloc(&self, layout: Layout) -> *mut u883     unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
84         return self.local_alloc(layout);
85         // self.local_alloc_zeroed(layout, 0)
86     }
87 
alloc_zeroed(&self, layout: Layout) -> *mut u888     unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
89         self.local_alloc_zeroed(layout)
90     }
91 
dealloc(&self, ptr: *mut u8, layout: Layout)92     unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
93         self.local_dealloc(ptr, layout);
94     }
95 }
96 
97 /// 内存分配错误处理函数
98 #[alloc_error_handler]
global_alloc_err_handler(layout: Layout) -> !99 pub fn global_alloc_err_handler(layout: Layout) -> ! {
100     panic!("global_alloc_error, layout: {:?}", layout);
101 }
102