Runtime polymorphic allocator for Rust (i.e. C++-style mixing for supporting smart pointers 😉).

dyn_allocator.rs

// Works with: nightly-2022-12-02

//! Runtime polymorphic allocator for Rust (i.e. C++-style mixing for supporting smart pointers 😉).
//!
//! Not sure why you'd ever want to use this, but here it is for when you need to.

#![allow(incomplete_features)] // generic_const_exprs
#![feature(allocator_api, ptr_metadata, generic_const_exprs)]
#![no_std]

use core::{
    alloc::{AllocError, Allocator, Layout},
    marker::PhantomData,
    mem::{align_of, size_of, MaybeUninit},
    panic::{RefUnwindSafe, UnwindSafe},
    ptr::{
        addr_of, addr_of_mut, drop_in_place, from_raw_parts, from_raw_parts_mut, metadata,
        DynMetadata, NonNull,
    },
};

mod private {
    use crate::AutoTraits;
    use core::mem::size_of;

    pub trait Sealed {}
    pub trait Extant {}
    impl<T> Extant for T {}
    pub trait Implements<AT: ?Sized + AutoTraits> {}

    pub const fn convert_width(n: usize) -> usize {
        (n / size_of::<*mut ()>()) + if n % size_of::<*mut ()>() == 0 { 0 } else { 1 }
    }
}
use private::{convert_width, Extant, Implements, Sealed};

impl Sealed for dyn Extant {}
impl Sealed for dyn Extant + Send {}
impl Sealed for dyn Extant + Sync {}
impl Sealed for dyn Extant + UnwindSafe {}
impl Sealed for dyn Extant + RefUnwindSafe {}
impl Sealed for dyn Extant + Send + Sync {}
impl Sealed for dyn Extant + Send + UnwindSafe {}
impl Sealed for dyn Extant + Send + RefUnwindSafe {}
impl Sealed for dyn Extant + Sync + UnwindSafe {}
impl Sealed for dyn Extant + Sync + RefUnwindSafe {}
impl Sealed for dyn Extant + UnwindSafe + RefUnwindSafe {}
impl Sealed for dyn Extant + Send + Sync + UnwindSafe {}
impl Sealed for dyn Extant + Send + Sync + RefUnwindSafe {}
impl Sealed for dyn Extant + Send + UnwindSafe + RefUnwindSafe {}
impl Sealed for dyn Extant + Sync + UnwindSafe + RefUnwindSafe {}
impl Sealed for dyn Extant + Send + Sync + UnwindSafe + RefUnwindSafe {}

pub trait AutoTraits: Sealed {}
impl AutoTraits for dyn Extant {}
impl AutoTraits for dyn Extant + Send {}
impl AutoTraits for dyn Extant + Sync {}
impl AutoTraits for dyn Extant + UnwindSafe {}
impl AutoTraits for dyn Extant + RefUnwindSafe {}
impl AutoTraits for dyn Extant + Send + Sync {}
impl AutoTraits for dyn Extant + Send + UnwindSafe {}
impl AutoTraits for dyn Extant + Send + RefUnwindSafe {}
impl AutoTraits for dyn Extant + Sync + UnwindSafe {}
impl AutoTraits for dyn Extant + Sync + RefUnwindSafe {}
impl AutoTraits for dyn Extant + UnwindSafe + RefUnwindSafe {}
impl AutoTraits for dyn Extant + Send + Sync + UnwindSafe {}
impl AutoTraits for dyn Extant + Send + Sync + RefUnwindSafe {}
impl AutoTraits for dyn Extant + Send + UnwindSafe + RefUnwindSafe {}
impl AutoTraits for dyn Extant + Sync + UnwindSafe + RefUnwindSafe {}
impl AutoTraits for dyn Extant + Send + Sync + UnwindSafe + RefUnwindSafe {}

