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 { 17 unsafe fn local_alloc(&self, layout: Layout) -> *mut u8; 18 unsafe fn local_alloc_zeroed(&self, layout: Layout) -> *mut u8; 19 unsafe fn local_dealloc(&self, ptr: *mut u8, layout: Layout); 20 } 21 22 pub struct KernelAllocator; 23 24 impl KernelAllocator { 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 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 { 58 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 65 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 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 { 83 unsafe fn alloc(&self, layout: Layout) -> *mut u8 { 84 return self.local_alloc(layout); 85 // self.local_alloc_zeroed(layout, 0) 86 } 87 88 unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 { 89 self.local_alloc_zeroed(layout) 90 } 91 92 unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { 93 self.local_dealloc(ptr, layout); 94 } 95 } 96 97 /// 内存分配错误处理函数 98 #[cfg(target_os = "none")] 99 #[alloc_error_handler] 100 pub fn global_alloc_err_handler(layout: Layout) -> ! { 101 panic!("global_alloc_error, layout: {:?}", layout); 102 } 103