1use std::alloc::Layout;
2use std::future::Future;
3use std::panic::AssertUnwindSafe;
4use std::pin::Pin;
5use std::ptr::{self, NonNull};
6use std::task::{Context, Poll};
7use std::{fmt, panic};
8
9/// A reusable `Pin<Box<dyn Future<Output = T> + Send>>`.
10///
11/// This type lets you replace the future stored in the box without
12/// reallocating when the size and alignment permits this.
13pub(crate) struct ReusableBoxFuture<T> {
14 boxed: NonNull<dyn Future<Output = T> + Send>,
15}
16
17impl<T> ReusableBoxFuture<T> {
18 /// Create a new `ReusableBoxFuture<T>` containing the provided future.
19 pub(crate) fn new<F>(future: F) -> Self
20 where
21 F: Future<Output = T> + Send + 'static,
22 {
23 let boxed: Box<dyn Future<Output = T> + Send> = Box::new(future);
24
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 /// Replaces 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(crate) fn set<F>(&mut self, future: F)
38 where
39 F: Future<Output = T> + Send + 'static,
40 {
41 if let Err(future) = self.try_set(future) {
42 *self = Self::new(future);
43 }
44 }
45
46 /// Replaces 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(crate) fn try_set<F>(&mut self, future: F) -> Result<(), F>
52 where
53 F: Future<Output = T> + Send + 'static,
54 {
55 // SAFETY: The pointer is not dangling.
56 let self_layout = {
57 let dyn_future: &(dyn Future<Output = T> + Send) = 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 /// Sets the current future.
74 ///
75 /// # Safety
76 ///
77 /// This function requires that the layout of the provided future is the
78 /// same as `self.layout`.
79 unsafe fn set_same_layout<F>(&mut self, future: F)
80 where
81 F: Future<Output = T> + Send + 'static,
82 {
83 // Drop the existing future, catching any panics.
84 let result = panic::catch_unwind(AssertUnwindSafe(|| {
85 ptr::drop_in_place(self.boxed.as_ptr());
86 }));
87
88 // Overwrite the future behind the pointer. This is safe because the
89 // allocation was allocated with the same size and alignment as the type F.
90 let self_ptr: *mut F = self.boxed.as_ptr() as *mut F;
91 ptr::write(self_ptr, future);
92
93 // Update the vtable of self.boxed. The pointer is not null because we
94 // just got it from self.boxed, which is not null.
95 self.boxed = NonNull::new_unchecked(self_ptr);
96
97 // If the old future's destructor panicked, resume unwinding.
98 match result {
99 Ok(()) => {}
100 Err(payload) => {
101 panic::resume_unwind(payload);
102 }
103 }
104 }
105
106 /// Gets a pinned reference to the underlying future.
107 pub(crate) fn get_pin(&mut self) -> Pin<&mut (dyn Future<Output = T> + Send)> {
108 // SAFETY: The user of this box cannot move the box, and we do not move it
109 // either.
110 unsafe { Pin::new_unchecked(self.boxed.as_mut()) }
111 }
112
113 /// Polls the future stored inside this box.
114 pub(crate) fn poll(&mut self, cx: &mut Context<'_>) -> Poll<T> {
115 self.get_pin().poll(cx)
116 }
117}
118
119impl<T> Future for ReusableBoxFuture<T> {
120 type Output = T;
121
122 /// Polls the future stored inside this box.
123 fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T> {
124 Pin::into_inner(self).get_pin().poll(cx)
125 }
126}
127
128// The future stored inside ReusableBoxFuture<T> must be Send.
129unsafe impl<T> Send for ReusableBoxFuture<T> {}
130
131// The only method called on self.boxed is poll, which takes &mut self, so this
132// struct being Sync does not permit any invalid access to the Future, even if
133// the future is not Sync.
134unsafe impl<T> Sync for ReusableBoxFuture<T> {}
135
136// Just like a Pin<Box<dyn Future>> is always Unpin, so is this type.
137impl<T> Unpin for ReusableBoxFuture<T> {}
138
139impl<T> Drop for ReusableBoxFuture<T> {
140 fn drop(&mut self) {
141 unsafe {
142 drop(Box::from_raw(self.boxed.as_ptr()));
143 }
144 }
145}
146
147impl<T> fmt::Debug for ReusableBoxFuture<T> {
148 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
149 f.debug_struct(name:"ReusableBoxFuture").finish()
150 }
151}
152
153#[cfg(test)]
154mod test {
155 use super::ReusableBoxFuture;
156 use futures::future::FutureExt;
157 use std::alloc::Layout;
158 use std::future::Future;
159 use std::pin::Pin;
160 use std::task::{Context, Poll};
161
162 #[test]
163 fn test_different_futures() {
164 let fut = async move { 10 };
165 // Not zero sized!
166 assert_eq!(Layout::for_value(&fut).size(), 1);
167
168 let mut b = ReusableBoxFuture::new(fut);
169
170 assert_eq!(b.get_pin().now_or_never(), Some(10));
171
172 b.try_set(async move { 20 })
173 .unwrap_or_else(|_| panic!("incorrect size"));
174
175 assert_eq!(b.get_pin().now_or_never(), Some(20));
176
177 b.try_set(async move { 30 })
178 .unwrap_or_else(|_| panic!("incorrect size"));
179
180 assert_eq!(b.get_pin().now_or_never(), Some(30));
181 }
182
183 #[test]
184 fn test_different_sizes() {
185 let fut1 = async move { 10 };
186 let val = [0u32; 1000];
187 let fut2 = async move { val[0] };
188 let fut3 = ZeroSizedFuture {};
189
190 assert_eq!(Layout::for_value(&fut1).size(), 1);
191 assert_eq!(Layout::for_value(&fut2).size(), 4004);
192 assert_eq!(Layout::for_value(&fut3).size(), 0);
193
194 let mut b = ReusableBoxFuture::new(fut1);
195 assert_eq!(b.get_pin().now_or_never(), Some(10));
196 b.set(fut2);
197 assert_eq!(b.get_pin().now_or_never(), Some(0));
198 b.set(fut3);
199 assert_eq!(b.get_pin().now_or_never(), Some(5));
200 }
201
202 struct ZeroSizedFuture {}
203 impl Future for ZeroSizedFuture {
204 type Output = u32;
205 fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<u32> {
206 Poll::Ready(5)
207 }
208 }
209
210 #[test]
211 fn test_zero_sized() {
212 let fut = ZeroSizedFuture {};
213 // Zero sized!
214 assert_eq!(Layout::for_value(&fut).size(), 0);
215
216 let mut b = ReusableBoxFuture::new(fut);
217
218 assert_eq!(b.get_pin().now_or_never(), Some(5));
219 assert_eq!(b.get_pin().now_or_never(), Some(5));
220
221 b.try_set(ZeroSizedFuture {})
222 .unwrap_or_else(|_| panic!("incorrect size"));
223
224 assert_eq!(b.get_pin().now_or_never(), Some(5));
225 assert_eq!(b.get_pin().now_or_never(), Some(5));
226 }
227}
228