reusable_box_future/
local_box_future.rs

1use alloc::alloc::Layout;
2use alloc::boxed::Box;
3use core::fmt;
4use core::future::Future;
5use core::mem::ManuallyDrop;
6use core::pin::Pin;
7use core::ptr::{self, NonNull};
8use core::task::{Context, Poll};
9
10/// A reusable `Pin<Box<dyn Future<Output = T>>`.
11///
12/// This type lets you replace the future stored in the box without
13/// reallocating when the size and alignment permits this.
14pub struct ReusableLocalBoxFuture<T> {
15    boxed: NonNull<dyn Future<Output = T>>,
16}
17
18impl<T> ReusableLocalBoxFuture<T> {
19    /// Create a new `ReusableLocalBoxFuture<T>` containing the provided future.
20    pub fn new<F>(future: F) -> Self
21    where
22        F: Future<Output = T> + 'static,
23    {
24        let boxed: Box<dyn Future<Output = T>> = Box::new(future);
25        let boxed = Box::into_raw(boxed);
26
27        // SAFETY: Box::into_raw does not return null pointers.
28        let boxed = unsafe { NonNull::new_unchecked(boxed) };
29
30        Self { boxed }
31    }
32
33    /// Replace the future currently stored in this box.
34    ///
35    /// This reallocates if and only if the layout of the provided future is
36    /// different from the layout of the currently stored future.
37    pub fn set<F>(&mut self, future: F)
38    where
39        F: Future<Output = T> + 'static,
40    {
41        if let Err(future) = self.try_set(future) {
42            *self = Self::new(future);
43        }
44    }
45
46    /// Replace the future currently stored in this box.
47    ///
48    /// This function never reallocates, but returns an error if the provided
49    /// future has a different size or alignment from the currently stored
50    /// future.
51    pub fn try_set<F>(&mut self, future: F) -> Result<(), F>
52    where
53        F: Future<Output = T> + 'static,
54    {
55        // SAFETY: The pointer is not dangling.
56        let self_layout = {
57            let dyn_future: &(dyn Future<Output = T>) = unsafe { self.boxed.as_ref() };
58            Layout::for_value(dyn_future)
59        };
60
61        if Layout::new::<F>() == self_layout {
62            // SAFETY: We just checked that the layout of F is correct.
63            unsafe {
64                self.set_same_layout(future);
65            }
66
67            Ok(())
68        } else {
69            Err(future)
70        }
71    }
72
73    /// Set the current future.
74    ///
75    /// # Safety
76    ///
77    /// This function requires that the layout of the provided future is the
78    /// same as `self.boxed` layout.
79    unsafe fn set_same_layout<F>(&mut self, future: F)
80    where
81        F: Future<Output = T> + 'static,
82    {
83        struct SetLayout<'a, F, T>
84        where
85            F: Future<Output = T> + 'static,
86        {
87            rbf: &'a mut ReusableLocalBoxFuture<T>,
88            new_future: ManuallyDrop<F>,
89        }
90
91        impl<'a, F, T> Drop for SetLayout<'a, F, T>
92        where
93            F: Future<Output = T> + 'static,
94        {
95            fn drop(&mut self) {
96                // By doing the replacement on `drop` we make sure the change
97                // will happen even if the existing future panics on drop.
98                //
99                // We could use `catch_unwind`, but it is not available in `no_std`.
100                unsafe {
101                    // Overwrite the future behind the pointer. This is safe because the
102                    // allocation was allocated with the same size and alignment as the type F.
103                    let fut_ptr: *mut F = self.rbf.boxed.as_ptr() as *mut F;
104                    ptr::write(fut_ptr, ManuallyDrop::take(&mut self.new_future));
105
106                    // Update the vtable of self.boxed. The pointer is not null because we
107                    // just got it from self.boxed, which is not null.
108                    self.rbf.boxed = NonNull::new_unchecked(fut_ptr);
109                }
110            }
111        }
112
113        let set_layout = SetLayout {
114            rbf: self,
115            new_future: ManuallyDrop::new(future),
116        };
117
118        // Drop the existing future.
119        ptr::drop_in_place(set_layout.rbf.boxed.as_ptr());
120        // Now `set_layout` will be dropped and do the replacement.
121    }
122
123    /// Get a pinned reference to the underlying future.
124    pub fn get_pin(&mut self) -> Pin<&mut (dyn Future<Output = T>)> {
125        // SAFETY: The user of this box cannot move the box, and we do not move it
126        // either.
127        unsafe { Pin::new_unchecked(self.boxed.as_mut()) }
128    }
129
130    /// Poll the future stored inside this box.
131    pub fn poll(&mut self, cx: &mut Context<'_>) -> Poll<T> {
132        self.get_pin().poll(cx)
133    }
134}
135
136impl<T> Future for ReusableLocalBoxFuture<T> {
137    type Output = T;
138
139    /// Poll the future stored inside this box.
140    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T> {
141        Pin::into_inner(self).get_pin().poll(cx)
142    }
143}
144
145// Just like a Pin<Box<dyn Future>> is always Unpin, so is this type.
146impl<T> Unpin for ReusableLocalBoxFuture<T> {}
147
148impl<T> Drop for ReusableLocalBoxFuture<T> {
149    fn drop(&mut self) {
150        unsafe {
151            drop(Box::from_raw(self.boxed.as_ptr()));
152        }
153    }
154}
155
156impl<T> fmt::Debug for ReusableLocalBoxFuture<T> {
157    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
158        f.debug_struct("ReusableLocalBoxFuture").finish()
159    }
160}
161
162#[cfg(test)]
163mod tests {
164    use super::*;
165    use static_assertions::{assert_impl_all, assert_not_impl_all};
166
167    #[test]
168    fn static_assertion() {
169        assert_impl_all!(ReusableLocalBoxFuture<()>: Unpin);
170        assert_not_impl_all!(ReusableLocalBoxFuture<()>: Sync, Send);
171    }
172}
reusable_box_future/local_box_future.rs

reusable_box_future/
local_box_future.rs
