kernel: pmm with overlapping check
This commit is contained in:
parent
86d80132be
commit
ae06f75cc2
SSH key fingerprint:
SHA256:bB1qbbqmDXZNT0KKD5c2Dfjg53JGhj7B3CFcLIzSqq8
|
|
@ -8,6 +8,7 @@
|
|||
ptr_cast_slice,
|
||||
likely_unlikely,
|
||||
int_roundings,
|
||||
slice_ptr_get,
|
||||
never_type
|
||||
)]
|
||||
#![cfg_attr(test, feature(custom_test_frameworks))]
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
|
|
|
|||
|
|
@ -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::<PageHeader>().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<PhysicalPageNode>,
|
||||
}
|
||||
|
||||
#[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<NonNull<Self>> {
|
||||
self.node.left().map(Self::unsafe_from_node)
|
||||
}
|
||||
|
||||
fn right(&self) -> Option<NonNull<Self>> {
|
||||
self.node.right().map(Self::unsafe_from_node)
|
||||
}
|
||||
|
||||
fn parent(&self) -> Option<NonNull<Self>> {
|
||||
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<NonNull<Self>>) {
|
||||
self.node.set_left(left.map(Self::unsafe_into_node))
|
||||
}
|
||||
|
||||
fn set_right(&self, right: Option<NonNull<Self>>) {
|
||||
self.node.set_right(right.map(Self::unsafe_into_node))
|
||||
}
|
||||
|
||||
fn set_parent(&self, parent: Option<NonNull<Self>>) {
|
||||
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> {
|
||||
PhyAddr::from_hhdm_virt(self)
|
||||
}
|
||||
|
||||
fn unsafe_from_node(node: NonNull<PhysicalPageNode>) -> NonNull<Self> {
|
||||
node.cast::<Self>()
|
||||
}
|
||||
|
||||
fn unsafe_into_node(ptr: NonNull<Self>) -> NonNull<PhysicalPageNode> {
|
||||
ptr.cast::<PhysicalPageNode>()
|
||||
}
|
||||
|
||||
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<usize> {
|
||||
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<Self> {
|
||||
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::<PageHeader>().read_volatile() };
|
||||
fn free_usable_regions_from_memory_map(
|
||||
memory_map: &[crate::boot::MemoryRegion],
|
||||
) -> RBTree<PageHeader> {
|
||||
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);
|
||||
|
||||
let None = tree.insert_node(node_ptr) else {
|
||||
panic!(
|
||||
"Failed to insert node for page_index = {}, count = {}",
|
||||
idx, count
|
||||
);
|
||||
};
|
||||
|
||||
tree
|
||||
});
|
||||
|
||||
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 usable_regions = memory_map
|
||||
.iter()
|
||||
.copied()
|
||||
.filter(|region| region.region_type.is_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;
|
||||
|
||||
let mut pmm = PhysicalMemoryAllocator {
|
||||
tree: RBTree::new(),
|
||||
};
|
||||
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::<PageHeader>()
|
||||
.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.
|
||||
|
||||
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;
|
||||
serial_println!("Found overlapping region: {:?}", overlapping);
|
||||
|
||||
pmm.free_region(idx as usize, count as usize);
|
||||
});
|
||||
let Range {
|
||||
start: overlap_idx,
|
||||
end: overlap_end,
|
||||
} = overlapping.page_range();
|
||||
let overlap_count = overlap_end - overlap_idx;
|
||||
|
||||
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::<PhysicalPageNode>();
|
||||
|
||||
let None = self.tree.insert_node(node_ptr) else {
|
||||
panic!(
|
||||
"Failed to insert node for page_index = {}, count = {}",
|
||||
page_index, count
|
||||
);
|
||||
// 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<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 {
|
||||
foundation::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(foundation::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());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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(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,
|
||||
|
|
|
|||
Loading…
Reference in a new issue