impl<T: ?Sized> Implements<dyn Extant> for T where T: Extant {}
impl<T: ?Sized> Implements<dyn Extant + Send> for T where T: Extant + Send {}
impl<T: ?Sized> Implements<dyn Extant + Sync> for T where T: Extant + Sync {}
impl<T: ?Sized> Implements<dyn Extant + UnwindSafe> for T where T: Extant + UnwindSafe {}
impl<T: ?Sized> Implements<dyn Extant + RefUnwindSafe> for T where T: Extant + RefUnwindSafe {}
impl<T: ?Sized> Implements<dyn Extant + Send + Sync> for T where T: Extant + Send + Sync {}
impl<T: ?Sized> Implements<dyn Extant + Send + UnwindSafe> for T where T: Extant + Send + UnwindSafe {}
impl<T: ?Sized> Implements<dyn Extant + Send + RefUnwindSafe> for T where
    T: Extant + Send + RefUnwindSafe
{
}
impl<T: ?Sized> Implements<dyn Extant + Sync + UnwindSafe> for T where T: Extant + Sync + UnwindSafe {}
impl<T: ?Sized> Implements<dyn Extant + Sync + RefUnwindSafe> for T where
    T: Extant + Sync + UnwindSafe
{
}
impl<T: ?Sized> Implements<dyn Extant + UnwindSafe + RefUnwindSafe> for T where
    T: Extant + UnwindSafe + RefUnwindSafe
{
}
impl<T: ?Sized> Implements<dyn Extant + Send + Sync + UnwindSafe> for T where
    T: Extant + Send + Sync + UnwindSafe
{
}
impl<T: ?Sized> Implements<dyn Extant + Send + Sync + RefUnwindSafe> for T where
    T: Extant + Send + Sync + RefUnwindSafe
{
}
impl<T: ?Sized> Implements<dyn Extant + Send + UnwindSafe + RefUnwindSafe> for T where
    T: Extant + Send + UnwindSafe + RefUnwindSafe
{
}
impl<T: ?Sized> Implements<dyn Extant + Sync + UnwindSafe + RefUnwindSafe> for T where
    T: Extant + Sync + UnwindSafe + RefUnwindSafe
{
}
impl<T: ?Sized> Implements<dyn Extant + Send + Sync + UnwindSafe + RefUnwindSafe> for T where
    T: Extant + Send + Sync + UnwindSafe + RefUnwindSafe
{
}

#[derive(Debug)]
pub struct DynAllocator<'a, AT: ?Sized + AutoTraits, const MAX_ALLOCATOR_SIZE: usize>
where
    [(); convert_width(MAX_ALLOCATOR_SIZE)]:,
{
    meta: DynMetadata<dyn 'a + Allocator>,
    instance: MaybeUninit<[*mut (); convert_width(MAX_ALLOCATOR_SIZE)]>,
    _phantom: PhantomData<AT>,
}

unsafe impl<AT: ?Sized + AutoTraits, const MAX_ALLOCATOR_SIZE: usize> Send
    for DynAllocator<'_, AT, MAX_ALLOCATOR_SIZE>
where
    [(); convert_width(MAX_ALLOCATOR_SIZE)]:,
{
}
unsafe impl<AT: ?Sized + AutoTraits, const MAX_ALLOCATOR_SIZE: usize> Sync
    for DynAllocator<'_, AT, MAX_ALLOCATOR_SIZE>
where
    [(); convert_width(MAX_ALLOCATOR_SIZE)]:,
{
}
impl<AT: ?Sized + AutoTraits, const MAX_ALLOCATOR_SIZE: usize> UnwindSafe
    for DynAllocator<'_, AT, MAX_ALLOCATOR_SIZE>
where
    [(); convert_width(MAX_ALLOCATOR_SIZE)]:,
{
}
impl<AT: ?Sized + AutoTraits, const MAX_ALLOCATOR_SIZE: usize> RefUnwindSafe
    for DynAllocator<'_, AT, MAX_ALLOCATOR_SIZE>
where
    [(); convert_width(MAX_ALLOCATOR_SIZE)]:,
{
}
impl<AT: ?Sized + AutoTraits, const MAX_ALLOCATOR_SIZE: usize> Unpin
    for DynAllocator<'_, AT, MAX_ALLOCATOR_SIZE>
where
    [(); convert_width(MAX_ALLOCATOR_SIZE)]:,
{
}

impl<'a, AT: ?Sized + AutoTraits + 'a, const MAX_ALLOCATOR_SIZE: usize>
    DynAllocator<'a, AT, MAX_ALLOCATOR_SIZE>
