use std::mem::MaybeUninit;
#[repr(align(64))]
struct Inline<const CAP: usize> {
_bytes: [u8; CAP],
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum Mode {
Empty,
Inline,
Heap,
}
pub(super) struct SmallSlot<const CAP: usize> {
inline: MaybeUninit<Inline<CAP>>,
ptr: *mut u8,
drop: unsafe fn(*mut u8),
mode: Mode,
}
impl<const CAP: usize> Default for SmallSlot<CAP> {
fn default() -> Self {
Self::new()
}
}
impl<const CAP: usize> SmallSlot<CAP> {
#[must_use]
pub(super) fn new() -> Self {
Self {
inline: MaybeUninit::uninit(),
ptr: std::ptr::null_mut(),
drop: Self::drop_nop,
mode: Mode::Empty,
}
}
pub(super) fn clear(&mut self) {
if self.mode == Mode::Empty {
return;
}
unsafe { (self.drop)(self.ptr) };
self.ptr = std::ptr::null_mut();
self.drop = Self::drop_nop;
self.mode = Mode::Empty;
}
pub(super) fn put<T: 'static>(&mut self, value: T) {
self.clear();
if std::mem::size_of::<T>() <= CAP
&& std::mem::align_of::<T>() <= std::mem::align_of::<Inline<CAP>>()
{
let ptr = self.inline.as_mut_ptr().cast::<u8>();
unsafe { (ptr as *mut T).write(value) };
self.ptr = ptr;
self.drop = Self::drop_inline::<T>;
self.mode = Mode::Inline;
return;
}
let raw = Box::into_raw(Box::new(value)) as *mut u8;
self.ptr = raw;
self.drop = Self::drop_heap::<T>;
self.mode = Mode::Heap;
}
unsafe fn drop_inline<T>(raw: *mut u8) {
unsafe { std::ptr::drop_in_place(raw as *mut T) };
}
unsafe fn drop_heap<T>(raw: *mut u8) {
unsafe {
drop(Box::from_raw(raw as *mut T));
}
}
unsafe fn drop_nop(_raw: *mut u8) {}
}
impl<const CAP: usize> Drop for SmallSlot<CAP> {
fn drop(&mut self) {
self.clear();
}
}