diff --git a/kernel/src/lib.rs b/kernel/src/lib.rs index fb109f7..7171d76 100644 --- a/kernel/src/lib.rs +++ b/kernel/src/lib.rs @@ -8,6 +8,7 @@ ptr_cast_slice, likely_unlikely, int_roundings, + slice_ptr_get, never_type )] #![cfg_attr(test, feature(custom_test_frameworks))] diff --git a/kernel/src/main.rs b/kernel/src/main.rs index 12b956f..563b64c 100644 --- a/kernel/src/main.rs +++ b/kernel/src/main.rs @@ -27,15 +27,15 @@ extern "C" fn _start() -> ! { kernel::limine::init_limine_boot_info(); - kernel::serial_println!( - "HHDM offset: 0x{:#x}", - kernel::memory::HHDM_BASE.get().unwrap() - ); + // kernel::serial_println!( + // "HHDM offset: 0x{:#x}", + // kernel::memory::HHDM_BASE.get().unwrap() + // ); - kernel::serial_println!( - "Memory map: {:#?}", - kernel::limine::MEMMAP_REQUEST.entries() - ); + // kernel::serial_println!( + // "Memory map: {:#?}", + // kernel::limine::MEMMAP_REQUEST.entries() + // ); GDT.load(); IDT.load(); diff --git a/kernel/src/memory.rs b/kernel/src/memory.rs index 50200e0..429b06e 100644 --- a/kernel/src/memory.rs +++ b/kernel/src/memory.rs @@ -4,10 +4,11 @@ use core::{ fmt::Debug, num::{NonZero, NonZeroUsize}, ptr::NonNull, + range::Range, }; use bit_field::BitField; -use rbtree::RBTree; +use rbtree::{RBTree, UnsafeNode}; use crate::{serial_println, sync::OnceLock, x86_64::PAGE_SIZE}; @@ -22,7 +23,7 @@ pub trait VirtAddrTranslationExt: Sized { } } -#[derive(Clone, Copy, PartialEq, Eq)] +#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] pub struct PhyAddr(pub u64); impl PhyAddr { @@ -119,11 +120,6 @@ impl VirtAddr { } } -pub struct PageChunk { - pub next: PhyAddr, - pub count: usize, -} - pub struct PhysicalMemoryManager { /// On amd64 platforms, there are at most 2^40 pages of addressable physical /// memory. We maintain a binary tree of free page chunks for each @@ -165,20 +161,21 @@ impl Debug for PhysicalMemoryManager { impl PhysicalMemoryManager { pub fn from_memory_map(memory_map: &[crate::boot::MemoryRegion]) -> Self { - let usable_regions = memory_map - .iter() - .filter(|region| region.region_type.is_usable()); + let mut pmm = PhysicalMemoryManager { + buddies: [(); 40].map(|_| RBTree::default()), + }; - let buddies = [(); 40].map(|_| RBTree::default()); + let mut free_usable_regions = free_usable_regions_from_memory_map(memory_map); - let mut pmm = PhysicalMemoryManager { buddies }; + serial_println!("Freeing usable regions: {:#?}", free_usable_regions); - usable_regions.for_each(|region| { - let idx = region.start / crate::x86_64::PAGE_SIZE as u64; - let count = region.length / crate::x86_64::PAGE_SIZE as u64; - - pmm.free_region(idx as usize, count as usize); - }); + for region in free_usable_regions + .drain() + .map(|p| unsafe { p.cast::().as_ref() }) + { + let (idx, count) = region.index_and_count(); + pmm.free_region(idx, count); + } pmm } @@ -226,6 +223,8 @@ impl PhysicalMemoryManager { unsafe { chunk.as_ref().phy().map(|phy| (phy, 1 << bin)) } } + // TODO: grow and shrink + fn free_region(&mut self, page_index: usize, mut count: usize) { serial_println!( "Freeing region: page_index = {:#x}, count = {:#x}", @@ -286,7 +285,7 @@ impl PhysicalMemoryManager { ); let None = self.buddies[bin].insert_node(node_ptr) else { panic!( - "Attempted to free a region that is already free: page_index = {}, count = {}", + "Attempted to free a region that is already free: page_index = {:#x}, count = {:#x}", page_index, bit ) }; @@ -294,10 +293,6 @@ impl PhysicalMemoryManager { } } -pub struct PhysicalMemoryAllocator { - tree: RBTree, -} - #[repr(C)] #[derive(Debug)] struct PageHeader { @@ -305,11 +300,79 @@ struct PageHeader { count: u64, } +unsafe impl UnsafeNode for PageHeader { + type Key = PhysicalPageNodeKey; + + fn left(&self) -> Option> { + self.node.left().map(Self::unsafe_from_node) + } + + fn right(&self) -> Option> { + self.node.right().map(Self::unsafe_from_node) + } + + fn parent(&self) -> Option> { + self.node.parent().map(Self::unsafe_from_node) + } + + fn key(&self) -> &Self::Key { + self.node.key() + } + + fn color(&self) -> rbtree::Color { + self.node.color() + } + + fn set_left(&self, left: Option>) { + self.node.set_left(left.map(Self::unsafe_into_node)) + } + + fn set_right(&self, right: Option>) { + self.node.set_right(right.map(Self::unsafe_into_node)) + } + + fn set_parent(&self, parent: Option>) { + self.node.set_parent(parent.map(Self::unsafe_into_node)) + } + + fn set_color(&self, color: rbtree::Color) { + self.node.set_color(color) + } +} + impl PageHeader { fn phy(&self) -> Option { PhyAddr::from_hhdm_virt(self) } + fn unsafe_from_node(node: NonNull) -> NonNull { + node.cast::() + } + + fn unsafe_into_node(ptr: NonNull) -> NonNull { + ptr.cast::() + } + + fn contains(&self, phy: PhyAddr) -> bool { + let start = self.phy().expect("PageHeader is not in HHDM"); + let end = start.page_add(self.count as usize); + + phy >= start && phy < end + } + + fn page_range(&self) -> Range { + let (start, count) = self.index_and_count(); + + (start..start + count).into() + } + + fn index_and_count(&self) -> (usize, usize) { + let start = self.phy().expect("PageHeader is not in HHDM").page_index() as usize; + let count = self.count as usize; + + (start, count) + } + fn new_from_page_idx_and_count(page_idx: usize, count: usize) -> NonNull { let phy = PhyAddr(page_idx as u64 * PAGE_SIZE as u64); let virt = phy.as_hhdm_virt(); @@ -326,72 +389,112 @@ impl PageHeader { } } -impl Debug for PhysicalMemoryAllocator { - fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { - f.debug_struct("PhysicalMemoryAllocator") - .field_with("tree", |f| { - let iter = self.tree.iter(); - write!(f, "[")?; - for key in iter { - let header = unsafe { (&raw const *key).cast::().read_volatile() }; +fn free_usable_regions_from_memory_map( + memory_map: &[crate::boot::MemoryRegion], +) -> RBTree { + let mut usable_regions = memory_map + .iter() + .copied() + .filter(|region| region.region_type.is_usable()) + .fold(RBTree::new(), |mut tree, region| { + let idx = region.start as usize / crate::x86_64::PAGE_SIZE; + let count = region.length as usize / crate::x86_64::PAGE_SIZE; - if let Some(phy) = header.phy() { - writeln!( - f, - "({:?}..{:?}), ", - phy, - phy.page_add(header.count as usize) - )?; - } - } - write!(f, "]") - }) - .finish() - } -} + let node_ptr = PageHeader::new_from_page_idx_and_count(idx, count); -impl PhysicalMemoryAllocator { - pub fn from_memory_map(memory_map: &[crate::boot::MemoryRegion]) -> Self { - let mapped_non_usable_regions = memory_map - .iter() - .filter(|region| region.region_type.mapped_non_usable()); + let None = tree.insert_node(node_ptr) else { + panic!( + "Failed to insert node for page_index = {}, count = {}", + idx, count + ); + }; - let usable_regions = memory_map - .iter() - .copied() - .filter(|region| region.region_type.is_usable()); - - let mut pmm = PhysicalMemoryAllocator { - tree: RBTree::new(), - }; - - usable_regions.for_each(|region| { - let idx = region.start / crate::x86_64::PAGE_SIZE as u64; - let count = region.length / crate::x86_64::PAGE_SIZE as u64; - - pmm.free_region(idx as usize, count as usize); + tree }); - pmm - } - - pub fn free_region(&mut self, page_index: usize, count: usize) { - serial_println!( - "Freeing region: page_index = {}, count = {}", - page_index, - count - ); - - let node_ptr = - PageHeader::new_from_page_idx_and_count(page_index, count).cast::(); - - let None = self.tree.insert_node(node_ptr) else { - panic!( - "Failed to insert node for page_index = {}, count = {}", - page_index, count - ); - }; + let mapped_non_usable_regions = memory_map + .iter() + .filter(|region| region.region_type.mapped_non_usable()); + + for region in mapped_non_usable_regions { + let idx = region.start as usize / crate::x86_64::PAGE_SIZE; + let count = region.length as usize / crate::x86_64::PAGE_SIZE; + let end = idx + count; + + while let Some(overlapping) = usable_regions + .range( + PhysicalPageNodeKey::from_page_index(idx) + ..PhysicalPageNodeKey::from_page_index(idx + count), + ) + .map(|p| unsafe { + p.phy() + .unwrap() + .as_hhdm_virt() + .as_ptr::() + .as_ref() + .unwrap() + }) + .find(|header| { + header.contains(PhyAddr::from_page_index(idx)) + || header.contains(PhyAddr::from_page_index(idx + count - 1)) + }) + .inspect(|header| { + usable_regions + .remove(header.node.key()) + .expect("node exists"); + }) + { + // We found an overlapping region. + + serial_println!("Found overlapping region: {:?}", overlapping); + + let Range { + start: overlap_idx, + end: overlap_end, + } = overlapping.page_range(); + let overlap_count = overlap_end - overlap_idx; + + // check if our region is a suffix or prefix of the overlapping region + if overlap_idx == idx { + if let Some(suffix_count) = overlap_count.checked_sub(count) { + let None = usable_regions.insert_node(PageHeader::new_from_page_idx_and_count( + overlap_idx + count, + suffix_count, + )) else { + panic!("pmm: double free") + }; + } + } else if overlap_end == end { + if let Some(suffix_count) = overlap_count.checked_sub(count) { + let None = usable_regions.insert_node(PageHeader::new_from_page_idx_and_count( + overlap_idx, + suffix_count, + )) else { + panic!("pmm: double free") + }; + } + } else { + // region is in the middle of the overlapping region + if let Some(prefix_count) = idx.checked_sub(overlap_idx) { + let None = usable_regions.insert_node(PageHeader::new_from_page_idx_and_count( + overlap_idx, + prefix_count, + )) else { + panic!("pmm: double free") + }; + } + + if let Some(suffix_count) = overlap_end.checked_sub(end) { + let None = usable_regions + .insert_node(PageHeader::new_from_page_idx_and_count(end, suffix_count)) + else { + panic!("pmm: double free") + }; + } + } + } } + usable_regions } #[derive(Debug, Default)] @@ -629,7 +732,315 @@ unsafe impl rbtree::UnsafeNode for PhysicalPageNode { } } +pub mod slab { + //! A slab allocator + + use core::{ + alloc::Layout, + cell::Cell, + hint::{cold_path, unlikely}, + num::NonZero, + ptr::NonNull, + }; + + const UNLINKED: NonNull<()> = unsafe { NonNull::new_unchecked(!0 as *mut ()) }; + + use alloc::alloc::Allocator; + + use crate::x86_64::PAGE_SIZE; + + pub struct Slab { + /// Size and alignment of each element in the slab. + element_size: usize, + /// Pointer to the first chunk in the slab. + head: Option>>, + alloc: A, + } + + struct SlabChunk { + /// Pointer to the next chunk in the slab. + next: Option>>, + /// Pointer to the slab that owns this chunk. + slab: NonNull>, + /// Linked list of free elements in this slab. When this is `None`, the + /// slab is full. + free: Cell>>, + /// Number of outstanding allocations from this slab. When this reaches + /// zero, the slab can be freed. + count: Cell, + } + + struct ChunkSlot(Option>); + + enum SlotResult { + Some(NonNull), + Last(NonNull), + None, + } + + impl SlabChunk { + fn pop_free_slot(&self) -> SlotResult { + let Some(slot) = self.free.get() else { + return SlotResult::None; + }; + + self.free.set(unsafe { slot.as_ref() }.next()); + self.count.update(|count| count + 1); + + match self.free.get() { + Some(_) => SlotResult::Some(slot.cast()), + None => SlotResult::Last(slot.cast()), + } + } + + /// returns `true` if the slab is now empty and can be freed + fn push_free_slot(&self, slot: NonNull) -> bool { + let slot = slot.cast::(); + let next = self.free.get(); + unsafe { slot.as_ptr().write(ChunkSlot(next)) }; + self.free.set(Some(slot)); + self.count.update(|count| count - 1); + + self.count.get() == 0 + } + } + + impl ChunkSlot { + fn next(&self) -> Option> { + self.0 + } + } + + impl Slab { + fn new(element_size: usize, alloc: A) -> Self { + assert!( + element_size.is_power_of_two(), + "element_size must be a power of two" + ); + + Self { + element_size, + head: None, + alloc, + } + } + + fn first_slot_offset(&self) -> usize { + foundation::mem::align_up(core::mem::size_of::>(), self.element_size) + } + + fn count_and_layout(&self) -> (usize, Layout) { + let (count, size, align) = { + let one_page_count = (PAGE_SIZE - self.first_slot_offset()) / self.element_size; + + if one_page_count < 3 { + let count = 3; + let size = + (self.first_slot_offset() + count * self.element_size).next_power_of_two(); + + assert!(size.is_multiple_of(PAGE_SIZE)); + assert!(size >= PAGE_SIZE); + assert!(size.is_multiple_of(self.element_size)); + + (count, size, size) + } else { + let count = one_page_count; + let size = PAGE_SIZE; + (count, size, self.element_size) + } + }; + + (count, unsafe { + Layout::from_size_align_unchecked(size, align) + }) + } + + fn alloc_chunk(&mut self) -> NonNull> { + // we want to limit chunks to 1 page unless the element size is so + // large that we can fit fewer than 3 elements in a page. + let (count, layout) = self.count_and_layout(); + + let Some(bytes) = self.alloc.allocate(layout).ok() else { + panic!() + }; + + let chunk = bytes.as_non_null_ptr().cast::>(); + unsafe { + let first_slot = chunk + .as_ptr() + .byte_add(self.first_slot_offset()) + .cast::(); + + for i in 0..(count - 1) { + let chunk = first_slot.byte_add(i * self.element_size); + let next = first_slot.byte_add((i + 1) * self.element_size); + chunk.write(ChunkSlot(Some(NonNull::new_unchecked(next)))); + } + first_slot + .byte_add((count - 1) * self.element_size) + .write(ChunkSlot(None)); + + chunk.write(SlabChunk { + next: self.head, + slab: NonNull::from(self), + free: Cell::new(Some(NonNull::new_unchecked(first_slot))), + count: Cell::new(0), + }); + } + + chunk + } + + #[cold] + fn alloc_chunk_cold(&mut self) -> NonNull> { + self.alloc_chunk() + } + + fn alloc_slot(&mut self) -> NonNull<[u8]> { + let mut chunk = match self.head { + Some(chunk) => chunk, + None => { + let chunk = self.alloc_chunk_cold(); + self.head = Some(chunk); + chunk + } + }; + + let chunk = unsafe { chunk.as_mut() }; + + let ptr = match chunk.pop_free_slot() { + SlotResult::Some(non_null) => non_null, + SlotResult::Last(non_null) => { + self.head = chunk.next.replace(UNLINKED.cast()); + non_null + } + SlotResult::None => { + panic!("SlabChunk is full, but it is still the head of the slab"); + } + }; + + ptr.cast_slice(self.element_size) + } + + fn free_slot(&mut self, slot: NonNull) { + let (_, layout) = self.count_and_layout(); + let mut chunk = slot + .map_addr(|addr| unsafe { + NonZero::new_unchecked(foundation::mem::align_down(addr.get(), layout.align())) + }) + .cast::>(); + + let chunk = unsafe { chunk.as_mut() }; + let linked = chunk.next != Some(UNLINKED.cast()); + + if chunk.push_free_slot(slot) { + if linked { + let mut head = self.head.expect("chunk is linked, so head exists"); + while let Some(next) = unsafe { head.as_ref().next } { + if next == chunk.into() { + unsafe { head.as_mut().next = chunk.next }; + break; + } + head = next; + } + } + + unsafe { + self.alloc + .deallocate(NonNull::from_mut(chunk).cast(), layout) + }; + } else if !linked { + chunk.next = self.head; + self.head = Some(chunk.into()); + } + } + } + + const SLAB_ALLOCATOR_BUCKETS: usize = 8; + pub struct SlabAllocator { + /// Slabs for each power-of-two from 16 bytes to 2048 bytes (inclusive). + slabs: [Slab; SLAB_ALLOCATOR_BUCKETS], + alloc: A, + } + + impl SlabAllocator { + pub fn new(alloc: A) -> Self { + let slabs = [ + Slab::new(16, alloc.clone()), + Slab::new(32, alloc.clone()), + Slab::new(64, alloc.clone()), + Slab::new(128, alloc.clone()), + Slab::new(256, alloc.clone()), + Slab::new(512, alloc.clone()), + Slab::new(1024, alloc.clone()), + Slab::new(2048, alloc.clone()), + ]; + + Self { slabs, alloc } + } + + fn slab_index_for_size(size: usize) -> Option { + if size == 0 { + return None; + } + + // the smallest slab is 16 bytes + let size = size.max(16); + + // get the index of the slab by calculating the log2 of the size and + // subtracting 4 (since 2^4 = 16) + let index = + (size.next_power_of_two().trailing_zeros() - 16usize.trailing_zeros()) as usize; + + // we have 8 slabs, so the index must be less than 8 + if index < SLAB_ALLOCATOR_BUCKETS { + Some(index) + } else { + None + } + } + + pub fn alloc(&mut self, layout: Layout) -> Option> { + if unlikely(layout.size() == 0) { + return Some(NonNull::dangling().cast_slice(0)); + } + + let size = layout.size().max(layout.align()); + + match Self::slab_index_for_size(size) { + Some(slab_index) => { + Some(unsafe { self.slabs.get_unchecked_mut(slab_index).alloc_slot() }) + } + None => { + // allocate directly from the backing allocator + self.alloc.allocate(layout).ok() + } + } + } + + pub fn dealloc(&mut self, ptr: NonNull, layout: Layout) { + if unlikely(layout.size() == 0) { + return; + } + + let size = layout.size().max(layout.align()); + + match Self::slab_index_for_size(size) { + Some(slab_index) => unsafe { + self.slabs.get_unchecked_mut(slab_index).free_slot(ptr) + }, + None => { + // deallocate directly to the backing allocator + unsafe { self.alloc.deallocate(ptr, layout) } + } + } + } + } +} + pub mod bump { + //! A bump allocator inspired by / taken from the `stumpalo` crate + use core::{ alloc::{Allocator, Layout}, cell::Cell,