1use std::alloc::Layout;
2use std::fmt;
3use std::future::{self, Future};
4use std::mem::{self, ManuallyDrop};
5use std::pin::Pin;
6use std::ptr;
7use std::task::{Context, Poll};
8
9/// A reusable `Pin<Box<dyn Future<Output = T> + Send + 'a>>`.
10///
11/// This type lets you replace the future stored in the box without
12/// reallocating when the size and alignment permits this.
13pub struct ReusableBoxFuture<'a, T> {
14 boxed: Pin<Box<dyn Future<Output = T> + Send + 'a>>,
15}
16
17impl<'a, T> ReusableBoxFuture<'a, T> {
18 /// Create a new `ReusableBoxFuture<T>` containing the provided future.
19 pub fn new<F>(future: F) -> Self
20 where
21 F: Future<Output = T> + Send + 'a,
22 {
23 Self {
24 boxed: Box::pin(future),
25 }
26 }
27
28 /// Replace the future currently stored in this box.
29 ///
30 /// This reallocates if and only if the layout of the provided future is
31 /// different from the layout of the currently stored future.
32 pub fn set<F>(&mut self, future: F)
33 where
34 F: Future<Output = T> + Send + 'a,
35 {
36 if let Err(future) = self.try_set(future) {
37 *self = Self::new(future);
38 }
39 }
40
41 /// Replace the future currently stored in this box.
42 ///
43 /// This function never reallocates, but returns an error if the provided
44 /// future has a different size or alignment from the currently stored
45 /// future.
46 pub fn try_set<F>(&mut self, future: F) -> Result<(), F>
47 where
48 F: Future<Output = T> + Send + 'a,
49 {
50 // If we try to inline the contents of this function, the type checker complains because
51 // the bound `T: 'a` is not satisfied in the call to `pending()`. But by putting it in an
52 // inner function that doesn't have `T` as a generic parameter, we implicitly get the bound
53 // `F::Output: 'a` transitively through `F: 'a`, allowing us to call `pending()`.
54 #[inline(always)]
55 fn real_try_set<'a, F>(
56 this: &mut ReusableBoxFuture<'a, F::Output>,
57 future: F,
58 ) -> Result<(), F>
59 where
60 F: Future + Send + 'a,
61 {
62 // future::Pending<T> is a ZST so this never allocates.
63 let boxed = mem::replace(&mut this.boxed, Box::pin(future::pending()));
64 reuse_pin_box(boxed, future, |boxed| this.boxed = Pin::from(boxed))
65 }
66
67 real_try_set(self, future)
68 }
69
70 /// Get a pinned reference to the underlying future.
71 pub fn get_pin(&mut self) -> Pin<&mut (dyn Future<Output = T> + Send)> {
72 self.boxed.as_mut()
73 }
74
75 /// Poll the future stored inside this box.
76 pub fn poll(&mut self, cx: &mut Context<'_>) -> Poll<T> {
77 self.get_pin().poll(cx)
78 }
79}
80
81impl<T> Future for ReusableBoxFuture<'_, T> {
82 type Output = T;
83
84 /// Poll the future stored inside this box.
85 fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T> {
86 Pin::into_inner(self).get_pin().poll(cx)
87 }
88}
89
90// The only method called on self.boxed is poll, which takes &mut self, so this
91// struct being Sync does not permit any invalid access to the Future, even if
92// the future is not Sync.
93unsafe impl<T> Sync for ReusableBoxFuture<'_, T> {}
94
95impl<T> fmt::Debug for ReusableBoxFuture<'_, T> {
96 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
97 f.debug_struct(name:"ReusableBoxFuture").finish()
98 }
99}
100
101fn reuse_pin_box<T: ?Sized, U, O, F>(boxed: Pin<Box<T>>, new_value: U, callback: F) -> Result<O, U>
102where
103 F: FnOnce(Box<U>) -> O,
104{
105 let layout = Layout::for_value::<T>(&*boxed);
106 if layout != Layout::new::<U>() {
107 return Err(new_value);
108 }
109
110 // SAFETY: We don't ever construct a non-pinned reference to the old `T` from now on, and we
111 // always drop the `T`.
112 let raw: *mut T = Box::into_raw(unsafe { Pin::into_inner_unchecked(boxed) });
113
114 // When dropping the old value panics, we still want to call `callback` — so move the rest of
115 // the code into a guard type.
116 let guard = CallOnDrop::new(|| {
117 let raw: *mut U = raw.cast::<U>();
118 unsafe { raw.write(new_value) };
119
120 // SAFETY:
121 // - `T` and `U` have the same layout.
122 // - `raw` comes from a `Box` that uses the same allocator as this one.
123 // - `raw` points to a valid instance of `U` (we just wrote it in).
124 let boxed = unsafe { Box::from_raw(raw) };
125
126 callback(boxed)
127 });
128
129 // Drop the old value.
130 unsafe { ptr::drop_in_place(raw) };
131
132 // Run the rest of the code.
133 Ok(guard.call())
134}
135
136struct CallOnDrop<O, F: FnOnce() -> O> {
137 f: ManuallyDrop<F>,
138}
139
140impl<O, F: FnOnce() -> O> CallOnDrop<O, F> {
141 fn new(f: F) -> Self {
142 let f: ManuallyDrop = ManuallyDrop::new(f);
143 Self { f }
144 }
145 fn call(self) -> O {
146 let mut this: ManuallyDrop> = ManuallyDrop::new(self);
147 let f: F = unsafe { ManuallyDrop::take(&mut this.f) };
148 f()
149 }
150}
151
152impl<O, F: FnOnce() -> O> Drop for CallOnDrop<O, F> {
153 fn drop(&mut self) {
154 let f: F = unsafe { ManuallyDrop::take(&mut self.f) };
155 f();
156 }
157}
158