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;