foundation crate

num_traits dep
This commit is contained in:
janis 2026-08-05 19:00:32 +02:00
parent 549b94d34c
commit 47288aeed1
Signed by: janis
SSH key fingerprint: SHA256:bB1qbbqmDXZNT0KKD5c2Dfjg53JGhj7B3CFcLIzSqq8
8 changed files with 1447 additions and 0 deletions

24
Cargo.lock generated
View file

@ -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"

View file

@ -0,0 +1,7 @@
[package]
name = "foundation"
version = "0.1.0"
edition = "2024"
[dependencies]
num-traits = { version = "0.2.19", default-features = false }

View file

@ -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<NonNull<Chunk>>,
size: usize,
}
const EMPTY_CHUNK: Chunk = Chunk {
next: None,
size: 0,
};
struct RestorePoint {
top: *mut u8,
chunk: NonNull<Chunk>,
}
pub struct Bump<A: Allocator> {
top: Cell<*mut u8>,
bottom: Cell<*mut Chunk>,
next_chunk: Cell<Option<NonNull<Chunk>>>,
backing_alloc: A,
}
unsafe impl<A: Allocator + Send> Send for Bump<A> {}
unsafe impl<A: Allocator> Allocator for BumpScope<'_, '_, A> {
fn allocate(
&self,
layout: Layout,
) -> Result<core::ptr::NonNull<[u8]>, 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<u8>, _layout: Layout) {}
}
impl<A: Allocator> Bump<A> {
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::<u8>()),
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>()),
"Chunk must be aligned to Chunk alignment"
);
assert!(
chunk.len() >= size_of::<Chunk>(),
"Chunk must be at least the size of Chunk"
);
let len = chunk.len();
let chunk = chunk.cast::<Chunk>();
let bottom = unsafe { chunk.as_ptr().add(1) };
let top = unsafe { bottom.byte_add(len - size_of::<Chunk>()) };
Self {
top: Cell::new(top.cast::<u8>()),
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<NonNull<u8>> {
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<F, T>(&self, f: F) -> Option<NonNull<T>>
where
F: FnOnce() -> T,
{
let p = self.alloc_layout_slow(Layout::new::<T>())?;
let p = p.cast::<T>();
unsafe { p.write(f()) };
Some(p)
}
#[cold]
fn alloc_layout_slow_cold(&self, layout: Layout) -> Option<NonNull<u8>> {
self.alloc_layout_slow(layout)
}
// #[inline(never)]
// fn alloc_layout_slow_no_inline(&self, layout: Layout) -> Option<NonNull<u8>> {
// self.alloc_layout_slow(layout)
// }
fn alloc_layout_slow(&self, layout: Layout) -> Option<NonNull<u8>> {
let extra_with_chunk = {
assert!(size_of::<Chunk>() == align_of::<Chunk>());
// 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::<Chunk>() {
layout.align() - size_of::<Chunk>()
} 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::<u8>();
// 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::<Chunk>(),
align_of::<Chunk>(),
);
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::<Chunk>(), align_of::<Chunk>())
})
.ok()?
.cast::<Chunk>();
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::<u8>();
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::<Chunk>()).saturating_mul(2) - size_of::<Chunk>();
cap.max(0x1000)
}
pub unsafe fn alloc_raw<T>(&self) -> Option<NonNull<T>> {
let layout = core::alloc::Layout::new::<T>();
unsafe { self.alloc_layout(layout).map(|p| p.cast::<T>()) }
}
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<A>,
_env: PhantomData<&'env &'env mut ()>,
_scope: PhantomData<&'scope &'scope mut ()>,
}
impl<'env, 'scope, A: Allocator> Deref for BumpScope<'env, 'scope, A> {
type Target = Bump<A>;
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<F, T>(&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::<T>(),
align_of::<T>()
)
})
}
pub fn try_alloc_with<F, T>(&self, f: F) -> Option<&'scope mut T>
where
F: FnOnce() -> T,
{
let layout = Layout::new::<T>();
let top = self.top.get();
let bottom = self.bottom.get().addr();
let extra = Bump::<A>::extra_bytes(top, layout.align());
let slow_path = Bump::<A>::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::<T>(), f());
NonNull::new_unchecked(new_top.cast::<T>())
}
}
} 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::<T>();
unsafe {
ptr.write(f());
}
ptr
};
Some(unsafe { ptr.as_mut() })
}
}
fn _asdf(bump: &mut Bump<super::PanicingAllocator>) {
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);
}

