diff --git a/Cargo.lock b/Cargo.lock index 34fdcba..b3c2b21 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,6 +2,12 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "autocfg" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" + [[package]] name = "bit_field" version = "0.10.3" @@ -14,13 +20,22 @@ version = "2.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b588b76d00fde79687d7646a9b5bdf3cc0f655e0bbd080335a95d7e96f3587da" +[[package]] +name = "foundation" +version = "0.1.0" +dependencies = [ + "num-traits", +] + [[package]] name = "kernel" version = "0.1.0" dependencies = [ "bit_field", "bitflags", + "foundation", "log", + "num-traits", "plain", "rbtree", "rustc-demangle", @@ -33,6 +48,15 @@ version = "0.4.33" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", +] + [[package]] name = "plain" version = "0.2.3" diff --git a/crates/foundation/Cargo.toml b/crates/foundation/Cargo.toml new file mode 100644 index 0000000..9ffa7a3 --- /dev/null +++ b/crates/foundation/Cargo.toml @@ -0,0 +1,7 @@ +[package] +name = "foundation" +version = "0.1.0" +edition = "2024" + +[dependencies] +num-traits = { version = "0.2.19", default-features = false } diff --git a/crates/foundation/src/alloc/bump.rs b/crates/foundation/src/alloc/bump.rs new file mode 100644 index 0000000..e4b3c64 --- /dev/null +++ b/crates/foundation/src/alloc/bump.rs @@ -0,0 +1,426 @@ + +//! A bump allocator inspired by / taken from the `stumpalo` crate + +use core::{ + alloc::{Allocator, Layout}, + cell::Cell, + marker::PhantomData, + ops::{Deref, DerefMut}, + ptr::{self, NonNull}, +}; + +#[repr(align(16))] +struct Chunk { + next: Option>, + size: usize, +} + +const EMPTY_CHUNK: Chunk = Chunk { + next: None, + size: 0, +}; + +struct RestorePoint { + top: *mut u8, + chunk: NonNull, +} + +pub struct Bump { + top: Cell<*mut u8>, + bottom: Cell<*mut Chunk>, + next_chunk: Cell>>, + backing_alloc: A, +} + +unsafe impl Send for Bump {} + +unsafe impl Allocator for BumpScope<'_, '_, A> { + fn allocate( + &self, + layout: Layout, + ) -> Result, core::alloc::AllocError> { + let virt = unsafe { self.alloc_layout(layout).ok_or(core::alloc::AllocError) }?; + + Ok(virt.cast_slice(layout.size())) + } + + unsafe fn deallocate(&self, _ptr: core::ptr::NonNull, _layout: Layout) {} +} + +impl Bump { + pub fn new_in(backing_alloc: A) -> Self { + let chunk = NonNull::from(&EMPTY_CHUNK); + let bottom = unsafe { chunk.as_ptr().add(1) }; + + Self { + top: Cell::new(bottom.cast::()), + bottom: Cell::new(bottom), + next_chunk: Cell::new(None), + backing_alloc, + } + } + + /// # Safety + /// The caller must ensure that the chunk is aligned to `Chunk` alignment. + pub unsafe fn from_raw_chunk_in(chunk: NonNull<[u8]>, backing_alloc: A) -> Self { + assert!( + chunk.as_ptr().addr().is_multiple_of(align_of::()), + "Chunk must be aligned to Chunk alignment" + ); + assert!( + chunk.len() >= size_of::(), + "Chunk must be at least the size of Chunk" + ); + + let len = chunk.len(); + let chunk = chunk.cast::(); + let bottom = unsafe { chunk.as_ptr().add(1) }; + let top = unsafe { bottom.byte_add(len - size_of::()) }; + + Self { + top: Cell::new(top.cast::()), + bottom: Cell::new(bottom), + next_chunk: Cell::new(None), + backing_alloc, + } + } + + pub fn as_scope<'env>(&'env mut self) -> BumpScope<'env, 'env, A> { + BumpScope { + bump: self, + _env: PhantomData, + _scope: PhantomData, + } + } + + pub fn scope<'env, F, R>(&'env mut self, f: F) -> R + where + F: for<'scope> FnOnce(&'scope mut BumpScope<'env, 'scope, A>) -> R + 'env, + R: 'env, + { + let restore = self.restore_point(); + + let mut scope = BumpScope { + bump: self, + _env: PhantomData, + _scope: PhantomData, + }; + + let result = f(&mut scope); + + unsafe { self.restore(restore) }; + + result + } + + fn restore_point(&self) -> RestorePoint { + RestorePoint { + top: self.top.get(), + chunk: NonNull::from(self.chunk()), + } + } + + unsafe fn restore(&self, restore_point: RestorePoint) { + // get current chunk and next chunk + let mut chunk = Some(NonNull::from(self.chunk())); + let mut head = self.next_chunk.get(); + + // walk the linked list of chunks used since the restore point and + // re-link them onto the free list. + while let Some(cnk) = chunk + && cnk != restore_point.chunk + { + let next = unsafe { ptr::replace(&raw mut (*cnk.as_ptr()).next, head) }; + + head = Some(cnk); + chunk = next; + } + + self.top.set(restore_point.top); + let bot = unsafe { restore_point.chunk.as_ptr().add(1) }; + self.bottom.set(bot); + self.next_chunk.set(head); + } + + #[inline] + pub unsafe fn alloc_layout(&self, layout: Layout) -> Option> { + let top = self.top.get(); + let bottom = self.bottom.get().addr(); + let extra = Self::extra_bytes(top, layout.align()); + + let slow_path = Self::out_of_mem(layout, top, bottom, false); + + if slow_path { + self.alloc_layout_slow_cold(layout) + } else { + let new_top = unsafe { top.byte_sub(extra).byte_sub(layout.size()) }; + self.top.set(new_top); + + Some(unsafe { NonNull::new_unchecked(new_top) }) + } + } + + #[inline(never)] + fn try_alloc_slow_with_no_inline(&self, f: F) -> Option> + where + F: FnOnce() -> T, + { + let p = self.alloc_layout_slow(Layout::new::())?; + let p = p.cast::(); + unsafe { p.write(f()) }; + Some(p) + } + + #[cold] + fn alloc_layout_slow_cold(&self, layout: Layout) -> Option> { + self.alloc_layout_slow(layout) + } + + // #[inline(never)] + // fn alloc_layout_slow_no_inline(&self, layout: Layout) -> Option> { + // self.alloc_layout_slow(layout) + // } + + fn alloc_layout_slow(&self, layout: Layout) -> Option> { + let extra_with_chunk = { + assert!(size_of::() == align_of::()); + // if the allocation is greater than the size of the chunk header, we need more bytes past the chunk header to align the allocation. + if layout.align() > size_of::() { + layout.align() - size_of::() + } else { + 0 + } + }; + + let min_size = layout.size() + extra_with_chunk; + + while let Some(mut chunk) = self.next_chunk.get() { + // SAFETY: `Bump` is not `Sync` and `alloc_layout_slow is not + // reentrant, we have exclusive access to `chunk` + let chunk = unsafe { chunk.as_mut() }; + let cap = chunk.size; + + if cap >= min_size { + let prev = self.chunk(); + self.next_chunk.set(chunk.next); + + // add the current chunk to the linked list of chunks, so it can be freed later. + chunk.next = Some(NonNull::from(prev)); + + unsafe { + // bottom is just past the chunk header + let bot = (&raw const *chunk).add(1).cast_mut(); + self.bottom.set(bot); + + // calculate top of chunk + let top = bot.byte_add(cap).cast::(); + + // sub allocation + let extra = Self::extra_bytes(top, layout.align()); + let top = top.byte_sub(extra).byte_sub(layout.size()); + self.top.set(top); + + return Some(NonNull::new_unchecked(top)); + }; + } else { + // chunk is too small, try the next one + self.next_chunk.set(chunk.next); + unsafe { + let layout = Layout::from_size_align_unchecked( + cap + size_of::(), + align_of::(), + ); + self.backing_alloc + .deallocate(NonNull::from(chunk).cast(), layout); + } + } + } + + let chunk = self.chunk(); + let prev = if chunk.size == 0 { + None + } else { + Some(NonNull::from(chunk)) + }; + + let size = self.next_chunk_size().max(min_size); + let new_chunk = self + .backing_alloc + .allocate(unsafe { + Layout::from_size_align_unchecked(size + size_of::(), align_of::()) + }) + .ok()? + .cast::(); + + unsafe { + new_chunk.write(Chunk { next: prev, size }); + + let bot = new_chunk.add(1).as_ptr(); + self.bottom.set(bot); + + let top = bot.byte_add(size).cast::(); + let extra = Self::extra_bytes(top, layout.align()); + let top = top.byte_sub(extra).byte_sub(layout.size()); + self.top.set(top); + + Some(NonNull::new_unchecked(top)) + } + } + + fn next_chunk_size(&self) -> usize { + let chunk = self.chunk(); + let cap = (chunk.size + size_of::()).saturating_mul(2) - size_of::(); + + cap.max(0x1000) + } + + pub unsafe fn alloc_raw(&self) -> Option> { + let layout = core::alloc::Layout::new::(); + unsafe { self.alloc_layout(layout).map(|p| p.cast::()) } + } + + fn out_of_mem(layout: Layout, top: *mut u8, bottom: usize, comptime: bool) -> bool { + let extra = Self::extra_bytes(top, layout.align()); + + let max_padding = layout.align() - 1; + let max_size = layout.size() + max_padding; + let top = top.addr(); + + // biggest possible virtual address on x86_64 is 2^57 - 1 + const MAX_ADDR: usize = ((1u64 << 57) - 1) as usize; + // offsets of < SAFE_SIZE are guaranteed not to overflow the address space. + const SAFE_SIZE: usize = usize::MAX - MAX_ADDR; + + if comptime && max_size < SAFE_SIZE { + if max_size <= 16 { + let top = top - extra; + let top = top - layout.size(); + bottom > top + } else if max_padding < 16 { + let top = top - extra; + bottom + layout.size() > top + } else { + bottom + layout.size() + max_padding > top + } + } else if max_padding + .checked_add(isize::MAX as usize) + .is_some_and(|p| p < SAFE_SIZE) + { + if max_size < 16 { + let top = top - extra; + bottom + layout.size() > top + } else { + bottom + extra + layout.size() > top + } + } else { + if max_padding < 16 { + let top = top - extra; + top.checked_sub(layout.size()).is_none_or(|t| t < bottom) + } else { + top.checked_sub(extra) + .and_then(|t| t.checked_sub(layout.size())) + .is_none_or(|t| t < bottom) + } + } + } + + fn extra_bytes(ptr: *mut u8, align: usize) -> usize { + assert!(align.is_power_of_two(), "Alignment must be a power of two"); + ptr.addr() & (align - 1) + } + + fn chunk(&self) -> &Chunk { + unsafe { self.bottom.get().sub(1).as_ref_unchecked() } + } +} + +#[repr(transparent)] +pub struct BumpScope<'env, 'scope, A: Allocator> { + bump: &'scope mut Bump, + _env: PhantomData<&'env &'env mut ()>, + _scope: PhantomData<&'scope &'scope mut ()>, +} + +impl<'env, 'scope, A: Allocator> Deref for BumpScope<'env, 'scope, A> { + type Target = Bump; + + fn deref(&self) -> &Self::Target { + self.bump + } +} + +impl<'env, 'scope, A: Allocator> DerefMut for BumpScope<'env, 'scope, A> { + fn deref_mut(&mut self) -> &mut Self::Target { + self.bump + } +} + +impl<'env, 'scope, A: Allocator> BumpScope<'env, 'scope, A> { + pub fn alloc_with(&self, f: F) -> &'scope mut T + where + F: FnOnce() -> T, + { + self.try_alloc_with(f).unwrap_or_else(|| { + panic!( + "Bump allocator out of memory when allocating {} bytes with alignment {}", + size_of::(), + align_of::() + ) + }) + } + + pub fn try_alloc_with(&self, f: F) -> Option<&'scope mut T> + where + F: FnOnce() -> T, + { + let layout = Layout::new::(); + let top = self.top.get(); + let bottom = self.bottom.get().addr(); + let extra = Bump::::extra_bytes(top, layout.align()); + + let slow_path = Bump::::out_of_mem(layout, top, bottom, false); + + let mut ptr = if layout.size() <= 16 { + if slow_path { + self.try_alloc_slow_with_no_inline(f)? + } else { + let new_top = unsafe { top.byte_sub(extra).byte_sub(layout.size()) }; + self.top.set(new_top); + unsafe { + ptr::write(new_top.cast::(), f()); + NonNull::new_unchecked(new_top.cast::()) + } + } + } else { + // in-place + let ptr = if slow_path { + self.alloc_layout_slow_cold(layout)? + } else { + let new_top = unsafe { top.byte_sub(extra).byte_sub(layout.size()) }; + self.top.set(new_top); + + unsafe { NonNull::new_unchecked(new_top) } + }; + + let ptr = ptr.cast::(); + unsafe { + ptr.write(f()); + } + + ptr + }; + + Some(unsafe { ptr.as_mut() }) + } +} + +fn _asdf(bump: &mut Bump) { + let mut bump = bump.as_scope(); + let x = bump.alloc_with(|| 3u64); + bump.scope(|bump| { + let _y = bump.alloc_with(|| 4u64); + let _z = bump.alloc_with(|| 5u64); + }); + + assert!(*x == 3); +} diff --git a/crates/foundation/src/alloc/mod.rs b/crates/foundation/src/alloc/mod.rs new file mode 100644 index 0000000..a6258f1 --- /dev/null +++ b/crates/foundation/src/alloc/mod.rs @@ -0,0 +1,23 @@ +use core::ptr::NonNull; + +use liballoc::alloc::Allocator; + +pub mod bump; +pub mod slab; + +#[derive(Debug, Default)] +pub struct PanicingAllocator; + +unsafe impl Allocator for PanicingAllocator { + fn allocate( + &self, + layout: core::alloc::Layout, + ) -> Result, core::alloc::AllocError> { + match layout.size() { + 0 => Ok(NonNull::slice_from_raw_parts(NonNull::dangling(), 0)), + _ => panic!("PanicingAllocator cannot allocate memory"), + } + } + + unsafe fn deallocate(&self, _ptr: NonNull, _layout: core::alloc::Layout) {} +} diff --git a/crates/foundation/src/alloc/slab.rs b/crates/foundation/src/alloc/slab.rs new file mode 100644 index 0000000..9fb65f2 --- /dev/null +++ b/crates/foundation/src/alloc/slab.rs @@ -0,0 +1,379 @@ +//! A slab allocator + +use core::{alloc::Layout, cell::Cell, hint::unlikely, mem::offset_of, num::NonZero, ptr::NonNull}; + +const UNLINKED: NonNull<()> = unsafe { NonNull::new_unchecked(!0 as *mut ()) }; + +use liballoc::alloc::Allocator; + +use crate::mem; + +const PAGE_SIZE: usize = 4096; + +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 { + 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(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()); + } + } +} + +impl Drop for Slab { + fn drop(&mut self) { + let (_, layout) = self.count_and_layout(); + + let mut chunk = self.head.take(); + while let Some(mut chunk_ptr) = chunk { + let chunk_ref = unsafe { chunk_ptr.as_mut() }; + chunk = chunk_ref.next.take(); + + unsafe { + self.alloc.deallocate(chunk_ptr.cast(), layout); + } + } + } +} + +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(_) => unsafe { + let chunk = ptr + .map_addr(|addr| { + NonZero::new_unchecked(mem::align_down(addr.get(), layout.align())) + }) + .cast::>(); + + let mut slab = (&raw const (*chunk.as_ptr()).slab).read(); + slab.as_mut().free_slot(ptr); + }, + None => { + // deallocate directly to the backing allocator + unsafe { self.alloc.deallocate(ptr, layout) } + } + } + } +} + +impl Drop for SlabAllocator { + fn drop(&mut self) {} +} + +#[cfg(test)] +mod tests { + use super::*; + + use std::{prelude::rust_2024::*, sync::Mutex}; + + struct SlabAllocatorWrapper(Mutex>); + + impl SlabAllocatorWrapper { + fn new(alloc: A) -> Self { + Self(Mutex::new(SlabAllocator::new(alloc))) + } + } + impl SlabAllocatorWrapper { + fn new_system() -> SlabAllocatorWrapper { + SlabAllocatorWrapper::new(std::alloc::System) + } + } + + unsafe impl Allocator for SlabAllocatorWrapper { + fn allocate( + &self, + layout: core::alloc::Layout, + ) -> Result, core::alloc::AllocError> { + self.0 + .lock() + .unwrap() + .alloc(layout) + .ok_or(core::alloc::AllocError) + } + + unsafe fn deallocate(&self, ptr: NonNull, layout: core::alloc::Layout) { + self.0.lock().unwrap().dealloc(ptr, layout) + } + } + + #[test] + fn test_slab_allocator() { + let alloc = SlabAllocatorWrapper::new_system(); + + let layout = Layout::from_size_align(16, 8).unwrap(); + let ptr = alloc.allocate(layout).unwrap(); + assert_eq!(ptr.len(), 16); + assert_eq!(ptr.addr().get() % 8, 0); + + let mut boxed = Box::new_in(42u32, &alloc); + assert_eq!(*boxed, 42); + assert_ne!(Box::as_non_null(&mut boxed), ptr.cast()); + + unsafe { alloc.deallocate(ptr.cast(), layout) }; + _ = boxed; + } +} diff --git a/crates/foundation/src/lib.rs b/crates/foundation/src/lib.rs new file mode 100644 index 0000000..5c005aa --- /dev/null +++ b/crates/foundation/src/lib.rs @@ -0,0 +1,99 @@ +#![no_std] +#![feature( + const_trait_impl, + const_default, + allocator_api, + ptr_cast_slice, + likely_unlikely, + slice_ptr_get +)] +#![cfg_attr(test, feature(box_as_non_null))] + +extern crate alloc as liballoc; + +#[cfg(test)] +extern crate std; + +pub mod alloc; +pub mod sync; +pub mod mem { + use num_traits::PrimInt; + + pub fn replace(v: &mut T, f: impl FnOnce(T) -> (T, R)) -> R { + struct Guard; + impl Drop for Guard { + fn drop(&mut self) { + panic!("replace() panicked"); + } + } + let guard = Guard; + let value = unsafe { core::ptr::read(v) }; + let (new_value, ret) = f(value); + unsafe { core::ptr::write(v, new_value) }; + core::mem::forget(guard); + ret + } + + pub fn align_up(value: T, alignment: T) -> T { + (value + alignment - T::one()) & !(alignment - T::one()) + } + + pub fn align_down(value: T, alignment: T) -> T { + value & !(alignment - T::one()) + } + + /// # Safety + pub unsafe fn volatile_copy(src: *const T, dst: *mut T, count: usize) { + unsafe { + for i in 0..count { + let src_ptr = src.add(i); + let dst_ptr = dst.add(i); + let val = src_ptr.read(); + dst_ptr.write_volatile(val); + } + } + } +} + +pub struct DropGuard(::core::mem::ManuallyDrop); +impl DropGuard { + pub fn new(f: F) -> Self { + DropGuard(::core::mem::ManuallyDrop::new(f)) + } + #[allow(dead_code)] + pub fn forget(self) { + let mut this = ::core::mem::ManuallyDrop::new(self); + unsafe { + ::core::mem::ManuallyDrop::drop(&mut this.0); + } + } +} + +impl Drop for DropGuard { + fn drop(&mut self) { + unsafe { ::core::ptr::read(&*self.0)() } + } +} + +#[derive(Default)] +pub struct PanicGuard; + +impl PanicGuard { + pub fn new() -> Self { + Self + } + + pub fn forget(self) { + core::mem::forget(self); + } +} + +impl Drop for PanicGuard { + #[track_caller] + fn drop(&mut self) { + panic!( + "PanicGuard dropped without being forgotten: {}", + core::panic::Location::caller() + ); + } +} diff --git a/crates/foundation/src/sync/mod.rs b/crates/foundation/src/sync/mod.rs new file mode 100644 index 0000000..d00f9ba --- /dev/null +++ b/crates/foundation/src/sync/mod.rs @@ -0,0 +1,487 @@ +mod spin_mutex { + use core::{ + cell::UnsafeCell, + sync::atomic::{AtomicBool, Ordering}, + }; + + pub struct SpinMutex { + lock: AtomicBool, + data: UnsafeCell, + } + + pub struct SpinMutexGuard<'a, T: ?Sized + 'a> { + mutex: &'a SpinMutex, + } + + impl core::ops::Deref for SpinMutexGuard<'_, T> { + type Target = T; + + fn deref(&self) -> &Self::Target { + unsafe { &*self.mutex.data.get() } + } + } + + impl core::ops::DerefMut for SpinMutexGuard<'_, T> { + fn deref_mut(&mut self) -> &mut Self::Target { + unsafe { &mut *self.mutex.data.get() } + } + } + + impl Drop for SpinMutexGuard<'_, T> { + fn drop(&mut self) { + unsafe { self.mutex.unlock() } + } + } + + unsafe impl Send for SpinMutex {} + // SAFETY: SpinMutex properly synchronises access to T, but T must be Send + // to be safely shared between threads. + unsafe impl Sync for SpinMutex {} + + // SAFETY: &SpinMutexGuard cannot yield a &mut T, and the SpinMutexGuard + // holds the lock for the duration of its lifetime, so it is safe to share + // between threads. + unsafe impl Sync for SpinMutexGuard<'_, T> {} + // SAFETY: SpinMutex may be unlocked on any thread, so it is safe to send + // the guard to another thread as long as T is Send. + unsafe impl Send for SpinMutexGuard<'_, T> {} + + impl SpinMutex { + pub const fn new(data: T) -> Self { + Self { + lock: AtomicBool::new(false), + data: UnsafeCell::new(data), + } + } + + pub fn into_inner(self) -> T { + self.data.into_inner() + } + } + + impl SpinMutex { + pub fn as_mut_ptr(&self) -> *mut T { + self.data.get() + } + + /// # Safety + /// The caller must be the logical owner of the locked mutex. + #[expect(clippy::mut_from_ref, reason = "unsafe mutex api")] + pub unsafe fn get_mut_unchecked(&self) -> &mut T { + unsafe { &mut *self.data.get() } + } + + /// # Safety + /// The caller must be the logical owner of the locked mutex. + pub unsafe fn get_unchecked(&self) -> &T { + unsafe { &*self.data.get() } + } + + pub fn try_lock(&self) -> bool { + self.lock + .compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed) + .is_ok() + } + pub fn try_lock_weak(&self) -> bool { + self.lock + .compare_exchange_weak(false, true, Ordering::Acquire, Ordering::Relaxed) + .is_ok() + } + pub fn lock(&self) -> SpinMutexGuard<'_, T> { + while !self.try_lock() { + core::hint::spin_loop(); + } + SpinMutexGuard { mutex: self } + } + + /// # Safety + /// The caller must be the logical owner of the locked mutex. + pub unsafe fn unlock(&self) { + self.lock.store(false, Ordering::Release); + } + } +} + +pub use spin_mutex::SpinMutex; + +mod once { + use core::{ + cell::{Cell, UnsafeCell}, + mem::{ManuallyDrop, MaybeUninit}, + sync::atomic::{AtomicU8, Ordering}, + }; + + const UNINITIALIZED: u8 = 0; + const INITIALIZING: u8 = 1; + const INITIALIZED: u8 = 2; + const POISONED: u8 = 3; + + pub struct OnceState { + poisoned: bool, + state_to_set: Cell, + } + + impl OnceState { + pub fn is_poisoned(&self) -> bool { + self.poisoned + } + pub fn poison(&self) { + self.state_to_set.set(POISONED); + } + } + + pub struct Once { + state: AtomicU8, + } + + const impl Default for Once { + fn default() -> Self { + Self::new() + } + } + + impl Once { + pub const fn new() -> Self { + Self { + state: AtomicU8::new(UNINITIALIZED), + } + } + + pub fn state(&self) -> u8 { + self.state.load(Ordering::Acquire) + } + + pub fn set_state(&self, state: u8) { + self.state.store(state, Ordering::Release); + } + + pub fn is_completed(&self) -> bool { + self.state.load(Ordering::Acquire) == INITIALIZED + } + + pub fn call_once(&self, f: F) + where + F: FnOnce(&OnceState), + { + if self.is_completed() { + return; + } + + self.call_once_slow(false, f); + } + + pub fn call_once_force(&self, f: F) + where + F: FnOnce(&OnceState), + { + if self.is_completed() { + return; + } + + self.call_once_slow(true, f); + } + + #[cold] + pub fn call_once_slow(&self, ignore_poison: bool, f: F) + where + F: FnOnce(&OnceState), + { + let mut state = self.state.load(Ordering::Acquire); + loop { + match state { + INITIALIZED => return, + POISONED if !ignore_poison => { + panic!("Once instance has previously been poisoned") + } + POISONED | UNINITIALIZED => { + match self.state.compare_exchange( + state, + INITIALIZING, + Ordering::Acquire, + // if we get `Err(INITIALIZED)`, we want to have + // acquired what the lock is protecting. + Ordering::Acquire, + ) { + Err(new) => { + state = new; + continue; + } + Ok(_) => { + // even though we don't have unwinding, for + // completeness sake we'll poison the lock if + // the closure panics. + let guard = crate::DropGuard::new(|| { + self.state.store(POISONED, Ordering::Release) + }); + + let state = OnceState { + poisoned: state == POISONED, + state_to_set: Cell::new(INITIALIZED), + }; + + f(&state); + guard.forget(); + + self.state + .store(state.state_to_set.take(), Ordering::Release); + return; + } + } + } + _ => { + assert_eq!(state, INITIALIZING); + loop { + state = self.state.load(Ordering::Acquire); + if state != INITIALIZING { + break; + } + core::hint::spin_loop(); + } + } + } + } + } + } + + pub struct OnceLock { + once: Once, + data: UnsafeCell>, + } + + unsafe impl Sync for OnceLock {} + + const impl Default for OnceLock { + fn default() -> Self { + Self { + once: Once::new(), + data: UnsafeCell::new(MaybeUninit::uninit()), + } + } + } + + impl OnceLock { + pub const fn new() -> Self { + Self { + once: Once::new(), + data: UnsafeCell::new(MaybeUninit::uninit()), + } + } + + pub const fn from(t: T) -> Self { + Self { + once: Once::new(), + data: UnsafeCell::new(MaybeUninit::new(t)), + } + } + + /// # Safety + /// The caller must ensure that the `OnceLock` has been initialized before calling this method. + pub unsafe fn get_unchecked(&self) -> &T { + unsafe { self.data.get().as_ref_unchecked().assume_init_ref() } + } + + /// # Safety + /// The caller must ensure that the `OnceLock` has been initialized before calling this method. + pub unsafe fn get_mut_unchecked(&mut self) -> &mut T { + unsafe { self.data.get().as_mut_unchecked().assume_init_mut() } + } + + pub fn get(&self) -> Option<&T> { + if self.is_completed() { + Some(unsafe { self.get_unchecked() }) + } else { + None + } + } + + pub fn get_mut(&mut self) -> Option<&mut T> { + if self.is_completed_mut() { + Some(unsafe { self.get_mut_unchecked() }) + } else { + None + } + } + + pub fn is_completed(&self) -> bool { + self.once.is_completed() + } + + pub fn is_completed_mut(&mut self) -> bool { + // avoid atomic load when we have mutable access to self + *self.once.state.get_mut() == INITIALIZED + } + + pub fn get_spinning(&self) -> Option<&T> { + loop { + match self.once.state() { + INITIALIZED => return Some(unsafe { self.get_unchecked() }), + INITIALIZING => core::hint::spin_loop(), + UNINITIALIZED | POISONED => return None, + _ => unreachable!(), + } + } + } + + pub fn initialize(&self, f: F) -> Result<(), E> + where + F: FnOnce() -> Result, + { + let mut res: Result<(), E> = Ok(()); + let cell = unsafe { &mut *self.data.get() }; + + self.once.call_once_force(|state| match f() { + Ok(value) => { + unsafe { cell.as_mut_ptr().write(value) }; + } + Err(e) => { + state.poison(); + res = Err(e); + } + }); + + res + } + + pub fn get_or_init(&self, f: F) -> &T + where + F: FnOnce() -> T, + { + match self.get_or_try_init(|| Ok::(f())) { + Ok(value) => value, + Err(_) => unreachable!(), + } + } + + pub fn get_mut_or_init(&mut self, f: F) -> &mut T + where + F: FnOnce() -> T, + { + match self.get_mut_or_try_init(|| Ok::(f())) { + Ok(value) => value, + Err(_) => unreachable!(), + } + } + + pub fn get_or_try_init(&self, f: F) -> Result<&T, E> + where + F: FnOnce() -> Result, + { + if !self.is_completed() { + self.initialize(f)?; + } + + unsafe { Ok(self.get_unchecked()) } + } + + pub fn get_mut_or_try_init(&mut self, f: F) -> Result<&mut T, E> + where + F: FnOnce() -> Result, + { + if self.get_mut().is_none() { + self.initialize(f)?; + } + + Ok(unsafe { self.get_mut_unchecked() }) + } + + pub fn try_insert(&self, value: T) -> Result<&T, (&T, T)> { + let mut val = Some(value); + let res = self.get_or_init(|| unsafe { val.take().unwrap_unchecked() }); + match val { + Some(value) => Err((res, value)), + None => Ok(res), + } + } + } + + impl Drop for OnceLock { + fn drop(&mut self) { + if self.is_completed() { + unsafe { self.data.get().as_mut_unchecked().assume_init_drop() } + } + } + } + + pub struct LazyLock T> { + once: Once, + f: UnsafeCell>, + t: UnsafeCell>, + } + + unsafe impl Sync for LazyLock {} + + impl LazyLock + where + F: FnOnce() -> T, + { + pub const fn new(f: F) -> Self { + Self { + once: Once::new(), + f: UnsafeCell::new(ManuallyDrop::new(f)), + t: UnsafeCell::new(MaybeUninit::uninit()), + } + } + + pub fn get(&self) -> Option<&T> { + if self.once.is_completed() { + Some(unsafe { (&*self.t.get()).assume_init_ref() }) + } else { + None + } + } + + fn force(this: &Self) -> &T { + this.once.call_once(|_| { + // SAFETY: because `call_once` will panic if poisoned, this + // closure will only be called once. + let f = unsafe { ManuallyDrop::take(&mut *this.f.get()) }; + let val = f(); + unsafe { (&mut *this.t.get()).write(val) }; + }); + + unsafe { (&*this.t.get()).assume_init_ref() } + } + + fn force_mut(this: &mut Self) -> &mut T { + this.once.call_once(|_| { + // SAFETY: because `call_once` will panic if poisoned, this + // closure will only be called once. + let f = unsafe { ManuallyDrop::take(&mut *this.f.get()) }; + let val = f(); + unsafe { (&mut *this.t.get()).write(val) }; + }); + + unsafe { (&mut *this.t.get()).assume_init_mut() } + } + } + + impl T> core::ops::Deref for LazyLock { + type Target = T; + + fn deref(&self) -> &Self::Target { + LazyLock::force(self) + } + } + + impl T> core::ops::DerefMut for LazyLock { + fn deref_mut(&mut self) -> &mut Self::Target { + LazyLock::force_mut(self) + } + } + + impl Drop for LazyLock { + fn drop(&mut self) { + match self.once.state() { + UNINITIALIZED => unsafe { + ManuallyDrop::drop(self.f.get_mut()); + }, + INITIALIZED => unsafe { + (&mut *self.t.get()).assume_init_drop(); + }, + _ => {} + } + } + } +} + +pub use once::{LazyLock, Once, OnceLock}; diff --git a/kernel/Cargo.toml b/kernel/Cargo.toml index 27aecf6..e883032 100644 --- a/kernel/Cargo.toml +++ b/kernel/Cargo.toml @@ -31,3 +31,5 @@ seq-macro = "0.3.6" rbtree = { path = "../crates/rbtree" } plain = "0.2.3" rustc-demangle = "0.1.28" +num-traits = { version = "0.2.19", default-features = false } +foundation = { version = "0.1.0", path = "../crates/foundation" }