1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96

//! Mostly safe no_std wrapper for `alloca` in Rust.
//!
//! This crate uses Rust lifetime system to ensure that stack allocated memory will not be used
//! after function return, but it does not make any guarantee about memory that is turned into
//! raw pointer and stored somewhere else.
//!
//! ### Example
//! ```rust
//! // allocate 128 bytes on the stack
//! alloca::with_alloca(128, |memory| {
//!     // memory: &mut [MaybeUninit<u8>]
//!     assert_eq!(memory.len(), 128);
//! });
//! ```

#![no_std]

use core::{
    ffi::c_void,
    mem::{self, ManuallyDrop, MaybeUninit},
    ptr, slice,
};

type Callback = unsafe extern "C-unwind" fn(ptr: *mut MaybeUninit<u8>, data: *mut c_void);

extern "C-unwind" {
    fn c_with_alloca(size: usize, callback: Callback, data: *mut c_void);
}

#[inline(always)]
fn get_trampoline<F: FnOnce(*mut MaybeUninit<u8>)>(_closure: &F) -> Callback {
    trampoline::<F>
}

unsafe extern "C-unwind" fn trampoline<F: FnOnce(*mut MaybeUninit<u8>)>(
    ptr: *mut MaybeUninit<u8>,
    data: *mut c_void,
) {
    let f = ManuallyDrop::take(&mut *(data as *mut ManuallyDrop<F>));
    f(ptr);
}

/// Allocates `[u8; size]` memory on stack and invokes `closure` with this slice as argument.
///
/// # Safety
/// This function is safe because `c_with_alloca` (which is internally used) will always returns non-null
/// pointer.
///
/// # Potential segfaults or UB
///
/// When using this function in wrong way your program might get UB or segfault "for free":
/// - Using memory allocated by `with_alloca` outside of it e.g closure is already returned but you somehow
/// managed to store pointer to memory and use it.
/// - Allocating more memory than thread stack size.
///
///   This will trigger segfault on stack overflow.
pub fn with_alloca<R>(size: usize, f: impl FnOnce(&mut [MaybeUninit<u8>]) -> R) -> R {
    let mut ret = MaybeUninit::uninit();

    let closure = |ptr| {
        let slice = unsafe { slice::from_raw_parts_mut(ptr, size) };
        ret.write(f(slice));
    };

    let trampoline = get_trampoline(&closure);
    let mut closure_data = ManuallyDrop::new(closure);

    unsafe {
        c_with_alloca(size, trampoline, &mut closure_data as *mut _ as *mut c_void);
        ret.assume_init()
    }
}

/// Same as `with_alloca` except it zeroes memory slice.
pub fn with_alloca_zeroed<R>(size: usize, f: impl FnOnce(&mut [u8]) -> R) -> R {
    with_alloca(size, |memory| unsafe {
        ptr::write_bytes(memory.as_mut_ptr(), 0, size);
        f(mem::transmute(memory))
    })
}

/// Allocates `T` on stack space.
pub fn alloca<T, R>(f: impl FnOnce(&mut MaybeUninit<T>) -> R) -> R {
    use mem::{align_of, size_of};

    with_alloca(size_of::<T>() + (align_of::<T>() - 1), |memory| unsafe {
        let mut raw_memory = memory.as_mut_ptr();
        if raw_memory as usize % align_of::<T>() != 0 {
            raw_memory = raw_memory.add(align_of::<T>() - raw_memory as usize % align_of::<T>());
        }
        f(&mut *raw_memory.cast::<MaybeUninit<T>>())
    })
}

#[cfg(test)]
mod tests;