use core::alloc::Layout;
#[cfg(feature = "alloc-xthread")]
use core::mem::ManuallyDrop;
#[cfg(feature = "alloc-xthread")]
use crate::alloc_core::os;
#[cfg(feature = "alloc-xthread")]
use crate::alloc_core::segment_header::{SegmentHeader, SegmentKind, SegmentMeta, SEGMENT_MAGIC};
use crate::alloc_core::AllocCore;
#[cfg(feature = "alloc-xthread")]
use super::thread_free::ThreadFreeStack;
pub struct Heap {
#[cfg(feature = "alloc-xthread")]
core: ManuallyDrop<AllocCore>,
#[cfg(not(feature = "alloc-xthread"))]
core: AllocCore,
#[cfg(feature = "alloc-xthread")]
thread_free: ManuallyDrop<ThreadFreeStack>,
}
impl Heap {
#[must_use]
pub fn new() -> Option<Self> {
let core = AllocCore::new()?;
Some(Self {
#[cfg(feature = "alloc-xthread")]
core: ManuallyDrop::new(core),
#[cfg(not(feature = "alloc-xthread"))]
core,
#[cfg(feature = "alloc-xthread")]
thread_free: ManuallyDrop::new(ThreadFreeStack::new()),
})
}
#[must_use]
pub fn alloc(&mut self, layout: Layout) -> *mut u8 {
let size = layout
.size()
.max(crate::alloc_core::size_classes::MIN_BLOCK);
let align = layout.align();
match classify(size, align) {
Some(class_idx) => self.alloc_small(class_idx),
None => {
let ptr = self.core.alloc(layout);
#[cfg(feature = "alloc-xthread")]
if !ptr.is_null() {
self.stamp_owner(ptr);
}
ptr
}
}
}
#[must_use]
pub fn alloc_zeroed(&mut self, layout: Layout) -> *mut u8 {
let ptr = self.alloc(layout);
if !ptr.is_null() {
crate::alloc_core::node::Node::zero(
ptr,
layout
.size()
.max(crate::alloc_core::size_classes::MIN_BLOCK),
);
}
ptr
}
pub fn dealloc(&mut self, ptr: *mut u8, layout: Layout) {
if ptr.is_null() {
return;
}
let size = layout
.size()
.max(crate::alloc_core::size_classes::MIN_BLOCK);
let align = layout.align();
match classify(size, align) {
Some(class_idx) => self.dealloc_small(ptr, class_idx),
None => self.core.dealloc(ptr, layout),
}
}
pub fn realloc(&mut self, ptr: *mut u8, old_layout: Layout, new_size: usize) -> *mut u8 {
if ptr.is_null() {
return core::ptr::null_mut();
}
let new_layout = match Layout::from_size_align(new_size, old_layout.align()) {
Ok(l) => l,
Err(_) => return core::ptr::null_mut(),
};
let new_ptr = self.alloc(new_layout);
if new_ptr.is_null() {
return core::ptr::null_mut();
}
let copy = old_layout.size().min(new_size);
crate::alloc_core::node::Node::copy_nonoverlapping(ptr, new_ptr, copy);
self.dealloc(ptr, old_layout);
new_ptr
}
#[cfg(feature = "alloc-xthread")]
pub fn dealloc_any_thread(ptr: *mut u8, layout: Layout) {
if ptr.is_null() {
return;
}
let base = os::segment_base_of_ptr(ptr);
if SegmentHeader::magic_at(base) != SEGMENT_MAGIC {
return; }
if SegmentHeader::kind_at(base) == SegmentKind::Large {
return;
}
if SegmentHeader::owner_thread_free_at(base).is_null() {
return;
}
let size = layout
.size()
.max(crate::alloc_core::size_classes::MIN_BLOCK);
let class_idx =
match crate::alloc_core::size_classes::SizeClasses::class_for(size, layout.align()) {
Some(c) => c as u32,
None => return,
};
let off = (ptr as usize - base as usize) as u32;
let packed = crate::alloc_core::remote_free_ring::pack_entry(off, class_idx);
let ring = SegmentMeta::new(base).remote_ring();
let _ = ring.push(packed);
}
fn alloc_small(&mut self, class_idx: usize) -> *mut u8 {
let block_size = crate::alloc_core::size_classes::SizeClasses::block_size(class_idx);
let layout = match Layout::from_size_align(block_size, block_size.min(16)) {
Ok(l) => l,
Err(_) => return core::ptr::null_mut(),
};
let ptr = self.core.alloc(layout);
#[cfg(feature = "alloc-xthread")]
if !ptr.is_null() {
self.stamp_owner(ptr);
}
ptr
}
fn dealloc_small(&mut self, ptr: *mut u8, class_idx: usize) {
#[cfg(not(feature = "alloc-xthread"))]
{
let block_size = crate::alloc_core::size_classes::SizeClasses::block_size(class_idx);
let layout = match Layout::from_size_align(block_size, block_size.min(16)) {
Ok(l) => l,
Err(_) => return,
};
self.core.dealloc(ptr, layout);
}
#[cfg(feature = "alloc-xthread")]
{
let base = os::segment_base_of_ptr(ptr);
if SegmentHeader::magic_at(base) != SEGMENT_MAGIC {
return; }
let our_head = self.thread_free.head_ptr();
let owner_tf = SegmentHeader::owner_thread_free_at(base);
if owner_tf.is_null() || owner_tf == our_head {
let block_size =
crate::alloc_core::size_classes::SizeClasses::block_size(class_idx);
let layout = match Layout::from_size_align(block_size, block_size.min(16)) {
Ok(l) => l,
Err(_) => return,
};
self.core.dealloc(ptr, layout);
} else {
let off = (ptr as usize - base as usize) as u32;
let packed = crate::alloc_core::remote_free_ring::pack_entry(off, class_idx as u32);
let ring = SegmentMeta::new(base).remote_ring();
let _ = ring.push(packed);
}
}
}
#[cfg(feature = "alloc-xthread")]
fn stamp_owner(&mut self, ptr: *mut u8) {
let base = os::segment_base_of_ptr(ptr);
let mut meta = SegmentMeta::new(base);
if SegmentHeader::owner_thread_free_at(base).is_null() {
meta.stamp_owner_thread_free(self.thread_free.head_ptr());
}
}
}
impl Drop for Heap {
fn drop(&mut self) {
#[cfg(feature = "alloc-xthread")]
{
}
}
}
fn classify(size: usize, align: usize) -> Option<usize> {
crate::alloc_core::size_classes::SizeClasses::class_for(size, align)
}