View file

@ -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<NonNull<[u8]>, 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<u8>, _layout: core::alloc::Layout) {}
}

View file

@ -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<A: Allocator + Clone> {
/// Size and alignment of each element in the slab.
element_size: usize,
/// Pointer to the first chunk in the slab.
head: Option<NonNull<SlabChunk<A>>>,
alloc: A,
}
struct SlabChunk<A: Allocator + Clone> {
/// Pointer to the next chunk in the slab.
next: Option<NonNull<SlabChunk<A>>>,
/// Pointer to the slab that owns this chunk.
slab: NonNull<Slab<A>>,
/// Linked list of free elements in this slab. When this is `None`, the
/// slab is full.
free: Cell<Option<NonNull<ChunkSlot>>>,
/// Number of outstanding allocations from this slab. When this reaches
/// zero, the slab can be freed.
count: Cell<usize>,
}
struct ChunkSlot(Option<NonNull<Self>>);
enum SlotResult {
Some(NonNull<u8>),
Last(NonNull<u8>),
None,
}
impl<A: Allocator + Clone> SlabChunk<A> {
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<u8>) -> bool {
let slot = slot.cast::<ChunkSlot>();
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<NonNull<Self>> {
self.0
}
}
impl<A: Allocator + Clone> Slab<A> {
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::<SlabChunk<A>>(), 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<SlabChunk<A>> {
// 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::<SlabChunk<A>>();
unsafe {
let first_slot = chunk
.as_ptr()
.byte_add(self.first_slot_offset())
.cast::<ChunkSlot>();
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<SlabChunk<A>> {
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<u8>) {
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::<SlabChunk<A>>();
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<A: Allocator + Clone> Drop for Slab<A> {
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<A: Allocator + Clone> {
/// Slabs for each power-of-two from 16 bytes to 2048 bytes (inclusive).
slabs: [Slab<A>; SLAB_ALLOCATOR_BUCKETS],
alloc: A,
}
impl<A: Allocator + Clone> SlabAllocator<A> {
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<usize> {
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<NonNull<[u8]>> {
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<u8>, 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::<SlabChunk<A>>();
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<A: Allocator + Clone> Drop for SlabAllocator<A> {
fn drop(&mut self) {}
}
#[cfg(test)]
mod tests {
use super::*;
use std::{prelude::rust_2024::*, sync::Mutex};
struct SlabAllocatorWrapper<A: Allocator + Clone>(Mutex<SlabAllocator<A>>);
impl<A: Allocator + Clone> SlabAllocatorWrapper<A> {
fn new(alloc: A) -> Self {
Self(Mutex::new(SlabAllocator::new(alloc)))
}
}
impl SlabAllocatorWrapper<std::alloc::System> {
fn new_system() -> SlabAllocatorWrapper<std::alloc::System> {
SlabAllocatorWrapper::new(std::alloc::System)
}
}
unsafe impl<A: Allocator + Clone> Allocator for SlabAllocatorWrapper<A> {
fn allocate(
&self,
layout: core::alloc::Layout,
) -> Result<NonNull<[u8]>, core::alloc::AllocError> {
self.0
.lock()
.unwrap()
.alloc(layout)
.ok_or(core::alloc::AllocError)
}
unsafe fn deallocate(&self, ptr: NonNull<u8>, 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;
}
}

View file

@ -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<T, R>(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<T: PrimInt>(value: T, alignment: T) -> T {
(value + alignment - T::one()) & !(alignment - T::one())
}
pub fn align_down<T: PrimInt>(value: T, alignment: T) -> T {
value & !(alignment - T::one())
}
/// # Safety
pub unsafe fn volatile_copy<T: Sized>(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<F: FnOnce()>(::core::mem::ManuallyDrop<F>);
impl<F: FnOnce()> DropGuard<F> {
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<F: FnOnce()> Drop for DropGuard<F> {
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()
);
}
}

View file

@ -0,0 +1,487 @@
mod spin_mutex {
use core::{
cell::UnsafeCell,
sync::atomic::{AtomicBool, Ordering},
};
pub struct SpinMutex<T: ?Sized> {
lock: AtomicBool,
data: UnsafeCell<T>,
}
pub struct SpinMutexGuard<'a, T: ?Sized + 'a> {
mutex: &'a SpinMutex<T>,
}
impl<T: ?Sized> core::ops::Deref for SpinMutexGuard<'_, T> {
type Target = T;
fn deref(&self) -> &Self::Target {
unsafe { &*self.mutex.data.get() }
}
}
impl<T: ?Sized> core::ops::DerefMut for SpinMutexGuard<'_, T> {
fn deref_mut(&mut self) -> &mut Self::Target {
unsafe { &mut *self.mutex.data.get() }
}
}
impl<T: ?Sized> Drop for SpinMutexGuard<'_, T> {
fn drop(&mut self) {
unsafe { self.mutex.unlock() }
}
}
unsafe impl<T: ?Sized + Send> Send for SpinMutex<T> {}
// SAFETY: SpinMutex properly synchronises access to T, but T must be Send
// to be safely shared between threads.
unsafe impl<T: ?Sized + Send> Sync for SpinMutex<T> {}
// 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<T: ?Sized + Sync> 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<T: ?Sized + Send> Send for SpinMutexGuard<'_, T> {}
impl<T> SpinMutex<T> {
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<T: ?Sized> SpinMutex<T> {
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<u8>,
}
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<F>(&self, f: F)
where
F: FnOnce(&OnceState),
{
if self.is_completed() {
return;
}
self.call_once_slow(false, f);
}
pub fn call_once_force<F>(&self, f: F)
where
F: FnOnce(&OnceState),
{
if self.is_completed() {
return;
}
self.call_once_slow(true, f);
}
#[cold]
pub fn call_once_slow<F>(&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<T> {
once: Once,
data: UnsafeCell<MaybeUninit<T>>,
}
unsafe impl<T: Send + Sync> Sync for OnceLock<T> {}
const impl<T> Default for OnceLock<T> {
fn default() -> Self {
Self {
once: Once::new(),
data: UnsafeCell::new(MaybeUninit::uninit()),
}
}
}
impl<T> OnceLock<T> {
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<F, E>(&self, f: F) -> Result<(), E>
where
F: FnOnce() -> Result<T, E>,
{
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<F>(&self, f: F) -> &T
where
F: FnOnce() -> T,
{
match self.get_or_try_init(|| Ok::<T, core::convert::Infallible>(f())) {
Ok(value) => value,
Err(_) => unreachable!(),
}
}
pub fn get_mut_or_init<F>(&mut self, f: F) -> &mut T
where
F: FnOnce() -> T,
{
match self.get_mut_or_try_init(|| Ok::<T, core::convert::Infallible>(f())) {
Ok(value) => value,
Err(_) => unreachable!(),
}
}
pub fn get_or_try_init<F, E>(&self, f: F) -> Result<&T, E>
where
F: FnOnce() -> Result<T, E>,
{
if !self.is_completed() {
self.initialize(f)?;
}
unsafe { Ok(self.get_unchecked()) }
}
pub fn get_mut_or_try_init<F, E>(&mut self, f: F) -> Result<&mut T, E>
where
F: FnOnce() -> Result<T, E>,
{
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<T> Drop for OnceLock<T> {
fn drop(&mut self) {
if self.is_completed() {
unsafe { self.data.get().as_mut_unchecked().assume_init_drop() }
}
}
}
pub struct LazyLock<T, F = fn() -> T> {
once: Once,
f: UnsafeCell<ManuallyDrop<F>>,
t: UnsafeCell<MaybeUninit<T>>,
}
unsafe impl<T: Sync + Send, F: Send> Sync for LazyLock<T, F> {}
impl<T, F> LazyLock<T, F>
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, F: FnOnce() -> T> core::ops::Deref for LazyLock<T, F> {
type Target = T;
fn deref(&self) -> &Self::Target {
LazyLock::force(self)
}
}
impl<T, F: FnOnce() -> T> core::ops::DerefMut for LazyLock<T, F> {
fn deref_mut(&mut self) -> &mut Self::Target {
LazyLock::force_mut(self)
}
}
impl<T, F> Drop for LazyLock<T, F> {
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};

View file

@ -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" }