gear_stack_buffer/
lib.rs

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

gear_stack_buffer/
lib.rs
