1#![no_std]
8
9use core::alloc::Layout;
10use core::arch::global_asm;
11use core::mem::{ManuallyDrop, MaybeUninit};
12use core::slice;
13
14use cfg_if::cfg_if;
15
16pub fn with_alloca_bytes<R, F: FnOnce(&mut [MaybeUninit<u8>]) -> R>(layout: Layout, f: F) -> R {
20 with_alloca_ptr(layout, |p| {
21 f(unsafe { slice::from_raw_parts_mut(p.cast(), layout.size()) })
22 })
23}
24
25pub fn with_alloca<R, T, F: FnOnce(&mut MaybeUninit<T>) -> R>(f: F) -> R {
26 with_alloca_ptr(Layout::new::<T>(), |p| f(unsafe { &mut *p.cast() }))
27}
28
29pub fn with_alloca_slice<R, T, F: FnOnce(&mut [MaybeUninit<T>]) -> R>(n: usize, f: F) -> R {
30 with_alloca_ptr(Layout::array::<T>(n).unwrap(), |p| {
31 f(unsafe { slice::from_raw_parts_mut(p.cast(), n) })
32 })
33}
34
35pub fn with_alloca_ptr<R, F: FnOnce(*mut u8) -> R>(layout: Layout, f: F) -> R {
36 unsafe extern "C" fn cont_fn<F: FnOnce(*mut u8)>(
37 reservation_start: *mut u8,
38 cont_arg: *mut ReserveOnStackContArg,
39 ) {
40 let f = ManuallyDrop::take(&mut *(cont_arg as *mut ManuallyDrop<F>));
41 f(reservation_start)
42 }
43
44 #[inline(always)]
45 fn get_cont_fn<F: FnOnce(*mut u8)>(_closure: &F) -> ReserveOnStackContFn {
46 cont_fn::<F>
47 }
48
49 let mut ret = MaybeUninit::uninit();
50
51 let closure = |p| {
52 ret.write(f(p));
53 };
54
55 let inst_cont_fn = get_cont_fn(&closure);
56 let mut closure_data = ManuallyDrop::new(closure);
57
58 unsafe {
59 reserve_on_stack(
60 layout,
61 inst_cont_fn,
62 &mut closure_data as *mut _ as *mut ReserveOnStackContArg,
63 );
64 ret.assume_init()
65 }
66}
67
68type ReserveOnStackContFn =
69 unsafe extern "C" fn(reservation_start: *mut u8, cont_arg: *mut ReserveOnStackContArg);
70
71enum ReserveOnStackContArg {}
72
73#[allow(clippy::missing_safety_doc)]
74unsafe fn reserve_on_stack(
75 layout: Layout,
76 cont_fn: ReserveOnStackContFn,
77 cont_arg: *mut ReserveOnStackContArg,
78) {
79 let reservation_size = layout.size();
80 let reservation_align_down_mask = !(layout.align() - 1);
81 __sel4_alloca__reserve_on_stack(
82 reservation_size,
83 reservation_align_down_mask,
84 cont_fn,
85 cont_arg,
86 )
87}
88
89extern "C" {
90 fn __sel4_alloca__reserve_on_stack(
91 reservation_size: usize,
92 reservation_align_down_mask: usize,
93 cont_fn: ReserveOnStackContFn,
94 cont_arg: *mut ReserveOnStackContArg,
95 );
96}
97
98macro_rules! common_asm_prefix {
99 () => {
100 r#"
101 .section .text
102 .global __sel4_alloca__reserve_on_stack
103 __sel4_alloca__reserve_on_stack:
104 "#
105 };
106}
107
108cfg_if! {
109 if #[cfg(target_arch = "aarch64")] {
110 global_asm! {
111 common_asm_prefix!(),
112 r#"
113 // preamble
114 stp fp, lr, [sp, #-16]!
115 mov fp, sp
116
117 mov x9, sp
118 sub x9, x9, x0 // x0: reservation_size
119 and x9, x9, x1 // x1: reservation_align_down_mask
120 mov x10, x9 // save reservation_start for later
121 and x9, x9, ~(16 - 1) // align stack
122 mov sp, x9
123
124 mov x0, x10 // pass reservation_start
125 mov x1, x3 // pass cont_arg
126 blr x2 // call cont_fn
127
128 // postamble
129 mov sp, fp
130 ldp fp, lr, [sp], #16
131 ret
132 "#
133 }
134 } else if #[cfg(target_arch = "arm")] {
135 global_asm! {
136 common_asm_prefix!(),
137 r#"
138 // preamble
139 push {{r11, lr}}
140 add r11, sp, #0
141
142 mov r4, sp
143 sub r4, r4, r0 // r0: reservation_size
144 and r4, r4, r1 // r1: reservation_align_down_mask
145 mov r6, r4 // save reservation_start for later
146 and r4, r4, ~(4 - 1) // align stack
147 mov sp, r4
148
149 mov r0, r6 // pass reservation_start
150 mov r1, r3 // pass cont_arg
151 blx r2 // call cont_fn
152
153 // postamble
154 sub sp, r11, #0
155 pop {{r11, pc}}
156 "#
157 }
158 } else if #[cfg(target_arch = "riscv64")] {
159 global_asm! {
160 common_asm_prefix!(),
161 r#"
162 // preamble
163 sd ra, -8(sp)
164 sd s0, -16(sp)
165 mv s0, sp
166 addi sp, sp, -16
167
168 mv t0, sp
169 sub t0, t0, a0 // a0: reservation_size
170 and t0, t0, a1 // a1: reservation_align_down_mask
171 mv t1, t0 // save reservation_start for later
172 and t0, t0, ~(16 - 1) // align stack
173 mv sp, t0
174
175 mv a0, t1 // pass reservation_start
176 mv a1, a3 // pass cont_arg
177 jalr a2 // call cont_fn
178
179 // postamble
180 mv sp, s0
181 ld ra, -8(sp)
182 ld s0, -16(sp)
183 ret
184 "#
185 }
186 } else if #[cfg(target_arch = "riscv32")] {
187 global_asm! {
188 common_asm_prefix!(),
189 r#"
190 // preamble
191 sw ra, -4(sp)
192 sw s0, -8(sp)
193 mv s0, sp
194 addi sp, sp, -16
195
196 mv t0, sp
197 sub t0, t0, a0 // a0: reservation_size
198 and t0, t0, a1 // a1: reservation_align_down_mask
199 mv t1, t0 // save reservation_start for later
200 and t0, t0, ~(16 - 1) // align stack
201 mv sp, t0
202
203 mv a0, t1 // pass reservation_start
204 mv a1, a3 // pass cont_arg
205 jalr a2 // call cont_fn
206
207 // postamble
208 mv sp, s0
209 lw ra, -4(sp)
210 lw s0, -8(sp)
211 ret
212 "#
213 }
214 } else if #[cfg(target_arch = "x86_64")] {
215 global_asm! {
216 common_asm_prefix!(),
217 r#"
218 // preamble
219 push rbp
220 mov rbp, rsp
221
222 mov r10, rsp
223 sub r10, rdi // rdi: reservation_size
224 and r10, rsi // rsi: reservation_align_down_mask
225 mov rax, r10 // save reservation_start for later
226 and r10, ~(16 - 1) // align stack
227 mov rsp, r10
228
229 mov rdi, rax // pass reservation_start
230 mov rsi, rcx // pass cont_arg
231 call rdx // call cont_fn
232
233 // postamble
234 leave
235 ret
236 "#
237 }
238 } else {
239 compile_error!("unsupported architecture");
240 }
241}