You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
108 lines
3.0 KiB
108 lines
3.0 KiB
//! 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)
|
|
}
|
|
}
|