rsh/src/ext.rs

//! Extensions
use super::*;
use std::mem::{self, MaybeUninit};
use std::iter;
use smallvec::SmallVec;

/// Max size of memory allowed to be allocated on the stack.
pub const STACK_MEM_ALLOC_MAX: usize = 4096;

/// A stack-allocated vector that spills onto the heap when needed.
pub type StackVec<T> = SmallVec<[T; STACK_MEM_ALLOC_MAX]>;

/// A maybe-atom that can spill into a vector.
pub type MaybeVec<T> = SmallVec<[T; 1]>;

/// Allocate a vector of `MaybeUninit<T>`.
pub fn vec_uninit<T>(sz: usize) -> Vec<MaybeUninit<T>>
{
    let mut mem: Vec<T> = Vec::with_capacity(sz);
    unsafe {
	mem.set_len(sz);
	mem::transmute(mem)
    }
}

/// Allocate a local buffer initialised from `init`.
pub fn alloc_local_with<T, U>(sz: usize, init: impl FnMut() -> T, within: impl FnOnce(&mut [T]) -> U) -> U
{
    if sz > STACK_MEM_ALLOC_MAX {
	let mut memory: Vec<T> = iter::repeat_with(init).take(sz).collect();
	within(&mut memory[..])
    } else {
	stackalloc::stackalloc_with(sz, init, within)
    }
}



/// Allocate a local zero-initialised byte buffer
pub fn alloc_local_bytes<U>(sz: usize, within: impl FnOnce(&mut [u8]) -> U) -> U
{
    if sz > STACK_MEM_ALLOC_MAX {
	let mut memory: Vec<MaybeUninit<u8>> = vec_uninit(sz);
	within(unsafe {
	    std::ptr::write_bytes(memory.as_mut_ptr(), 0, sz);
	    stackalloc::helpers::slice_assume_init_mut(&mut memory[..])
	})
    } else {
	stackalloc::alloca_zeroed(sz, within)
    }
}


/// Allocate a local zero-initialised buffer
pub fn alloc_local_zeroed<T, U>(sz: usize, within: impl FnOnce(&mut [MaybeUninit<T>]) -> U) -> U
{
    let sz_bytes = mem::size_of::<T>() * sz;
    if sz > STACK_MEM_ALLOC_MAX {
	let mut memory = vec_uninit(sz);
	unsafe {
	    std::ptr::write_bytes(memory.as_mut_ptr(), 0, sz_bytes);
	}
	within(&mut memory[..])
    } else {
	stackalloc::alloca_zeroed(sz_bytes, move |buf| {
	    unsafe {
		debug_assert_eq!(buf.len() / mem::size_of::<T>(), sz);
		within(std::slice::from_raw_parts_mut(buf.as_mut_ptr() as *mut MaybeUninit<T>, sz))
	    }
	})
    }
}


/// Allocate a local uninitialised buffer
pub fn alloc_local_uninit<T, U>(sz: usize, within: impl FnOnce(&mut [MaybeUninit<T>]) -> U) -> U
{
    if sz > STACK_MEM_ALLOC_MAX {
	let mut memory = vec_uninit(sz);
	within(&mut memory[..])
    } else {
	stackalloc::stackalloc_uninit(sz, within)
    }
}


/// Allocate a local buffer initialised with `init`.
pub fn alloc_local<T: Clone, U>(sz: usize, init: T, within: impl FnOnce(&mut [T]) -> U) -> U
{
    if sz > STACK_MEM_ALLOC_MAX {
	let mut memory: Vec<T> = iter::repeat(init).take(sz).collect();
	within(&mut memory[..])
    } else {
	stackalloc::stackalloc(sz, init, within)
    }
}

/// Allocate a local buffer initialised with `T::default()`.
pub fn alloc_local_with_default<T: Default, U>(sz: usize, within: impl FnOnce(&mut [T]) -> U) -> U
{
    if sz > STACK_MEM_ALLOC_MAX {
	let mut memory: Vec<T> = iter::repeat_with(Default::default).take(sz).collect();
	within(&mut memory[..])
    } else {
	stackalloc::stackalloc_with_default(sz, within)
    }
}

/// Create a blanket-implementing trait that is a subtrait of any number of traits.
///
/// # Usage
/// To subtype a trait with `Sized` types only:
/// ```
/// # use rsh::subtrait;
/// subtrait!(pub SizedUnpinAsyncRead; "This subtrait must be sized": std::marker::Unpin + tokio::io::AsyncRead);
/// ```
/// To impl the trait for unsized types as well:
/// ```
/// # use rsh::subtrait;
/// subtrait!(pub dyn UnsizedUnpinAsyncRead; "Subtraits `Unpin` and `AsyncRead`. Also implemented for unsized types.": std::marker::Unpin + tokio::io::AsyncRead);
/// ```
/// # Downcastable subtrait
/// To create a subtrait that can be downcasted (correctly coerces to `dyn Any`), use the `virtual` keyword.
/// ```
/// # use rsh::subtrait;
/// subtrait!(pub virtual DowncastableUnpinAsyncRead; "Subtraits `Unpin` and `AsyncRead`. Also dynamically downcastable": std::marker::Unpin + tokio::io::AsyncRead);
/// subtrait!(pub virtual AnotherAny; "A dynamically downcastable trait that subtraits nothing else but `Any`");
/// ```
#[macro_export] macro_rules! subtrait {
    (@ $v:vis $name:ident $(; $doc:literal)?: $($t:path),+) => {
	$(#[doc = $doc])?
	    $v trait $name: $($t+)+ {}
    };
    ($v:vis $name:ident $(; $doc:literal)?: $($t:path),+) => {
	$crate::subtrait!(@ $v $name $(; $doc)?: Sized, $($t),+);
	impl<T: $($t+)+> $name for T{}
    };
    ($v:vis dyn $name:ident $(; $doc:literal)?: $($t:path),+) => {
	$crate::subtrait!(@ $v $name $(; $doc)?: $($t),+);
	impl<T: ?Sized + $($t+)+> $name for T{}
    };
    ($v:vis virtual $name:ident $(; $doc:literal)?$(: $($t:path),*)?) => {
	$crate::subtrait!(@ $v $name $(; $doc)?: ::mopa::Any $(, $($t),*)?);
	const _:() = {
	    fn _assert_object_safe(_: &dyn $name){}
	};
	mopafy!($name);
	impl<T: ?Sized + ::mopa::Any $(+ $($t+)+)?> $name for T {}
    };
}

// Static testing of `subtrait`
const _:() = {
    use mopa::*;
    subtrait!(dyn Test; "test dynamic subtrait": std::marker::Unpin);
    subtrait!(virtual TestAny; "test downcastable subtrait of `Read` and `Write`": std::io::Read);
    subtrait!(virtual TestAny2; "test downcastable subtrait of only `Any`");
    const fn _a<T: Test+?Sized>(){}
    const fn _b<T: TestAny+?Sized>(){}
    fn _c(a: Box<dyn TestAny>)
    {
	if a.is::<std::io::Stdin>() {
	    let _ref: &std::io::Stdin = a.downcast_ref::<std::io::Stdin>().unwrap();
	    drop(_ref);
	    let _unbox: std::io::Stdin = *a.downcast().map_err(|_| ()).unwrap();
	}
    }
    fn _d(a: &dyn TestAny)
    {
	if a.is::<std::io::Stdin>() {
	    let _ref: &std::io::Stdin = a.downcast_ref::<std::io::Stdin>().unwrap();
	}
    }

    const _TEST: () = _a::<dyn Test>();
    const _TEST2: () = _b::<dyn TestAny>();
};