embedded_jsonrpc/stackfuture/mod.rs
1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT License.
3
4//! This crate defines a `StackFuture` wrapper around futures that stores the wrapped
5//! future in space provided by the caller. This can be used to emulate dyn async traits
6//! without requiring heap allocation.
7//!
8//! For more details, see the documentation on the [`StackFuture`] struct.
9
10// std is needed to run tests, but otherwise we don't need it.
11
12use core::fmt::Debug;
13use core::fmt::Display;
14use core::future::Future;
15use core::marker::PhantomData;
16use core::mem;
17use core::mem::MaybeUninit;
18use core::pin::Pin;
19use core::ptr;
20use core::task::Context;
21use core::task::Poll;
22
23/// A wrapper that stores a future in space allocated by the container
24///
25/// Often this space comes from the calling function's stack, but it could just
26/// as well come from some other allocation.
27///
28/// A `StackFuture` can be used to emulate async functions in dyn Trait objects.
29/// For example:
30///
31/// ```
32/// #
33/// use embedded_jsonrpc::stackfuture::StackFuture;
34///
35/// trait PseudoAsyncTrait {
36/// fn do_something(&self) -> StackFuture<'_, (), { 512 }>;
37/// }
38///
39/// impl PseudoAsyncTrait for i32 {
40/// fn do_something(&self) -> StackFuture<'_, (), { 512 }> {
41/// StackFuture::from(async {
42/// // function body goes here
43/// })
44/// }
45/// }
46///
47/// async fn use_dyn_async_trait(x: &dyn PseudoAsyncTrait) {
48/// x.do_something().await;
49/// }
50///
51/// async fn call_with_dyn_async_trait() {
52/// use_dyn_async_trait(&42).await;
53/// }
54/// ```
55///
56/// This example defines `PseudoAsyncTrait` with a single method `do_something`.
57/// The `do_something` method can be called as if it were declared as
58/// `async fn do_something(&self)`. To implement `do_something`, the easiest thing
59/// to do is to wrap the body of the function in `StackFuture::from(async { ... })`,
60/// which creates an anonymous future for the body and stores it in a `StackFuture`.
61///
62/// Because `StackFuture` does not know the size of the future it wraps, the maximum
63/// size of the future must be specified in the `STACK_SIZE` parameter. In the example
64/// here, we've used a stack size of 512, which is probably much larger than necessary
65/// but would accommodate many futures besides the simple one we've shown here.
66///
67/// `StackFuture` ensures when wrapping a future that enough space is available, and
68/// it also respects any alignment requirements for the wrapped future. Note that the
69/// wrapped future's alignment must be less than or equal to that of the overall
70/// `StackFuture` struct.
71#[repr(C, align(8))] // Ensures the data first does not have any padding before it in the struct
72pub struct StackFuture<'a, T, const STACK_SIZE: usize> {
73 /// An array of bytes that is used to store the wrapped future.
74 data: [MaybeUninit<u8>; STACK_SIZE],
75 /// Since the type of `StackFuture` does not know the underlying future that it is wrapping,
76 /// we keep a manual vtable that serves pointers to Poll::poll and Drop::drop. These are
77 /// generated and filled in by `StackFuture::from`.
78 ///
79 /// This field stores a pointer to the poll function wrapper.
80 poll_fn: fn(this: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T>,
81 /// Stores a pointer to the drop function wrapper
82 ///
83 /// See the documentation on `poll_fn` for more details.
84 drop_fn: fn(this: &mut Self),
85 /// StackFuture can be used similarly to a `dyn Future`. We keep a PhantomData
86 /// here so the type system knows this.
87 _phantom: PhantomData<dyn Future<Output = T> + Send + 'a>,
88}
89
90impl<'a, T, const STACK_SIZE: usize> StackFuture<'a, T, { STACK_SIZE }> {
91 /// Creates a `StackFuture` from an existing future
92 ///
93 /// See the documentation on [`StackFuture`] for examples of how to use this.
94 ///
95 /// The size and alignment requirements are statically checked, so it is a compiler error
96 /// to use this with a future that does not fit within the [`StackFuture`]'s size and
97 /// alignment requirements.
98 ///
99 /// The following example illustrates a compile error for a future that is too large.
100 /// ```compile_fail
101 /// # use stackfuture::StackFuture;
102 /// // Fails because the future contains a large array and is therefore too big to fit in
103 /// // a 16-byte `StackFuture`.
104 /// let f = StackFuture::<_, { 16 }>::from(async {
105 /// let x = [0u8; 4096];
106 /// async {}.await;
107 /// println!("{}", x.len());
108 /// });
109 /// # #[cfg(miri)] break rust; // FIXME: miri doesn't detect this breakage for some reason...
110 /// ```
111 ///
112 /// The example below illustrates a compiler error for a future whose alignment is too large.
113 /// ```compile_fail
114 /// # use stackfuture::StackFuture;
115 ///
116 /// #[derive(Debug)]
117 /// #[repr(align(256))]
118 /// struct BigAlignment(usize);
119 ///
120 /// // Fails because the future contains a large array and is therefore too big to fit in
121 /// // a 16-byte `StackFuture`.
122 /// let f = StackFuture::<_, { 16 }>::from(async {
123 /// let x = BigAlignment(42);
124 /// async {}.await;
125 /// println!("{x:?}");
126 /// });
127 /// # #[cfg(miri)] break rust; // FIXME: miri doesn't detect this breakage for some reason...
128 /// ```
129 pub fn from<F>(future: F) -> Self
130 where
131 F: Future<Output = T> + Send + 'a, // the bounds here should match those in the _phantom field
132 {
133 // Ideally we would provide this as:
134 //
135 // impl<'a, F, const STACK_SIZE: usize> From<F> for StackFuture<'a, F::Output, { STACK_SIZE }>
136 // where
137 // F: Future + Send + 'a
138 //
139 // However, libcore provides a blanket `impl<T> From<T> for T`, and since `StackFuture: Future`,
140 // both impls end up being applicable to do `From<StackFuture> for StackFuture`.
141
142 // Statically assert that `F` meets all the size and alignment requirements
143 #[allow(clippy::let_unit_value)]
144 let _ = AssertFits::<F, STACK_SIZE>::ASSERT;
145
146 Self::try_from(future).unwrap()
147 }
148
149 /// Attempts to create a `StackFuture` from an existing future
150 ///
151 /// If the `StackFuture` is not large enough to hold `future`, this function returns an
152 /// `Err` with the argument `future` returned to you.
153 ///
154 /// Panics
155 ///
156 /// If we cannot satisfy the alignment requirements for `F`, this function will panic.
157 pub fn try_from<F>(future: F) -> Result<Self, IntoStackFutureError<F>>
158 where
159 F: Future<Output = T> + Send + 'a, // the bounds here should match those in the _phantom field
160 {
161 if Self::has_space_for_val(&future) && Self::has_alignment_for_val(&future) {
162 let mut result = StackFuture {
163 data: [MaybeUninit::uninit(); STACK_SIZE],
164 poll_fn: Self::poll_inner::<F>,
165 drop_fn: Self::drop_inner::<F>,
166 _phantom: PhantomData,
167 };
168
169 // Ensure result.data is at the beginning of the struct so we don't need to do
170 // alignment adjustments.
171 assert_eq!(result.data.as_ptr() as usize, &result as *const _ as usize);
172
173 // SAFETY: result.as_mut_ptr returns a pointer into result.data, which is an
174 // uninitialized array of bytes. result.as_mut_ptr ensures the returned pointer
175 // is correctly aligned, and the if expression we are in ensures the buffer is
176 // large enough.
177 //
178 // Because `future` is bound by `'a` and `StackFuture` is also bound by `'a`,
179 // we can be sure anything that `future` closes over will also outlive `result`.
180 unsafe { result.as_mut_ptr::<F>().write(future) };
181
182 Ok(result)
183 } else {
184 Err(IntoStackFutureError::new::<Self>(future))
185 }
186 }
187
188 /// A wrapper around the inner future's poll function, which we store in the poll_fn field
189 /// of this struct.
190 fn poll_inner<F: Future>(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<F::Output> {
191 self.as_pin_mut_ref::<F>().poll(cx)
192 }
193
194 /// A wrapper around the inner future's drop function, which we store in the drop_fn field
195 /// of this struct.
196 fn drop_inner<F>(&mut self) {
197 // SAFETY: *this.as_mut_ptr() was previously written as type F
198 unsafe { ptr::drop_in_place(self.as_mut_ptr::<F>()) }
199 }
200
201 /// Returns a pointer into self.data that meets the alignment requirements for type `F`
202 ///
203 /// Before writing to the returned pointer, the caller must ensure that self.data is large
204 /// enough to hold F and any required padding.
205 fn as_mut_ptr<F>(&mut self) -> *mut F {
206 assert!(Self::has_space_for::<F>());
207 // SAFETY: Self is laid out so that the space for the future comes at offset 0.
208 // This is checked by an assertion in Self::from. Thus it's safe to cast a pointer
209 // to Self into a pointer to the wrapped future.
210 unsafe { mem::transmute(self) }
211 }
212
213 /// Returns a pinned mutable reference to a type F stored in self.data
214 fn as_pin_mut_ref<F>(self: Pin<&mut Self>) -> Pin<&mut F> {
215 // SAFETY: `StackFuture` is only created by `StackFuture::from`, which
216 // writes an `F` to `self.as_mut_ptr(), so it's okay to cast the `*mut F`
217 // to an `&mut F` with the same lifetime as `self`.
218 //
219 // For pinning, since self is already pinned, we know the wrapped future
220 // is also pinned.
221 //
222 // This function is only doing pointer arithmetic and casts, so we aren't moving
223 // any pinned data.
224 unsafe { self.map_unchecked_mut(|this| &mut *this.as_mut_ptr()) }
225 }
226
227 /// Computes how much space is required to store a value of type `F`
228 const fn required_space<F>() -> usize {
229 mem::size_of::<F>()
230 }
231
232 /// Determines whether this `StackFuture` can hold a value of type `F`
233 pub const fn has_space_for<F>() -> bool {
234 Self::required_space::<F>() <= STACK_SIZE
235 }
236
237 /// Determines whether this `StackFuture` can hold the referenced value
238 pub const fn has_space_for_val<F>(_: &F) -> bool {
239 Self::has_space_for::<F>()
240 }
241
242 /// Determines whether this `StackFuture`'s alignment is compatible with the
243 /// type `F`.
244 pub const fn has_alignment_for<F>() -> bool {
245 mem::align_of::<F>() <= mem::align_of::<Self>()
246 }
247
248 /// Determines whether this `StackFuture`'s alignment is compatible with the
249 /// referenced value.
250 pub const fn has_alignment_for_val<F>(_: &F) -> bool {
251 Self::has_alignment_for::<F>()
252 }
253}
254
255impl<T, const STACK_SIZE: usize> Future for StackFuture<'_, T, { STACK_SIZE }> {
256 type Output = T;
257
258 fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
259 // SAFETY: This is doing pin projection. We unpin self so we can
260 // access self.poll_fn, and then re-pin self to pass it into poll_in.
261 // The part of the struct that needs to be pinned is data, since it
262 // contains a potentially self-referential future object, but since we
263 // do not touch that while self is unpinned and we do not move self
264 // while unpinned we are okay.
265 unsafe {
266 let this = self.get_unchecked_mut();
267 (this.poll_fn)(Pin::new_unchecked(this), cx)
268 }
269 }
270}
271
272impl<T, const STACK_SIZE: usize> Drop for StackFuture<'_, T, { STACK_SIZE }> {
273 fn drop(&mut self) {
274 (self.drop_fn)(self);
275 }
276}
277
278struct AssertFits<F, const STACK_SIZE: usize>(PhantomData<F>);
279
280impl<F, const STACK_SIZE: usize> AssertFits<F, STACK_SIZE> {
281 const ASSERT: () = {
282 if !StackFuture::<F, STACK_SIZE>::has_space_for::<F>() {
283 panic!("F is too large");
284 }
285
286 if !StackFuture::<F, STACK_SIZE>::has_alignment_for::<F>() {
287 panic!("F has incompatible alignment");
288 }
289 };
290}
291
292/// Captures information about why a future could not be converted into a [`StackFuture`]
293///
294/// It also contains the original future so that callers can still run the future in error
295/// recovery paths, such as by boxing the future instead of wrapping it in [`StackFuture`].
296pub struct IntoStackFutureError<F> {
297 /// The size of the StackFuture we tried to convert the future into
298 maximum_size: usize,
299 /// The StackFuture's alignment
300 maximum_alignment: usize,
301 /// The future that was attempted to be wrapped
302 future: F,
303}
304
305impl<F> IntoStackFutureError<F> {
306 fn new<Target>(future: F) -> Self {
307 Self {
308 maximum_size: mem::size_of::<Target>(),
309 maximum_alignment: mem::align_of::<Target>(),
310 future,
311 }
312 }
313
314 /// Returns true if the target [`StackFuture`] was too small to hold the given future.
315 pub fn insufficient_space(&self) -> bool {
316 self.maximum_size < mem::size_of_val(&self.future)
317 }
318
319 /// Returns true if the target [`StackFuture`]'s alignment was too small to accommodate the given future.
320 pub fn alignment_too_small(&self) -> bool {
321 self.maximum_alignment < mem::align_of_val(&self.future)
322 }
323
324 /// Returns the alignment of the wrapped future.
325 pub fn required_alignment(&self) -> usize {
326 mem::align_of_val(&self.future)
327 }
328
329 /// Returns the size of the wrapped future.
330 pub fn required_space(&self) -> usize {
331 mem::size_of_val(&self.future)
332 }
333
334 /// Returns the alignment of the target [`StackFuture`], which is also the maximum alignment
335 /// that can be wrapped.
336 pub const fn available_alignment(&self) -> usize {
337 self.maximum_alignment
338 }
339
340 /// Returns the amount of space that was available in the target [`StackFuture`].
341 pub const fn available_space(&self) -> usize {
342 self.maximum_size
343 }
344
345 /// Returns the underlying future that caused this error
346 ///
347 /// Can be used to try again, either by directly awaiting the future, wrapping it in a `Box`,
348 /// or some other method.
349 #[allow(dead_code)]
350 fn into_inner(self) -> F {
351 self.future
352 }
353}
354
355impl<F> Display for IntoStackFutureError<F> {
356 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
357 match (self.alignment_too_small(), self.insufficient_space()) {
358 (true, true) => write!(f,
359 "cannot create StackFuture, required size is {}, available space is {}; required alignment is {} but maximum alignment is {}",
360 self.required_space(),
361 self.available_space(),
362 self.required_alignment(),
363 self.available_alignment()
364 ),
365 (true, false) => write!(f,
366 "cannot create StackFuture, required alignment is {} but maximum alignment is {}",
367 self.required_alignment(),
368 self.available_alignment()
369 ),
370 (false, true) => write!(f,
371 "cannot create StackFuture, required size is {}, available space is {}",
372 self.required_space(),
373 self.available_space()
374 ),
375 // If we have space and alignment, then `try_from` would have succeeded
376 (false, false) => unreachable!(),
377 }
378 }
379}
380
381impl<F> Debug for IntoStackFutureError<F> {
382 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
383 f.debug_struct("IntoStackFutureError")
384 .field("maximum_size", &self.maximum_size)
385 .field("maximum_alignment", &self.maximum_alignment)
386 .field("future", &core::any::type_name::<F>())
387 .finish()
388 }
389}
390
391#[cfg(test)]
392mod tests;