use core::alloc::{AllocError, Allocator, Layout};
use core::cell::{Cell, UnsafeCell};
use core::mem::MaybeUninit;
use core::ptr::{self, NonNull};
use crate::{MimallocArena, alloc_result, mimalloc};
#[repr(C)] pub struct StackFallback<const N: usize, A: Allocator = std::alloc::Global> {
cur: Cell<usize>,
#[cfg(debug_assertions)]
got: Cell<bool>,
fallback: A,
buf: UnsafeCell<[MaybeUninit<u8>; N]>,
}
impl<const N: usize, A: Allocator> StackFallback<N, A> {
#[inline]
pub const fn new(fallback: A) -> Self {
Self {
cur: Cell::new(0),
#[cfg(debug_assertions)]
got: Cell::new(false),
fallback,
buf: UnsafeCell::new([MaybeUninit::uninit(); N]),
}
}
#[inline]
pub fn get(&self) -> &Self {
#[cfg(debug_assertions)]
{
assert!(!self.got.replace(true), "StackFallback::get called twice");
}
self.cur.set(0);
self
}
#[inline]
pub fn reset(&mut self) {
self.cur.set(0);
#[cfg(debug_assertions)]
self.got.set(false);
}
#[inline]
pub fn fallback(&self) -> &A {
&self.fallback
}
#[inline(always)]
fn buf_base(&self) -> *mut u8 {
self.buf.get().cast::<u8>()
}
#[inline]
pub fn owns(&self, p: *const u8) -> bool {
let base = self.buf_base().addr();
let q = p.addr();
q >= base && q < base.wrapping_add(N)
}
#[inline(always)]
fn is_last(&self, p: *const u8, len: usize) -> bool {
p.addr().wrapping_add(len) == self.buf_base().addr().wrapping_add(self.cur.get())
}
#[inline]
fn bump(&self, layout: Layout) -> Option<NonNull<u8>> {
let base = self.buf_base().addr();
let align = layout.align();
let adjusted = base.wrapping_add(self.cur.get()).checked_add(align - 1)? & !(align - 1);
let start = adjusted.wrapping_sub(base);
let end = start.checked_add(layout.size())?;
if end > N {
return None;
}
self.cur.set(end);
Some(unsafe { NonNull::new_unchecked(self.buf_base().add(start)) })
}
#[inline]
#[allow(clippy::mut_from_ref)]
pub fn alloc<T>(&self, val: T) -> &mut T {
let p = (&self)
.allocate(Layout::new::<T>())
.unwrap_or_else(|_| crate::out_of_memory())
.cast::<T>();
unsafe {
p.as_ptr().write(val);
&mut *p.as_ptr()
}
}
}
impl<const N: usize> StackFallback<N, std::alloc::Global> {
#[inline]
pub const fn with_global() -> Self {
Self::new(std::alloc::Global)
}
}
unsafe impl<const N: usize, A: Allocator> Allocator for &StackFallback<N, A> {
#[inline]
fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
if let Some(p) = self.bump(layout) {
return Ok(NonNull::slice_from_raw_parts(p, layout.size()));
}
self.fallback.allocate(layout)
}
#[inline]
unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
if self.owns(ptr.as_ptr()) {
if self.is_last(ptr.as_ptr(), layout.size()) {
self.cur.set(self.cur.get() - layout.size());
}
} else {
unsafe { self.fallback.deallocate(ptr, layout) }
}
}
#[inline]
unsafe fn grow(
&self,
ptr: NonNull<u8>,
old: Layout,
new: Layout,
) -> Result<NonNull<[u8]>, AllocError> {
if self.owns(ptr.as_ptr()) {
if self.is_last(ptr.as_ptr(), old.size()) {
let add = new.size() - old.size();
if self.cur.get() + add <= N {
self.cur.set(self.cur.get() + add);
return Ok(NonNull::slice_from_raw_parts(ptr, new.size()));
}
}
let newp = self.allocate(new)?;
unsafe {
ptr::copy_nonoverlapping(ptr.as_ptr(), newp.as_ptr().cast::<u8>(), old.size());
self.deallocate(ptr, old);
}
Ok(newp)
} else {
unsafe { self.fallback.grow(ptr, old, new) }
}
}
#[inline]
unsafe fn shrink(
&self,
ptr: NonNull<u8>,
old: Layout,
new: Layout,
) -> Result<NonNull<[u8]>, AllocError> {
if self.owns(ptr.as_ptr()) {
if self.is_last(ptr.as_ptr(), old.size()) {
self.cur.set(self.cur.get() - (old.size() - new.size()));
}
Ok(NonNull::slice_from_raw_parts(ptr, new.size()))
} else {
unsafe { self.fallback.shrink(ptr, old, new) }
}
}
}
#[derive(Clone, Copy)]
pub struct ArenaPtr {
arena: *const MimallocArena,
}
impl ArenaPtr {
#[inline]
pub const fn new(arena: *const MimallocArena) -> Self {
Self { arena }
}
#[inline]
pub const fn global() -> Self {
Self { arena: ptr::null() }
}
#[inline]
pub fn arena(&self) -> *const MimallocArena {
self.arena
}
#[inline]
pub fn set_arena(&mut self, arena: *const MimallocArena) {
self.arena = arena;
}
#[inline]
fn arena_ref(&self) -> Option<&MimallocArena> {
unsafe { self.arena.as_ref() }
}
}
unsafe impl Allocator for ArenaPtr {
#[inline]
fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
match self.arena_ref() {
Some(a) => a.allocate(layout),
None => {
let p = mimalloc::mi_malloc_auto_align(layout.size(), layout.align());
alloc_result(p, layout.size())
}
}
}
#[inline]
unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
unsafe {
match self.arena_ref() {
Some(a) => a.deallocate(ptr, layout),
None => mimalloc::mi_free(ptr.as_ptr().cast()),
}
}
}
#[inline]
unsafe fn grow(
&self,
ptr: NonNull<u8>,
old: Layout,
new: Layout,
) -> Result<NonNull<[u8]>, AllocError> {
unsafe {
match self.arena_ref() {
Some(a) => a.grow(ptr, old, new),
None => alloc_result(
mimalloc::mi_realloc_aligned(ptr.as_ptr().cast(), new.size(), new.align()),
new.size(),
),
}
}
}
#[inline]
unsafe fn shrink(
&self,
ptr: NonNull<u8>,
old: Layout,
new: Layout,
) -> Result<NonNull<[u8]>, AllocError> {
unsafe { self.grow(ptr, old, new) }
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::alloc::Global;
struct Counting {
allocs: Cell<usize>,
deallocs: Cell<usize>,
}
impl Counting {
const fn new() -> Self {
Self {
allocs: Cell::new(0),
deallocs: Cell::new(0),
}
}
}
unsafe impl Allocator for Counting {
fn allocate(&self, l: Layout) -> Result<NonNull<[u8]>, AllocError> {
self.allocs.set(self.allocs.get() + 1);
Global.allocate(l)
}
unsafe fn deallocate(&self, p: NonNull<u8>, l: Layout) {
self.deallocs.set(self.deallocs.get() + 1);
unsafe { Global.deallocate(p, l) }
}
}
#[test]
fn alloc_within_stack() {
let sf = StackFallback::<64, _>::new(Counting::new());
let a = &sf;
let p = a.allocate(Layout::from_size_align(16, 1).unwrap()).unwrap();
assert_eq!(p.len(), 16);
assert!(sf.owns(p.cast::<u8>().as_ptr()));
assert_eq!(sf.fallback().allocs.get(), 0);
let q = a.allocate(Layout::from_size_align(32, 1).unwrap()).unwrap();
assert!(sf.owns(q.cast::<u8>().as_ptr()));
assert_eq!(sf.fallback().allocs.get(), 0);
}
#[test]
fn overflow_to_fallback() {
let sf = StackFallback::<32, _>::new(Counting::new());
let a = &sf;
let _p = a.allocate(Layout::from_size_align(24, 1).unwrap()).unwrap();
assert_eq!(sf.fallback().allocs.get(), 0);
let q = a.allocate(Layout::from_size_align(16, 1).unwrap()).unwrap();
assert!(!sf.owns(q.cast::<u8>().as_ptr()));
assert_eq!(sf.fallback().allocs.get(), 1);
unsafe { a.deallocate(q.cast(), Layout::from_size_align(16, 1).unwrap()) };
assert_eq!(sf.fallback().deallocs.get(), 1);
}
#[test]
fn dealloc_range_check() {
let sf = StackFallback::<64, _>::new(Counting::new());
let a = &sf;
let l8 = Layout::from_size_align(8, 1).unwrap();
let p = a.allocate(l8).unwrap().cast::<u8>();
let q = a.allocate(l8).unwrap().cast::<u8>();
assert_eq!(sf.cur.get(), 16);
unsafe { a.deallocate(p, l8) };
assert_eq!(sf.cur.get(), 16);
unsafe { a.deallocate(q, l8) };
assert_eq!(sf.cur.get(), 8);
assert_eq!(sf.fallback().deallocs.get(), 0);
}
#[test]
fn grow_in_place() {
let sf = StackFallback::<64, _>::new(Counting::new());
let a = &sf;
let old = Layout::from_size_align(8, 1).unwrap();
let new = Layout::from_size_align(24, 1).unwrap();
let p = a.allocate(old).unwrap().cast::<u8>();
let g = unsafe { a.grow(p, old, new) }.unwrap();
assert_eq!(g.cast::<u8>().as_ptr(), p.as_ptr());
assert_eq!(sf.cur.get(), 24);
assert_eq!(sf.fallback().allocs.get(), 0);
}
#[test]
fn grow_spills_to_fallback() {
let sf = StackFallback::<32, _>::new(Counting::new());
let a = &sf;
let old = Layout::from_size_align(16, 1).unwrap();
let new = Layout::from_size_align(48, 1).unwrap();
let p = a.allocate(old).unwrap().cast::<u8>();
unsafe { ptr::write_bytes(p.as_ptr(), 0xAB, 16) };
let g = unsafe { a.grow(p, old, new) }.unwrap();
assert!(!sf.owns(g.cast::<u8>().as_ptr()));
assert_eq!(sf.fallback().allocs.get(), 1);
let bytes = unsafe { core::slice::from_raw_parts(g.cast::<u8>().as_ptr(), 16) };
assert!(bytes.iter().all(|&b| b == 0xAB));
unsafe { a.deallocate(g.cast(), new) };
}
#[test]
fn alignment() {
let sf = StackFallback::<128, _>::new(Counting::new());
let a = &sf;
let _ = a.allocate(Layout::from_size_align(1, 1).unwrap()).unwrap();
let p = a
.allocate(Layout::from_size_align(8, 16).unwrap())
.unwrap()
.cast::<u8>();
assert_eq!(p.as_ptr().addr() % 16, 0);
assert!(sf.owns(p.as_ptr()));
let q = a
.allocate(Layout::from_size_align(8, 256).unwrap())
.unwrap()
.cast::<u8>();
assert_eq!(q.as_ptr().addr() % 256, 0);
assert!(!sf.owns(q.as_ptr()));
unsafe { a.deallocate(q, Layout::from_size_align(8, 256).unwrap()) };
}
#[test]
fn shrink_rewinds_last() {
let sf = StackFallback::<64, _>::new(Counting::new());
let a = &sf;
let old = Layout::from_size_align(32, 1).unwrap();
let new = Layout::from_size_align(8, 1).unwrap();
let p = a.allocate(old).unwrap().cast::<u8>();
assert_eq!(sf.cur.get(), 32);
let s = unsafe { a.shrink(p, old, new) }.unwrap();
assert_eq!(s.cast::<u8>().as_ptr(), p.as_ptr());
assert_eq!(sf.cur.get(), 8);
}
#[test]
fn vec_roundtrip() {
let sf = StackFallback::<256>::with_global();
let mut v: Vec<u32, _> = Vec::new_in(&sf);
for i in 0..8 {
v.push(i);
}
assert!(sf.owns(v.as_ptr().cast()));
for i in 8..200 {
v.push(i);
}
assert!(!sf.owns(v.as_ptr().cast()));
assert_eq!(v.iter().copied().sum::<u32>(), (0..200).sum());
}
#[test]
#[cfg(debug_assertions)]
#[should_panic(expected = "StackFallback::get called twice")]
fn get_twice_panics() {
let sf = StackFallback::<16>::with_global();
let _ = sf.get();
let _ = sf.get();
}
#[test]
fn reset_clears_guard() {
let mut sf = StackFallback::<16>::with_global();
{
let a = sf.get();
let _ = a.allocate(Layout::from_size_align(8, 1).unwrap()).unwrap();
}
sf.reset();
let _ = sf.get(); assert_eq!(sf.cur.get(), 0);
}
}