gear_stack_buffer/
lib.rs

1// Copyright (C) Gear Technologies Inc.
2// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
3
4//! Stack allocations utils.
5
6#![no_std]
7
8extern crate alloc;
9
10use alloc::vec::Vec;
11use core::{
12    ffi::c_void,
13    mem::{ManuallyDrop, MaybeUninit},
14    slice,
15};
16
17/// The maximum buffer size that can be allocated on the stack.
18/// This is currently limited to 64 KiB.
19pub const MAX_BUFFER_SIZE: usize = 64 * 1024;
20
21/// A closure data type that is used in the native library to pass
22/// a pointer to allocated stack memory.
23type Callback = unsafe extern "C" fn(ptr: *mut MaybeUninit<u8>, data: *mut c_void);
24
25#[cfg(any(feature = "compile-alloca", target_arch = "wasm32"))]
26unsafe extern "C" {
27    /// Function from the native library that manipulates the stack pointer directly.
28    /// Can be used to dynamically allocate stack space.
29    fn c_with_alloca(size: usize, callback: Callback, data: *mut c_void);
30}
31
32/// This is a polyfill function that is used when the native library is unavailable.
33/// The maximum size that can be allocated on the stack is limited
34/// by the [`MAX_BUFFER_SIZE`] constant.
35#[cfg(not(any(feature = "compile-alloca", target_arch = "wasm32")))]
36unsafe extern "C" fn c_with_alloca(_size: usize, callback: Callback, data: *mut c_void) {
37    let mut buffer = [MaybeUninit::uninit(); MAX_BUFFER_SIZE];
38    unsafe { callback(buffer.as_mut_ptr(), data) };
39}
40
41/// Helper function to create a trampoline between C and Rust code.
42#[inline(always)]
43fn get_trampoline<F: FnOnce(*mut MaybeUninit<u8>)>(_closure: &F) -> Callback {
44    trampoline::<F>
45}
46
47/// A function that serves as a trampoline between C and Rust code.
48/// It is mainly used to switch from `fn()` to `FnOnce()`,
49/// which allows local variables to be captured.
50unsafe extern "C" fn trampoline<F: FnOnce(*mut MaybeUninit<u8>)>(
51    ptr: *mut MaybeUninit<u8>,
52    data: *mut c_void,
53) {
54    // This code gets `*mut ManuallyDrop<F>`, then takes ownership of the `F` function
55    // and executes it with a pointer to the allocated stack memory.
56    let f = unsafe { ManuallyDrop::take(&mut *(data as *mut ManuallyDrop<F>)) };
57    f(ptr);
58}
59
60/// This is a higher-level function for dynamically allocating space on the stack.
61fn with_alloca<T>(size: usize, f: impl FnOnce(&mut [MaybeUninit<u8>]) -> T) -> T {
62    let mut ret = MaybeUninit::uninit();
63
64    let closure = |ptr| {
65        let slice = unsafe { slice::from_raw_parts_mut(ptr, size) };
66        ret.write(f(slice));
67    };
68
69    // The `closure` variable is passed as `*mut ManuallyDrop<F>` to the trampoline function.
70    let trampoline = get_trampoline(&closure);
71    let mut closure_data = ManuallyDrop::new(closure);
72
73    unsafe {
74        c_with_alloca(size, trampoline, &mut closure_data as *mut _ as *mut c_void);
75        ret.assume_init()
76    }
77}
78
79/// Calls function `f` with provided uninitialized byte buffer allocated on stack.
80/// ### IMPORTANT
81/// If buffer size is too big (currently bigger than 0x10000 bytes),
82/// then allocation will be on heap.
83/// If buffer is small enough to be allocated on stack, then real allocated
84/// buffer size will be `size` aligned to 16 bytes.
85pub fn with_byte_buffer<T>(size: usize, f: impl FnOnce(&mut [MaybeUninit<u8>]) -> T) -> T {
86    if size <= MAX_BUFFER_SIZE {
87        with_alloca(size, f)
88    } else {
89        f(Vec::with_capacity(size).spare_capacity_mut())
90    }
91}
gear_stack_buffer/lib.rs

gear_stack_buffer/
lib.rs
