use core::{alloc, mem};
use alloc_traits::{NonZeroLayout, LocalAlloc};
use super::{
boxed::Box,
fixed_vec::FixedVec,
rc::Rc,
uninit::Uninit,
};
#[derive(Debug)]
pub struct LeakedAllocation<'a, T: ?Sized=()> {
pub uninit: Uninit<'a, T>,
}
pub trait LocalAllocLeakExt<'alloc>: LocalAlloc<'alloc> {
fn alloc_layout(&'alloc self, layout: NonZeroLayout)
-> Option<LeakedAllocation<'alloc>>
{
let alloc = self.alloc(layout)?;
let uninit = unsafe {
Uninit::from_memory(alloc.ptr, alloc.layout.size().into())
};
Some(LeakedAllocation {
uninit,
})
}
fn alloc_t<V>(&'alloc self) -> Option<LeakedAllocation<'alloc, V>> {
match NonZeroLayout::new::<V>() {
None => Some(LeakedAllocation::zst_fake_alloc()),
Some(alloc) => {
let allocation = self.alloc_layout(alloc)?;
let right_type = allocation.cast().unwrap();
Some(right_type)
},
}
}
fn boxed<V>(&'alloc self, val: V) -> Option<Box<'alloc, V>> {
let alloc = self.alloc_t::<V>()?;
Some(Box::new(val, alloc.uninit))
}
fn fixed_vec<V>(&'alloc self, capacity: usize) -> Option<FixedVec<'alloc, V>> {
let size = mem::size_of::<V>().checked_mul(capacity)?;
let layout = alloc::Layout::from_size_align(size, mem::align_of::<V>()).ok()?;
let uninit = match NonZeroLayout::from_layout(layout.into()) {
None => Uninit::empty(),
Some(layout) => {
let allocation = self.alloc_layout(layout)?;
let right_type = allocation.cast_slice().unwrap();
right_type.uninit
}
};
Some(FixedVec::new(uninit))
}
fn rc<V>(&'alloc self, val: V) -> Option<Rc<'alloc, V>> {
let layout = Rc::<V>::layout();
let layout = NonZeroLayout::from_layout(layout.into()).unwrap();
let alloc = self.alloc_layout(layout)?;
Some(Rc::new(val, alloc.uninit))
}
}
impl<'alloc, T> LocalAllocLeakExt<'alloc> for T
where T: LocalAlloc<'alloc>,
{ }
impl<Zst> LeakedAllocation<'_, Zst> {
pub fn zst_fake_alloc() -> Self {
LeakedAllocation {
uninit: Uninit::invent_for_zst(),
}
}
}
impl<'a, T> LeakedAllocation<'a, T> {
fn cast<U>(self) -> Option<LeakedAllocation<'a, U>> {
Some(LeakedAllocation {
uninit: self.uninit.cast().ok()?,
})
}
fn cast_slice<U>(self) -> Option<LeakedAllocation<'a, [U]>> {
Some(LeakedAllocation {
uninit: self.uninit.cast_slice().ok()?,
})
}
}