where
    [(); convert_width(MAX_ALLOCATOR_SIZE)]:,
{
    pub fn new<A: Allocator + 'static>(a: A) -> Self
    where
        A: Implements<AT>,
        [(); MAX_ALLOCATOR_SIZE - size_of::<A>()]:,

        // This causes "cycle detected" with convert_width, but the result here is the same either way.
        [(); align_of::<MaybeUninit<[*mut (); MAX_ALLOCATOR_SIZE]>>() - align_of::<A>()]:,
    {
        Self {
            meta: metadata::<dyn Allocator>(addr_of!(a)),
            instance: unsafe {
                let mut memory = MaybeUninit::<[_; convert_width(MAX_ALLOCATOR_SIZE)]>::uninit();
                memory.as_mut_ptr().cast::<A>().write(a);
                memory
            },
            _phantom: PhantomData,
        }
    }

    pub fn widen<AT2: ?Sized + AutoTraits + 'a, const MAX_ALLOCATOR_SIZE_2: usize>(
        self,
    ) -> DynAllocator<'a, AT2, MAX_ALLOCATOR_SIZE_2>
    where
        [(); convert_width(MAX_ALLOCATOR_SIZE_2)]:,
        [(); MAX_ALLOCATOR_SIZE_2 - MAX_ALLOCATOR_SIZE]:,
        &'a AT: Into<&'a AT2>,
    {
        DynAllocator {
            meta: self.meta,
            instance: unsafe {
                let mut memory = MaybeUninit::<[_; convert_width(MAX_ALLOCATOR_SIZE_2)]>::uninit();
                memory
                    .as_mut_ptr()
                    .cast::<MaybeUninit<[*mut (); convert_width(MAX_ALLOCATOR_SIZE)]>>()
                    .write(self.instance);
                memory
            },
            _phantom: PhantomData,
        }
    }
}

impl<AT: ?Sized + AutoTraits, const MAX_ALLOCATOR_SIZE: usize> Drop
    for DynAllocator<'_, AT, MAX_ALLOCATOR_SIZE>
where
    [(); convert_width(MAX_ALLOCATOR_SIZE)]:,
{
    fn drop(&mut self) {
        unsafe {
            drop_in_place(from_raw_parts_mut::<dyn Allocator>(
                addr_of_mut!(self.instance).cast(),
                self.meta,
            ));
        }
    }
}

impl<'a, AT: ?Sized + AutoTraits, const MAX_ALLOCATOR_SIZE: usize> AsRef<dyn Allocator + 'a>
    for DynAllocator<'a, AT, MAX_ALLOCATOR_SIZE>
where
    [(); convert_width(MAX_ALLOCATOR_SIZE)]:,
{
    fn as_ref(&self) -> &(dyn Allocator + 'a) {
        unsafe { &*from_raw_parts(self.instance.as_ptr().cast(), self.meta) }
    }
}

unsafe impl<AT: ?Sized + AutoTraits, const MAX_ALLOCATOR_SIZE: usize> Allocator
    for DynAllocator<'_, AT, MAX_ALLOCATOR_SIZE>
where
    [(); convert_width(MAX_ALLOCATOR_SIZE)]:,
{
    fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
        self.as_ref().allocate(layout)
    }

    unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
        self.as_ref().deallocate(ptr, layout)
    }

    fn allocate_zeroed(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
        self.as_ref().allocate_zeroed(layout)
    }

    unsafe fn grow(
        &self,
        ptr: NonNull<u8>,
        old_layout: Layout,
        new_layout: Layout,
    ) -> Result<NonNull<[u8]>, AllocError> {
        self.as_ref().grow(ptr, old_layout, new_layout)
    }

    unsafe fn grow_zeroed(
        &self,
        ptr: NonNull<u8>,
        old_layout: Layout,
        new_layout: Layout,
    ) -> Result<NonNull<[u8]>, AllocError> {
        self.as_ref().grow_zeroed(ptr, old_layout, new_layout)
    }

    unsafe fn shrink(
        &self,
        ptr: NonNull<u8>,
        old_layout: Layout,
        new_layout: Layout,
    ) -> Result<NonNull<[u8]>, AllocError> {
        self.as_ref().shrink(ptr, old_layout, new_layout)
    }

    fn by_ref(&self) -> &Self
    where
        Self: Sized,
    {
        self
    }
}