alloca/
lib.rs

1//! Mostly safe no_std wrapper for `alloca` in Rust.
2//!
3//! This crate uses Rust lifetime system to ensure that stack allocated memory will not be used
4//! after function return, but it does not make any guarantee about memory that is turned into
5//! raw pointer and stored somewhere else.
6//!
7//! ### Example
8//! ```rust
9//! // allocate 128 bytes on the stack
10//! alloca::with_alloca(128, |memory| {
11//!     // memory: &mut [MaybeUninit<u8>]
12//!     assert_eq!(memory.len(), 128);
13//! });
14//! ```
15
16#![no_std]
17
18use core::{
19    ffi::c_void,
20    mem::{self, ManuallyDrop, MaybeUninit},
21    ptr, slice,
22};
23
24type Callback = unsafe extern "C-unwind" fn(ptr: *mut MaybeUninit<u8>, data: *mut c_void);
25
26extern "C-unwind" {
27    fn c_with_alloca(size: usize, callback: Callback, data: *mut c_void);
28}
29
30#[inline(always)]
31fn get_trampoline<F: FnOnce(*mut MaybeUninit<u8>)>(_closure: &F) -> Callback {
32    trampoline::<F>
33}
34
35unsafe extern "C-unwind" fn trampoline<F: FnOnce(*mut MaybeUninit<u8>)>(
36    ptr: *mut MaybeUninit<u8>,
37    data: *mut c_void,
38) {
39    let f = ManuallyDrop::take(&mut *(data as *mut ManuallyDrop<F>));
40    f(ptr);
41}
42
43/// Allocates `[u8; size]` memory on stack and invokes `closure` with this slice as argument.
44///
45/// # Safety
46/// This function is safe because `c_with_alloca` (which is internally used) will always returns non-null
47/// pointer.
48///
49/// # Potential segfaults or UB
50///
51/// When using this function in wrong way your program might get UB or segfault "for free":
52/// - Using memory allocated by `with_alloca` outside of it e.g closure is already returned but you somehow
53/// managed to store pointer to memory and use it.
54/// - Allocating more memory than thread stack size.
55///
56///   This will trigger segfault on stack overflow.
57pub fn with_alloca<R>(size: usize, f: impl FnOnce(&mut [MaybeUninit<u8>]) -> R) -> R {
58    let mut ret = MaybeUninit::uninit();
59
60    let closure = |ptr| {
61        let slice = unsafe { slice::from_raw_parts_mut(ptr, size) };
62        ret.write(f(slice));
63    };
64
65    let trampoline = get_trampoline(&closure);
66    let mut closure_data = ManuallyDrop::new(closure);
67
68    unsafe {
69        c_with_alloca(size, trampoline, &mut closure_data as *mut _ as *mut c_void);
70        ret.assume_init()
71    }
72}
73
74/// Same as `with_alloca` except it zeroes memory slice.
75pub fn with_alloca_zeroed<R>(size: usize, f: impl FnOnce(&mut [u8]) -> R) -> R {
76    with_alloca(size, |memory| unsafe {
77        ptr::write_bytes(memory.as_mut_ptr(), 0, size);
78        f(mem::transmute(memory))
79    })
80}
81
82/// Allocates `T` on stack space.
83pub fn alloca<T, R>(f: impl FnOnce(&mut MaybeUninit<T>) -> R) -> R {
84    use mem::{align_of, size_of};
85
86    with_alloca(size_of::<T>() + (align_of::<T>() - 1), |memory| unsafe {
87        let mut raw_memory = memory.as_mut_ptr();
88        if raw_memory as usize % align_of::<T>() != 0 {
89            raw_memory = raw_memory.add(align_of::<T>() - raw_memory as usize % align_of::<T>());
90        }
91        f(&mut *raw_memory.cast::<MaybeUninit<T>>())
92    })
93}
94
95#[cfg(test)]
96mod tests;
alloca/lib.rs

alloca/
lib.rs
