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 + '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: NonNull<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 let boxed: Box<dyn Future<Output = T> + Send + 'a> = Box::new(future);
24
25 let boxed = NonNull::from(Box::leak(boxed));
26
27 Self { boxed }
28 }
29
30 /// Replace the future currently stored in this box.
31 ///
32 /// This reallocates if and only if the layout of the provided future is
33 /// different from the layout of the currently stored future.
34 pub fn set<F>(&mut self, future: F)
35 where
36 F: Future<Output = T> + Send + 'a,
37 {
38 if let Err(future) = self.try_set(future) {
39 *self = Self::new(future);
40 }
41 }
42
43 /// Replace the future currently stored in this box.
44 ///
45 /// This function never reallocates, but returns an error if the provided
46 /// future has a different size or alignment from the currently stored
47 /// future.
48 pub fn try_set<F>(&mut self, future: F) -> Result<(), F>
49 where
50 F: Future<Output = T> + Send + 'a,
51 {
52 // SAFETY: The pointer is not dangling.
53 let self_layout = {
54 let dyn_future: &(dyn Future<Output = T> + Send) = unsafe { self.boxed.as_ref() };
55 Layout::for_value(dyn_future)
56 };
57
58 if Layout::new::<F>() == self_layout {
59 // SAFETY: We just checked that the layout of F is correct.
60 unsafe {
61 self.set_same_layout(future);
62 }
63
64 Ok(())
65 } else {
66 Err(future)
67 }
68 }
69
70 /// Set the current future.
71 ///
72 /// # Safety
73 ///
74 /// This function requires that the layout of the provided future is the
75 /// same as `self.layout`.
76 unsafe fn set_same_layout<F>(&mut self, future: F)
77 where
78 F: Future<Output = T> + Send + 'a,
79 {
80 // Drop the existing future, catching any panics.
81 let result = panic::catch_unwind(AssertUnwindSafe(|| {
82 ptr::drop_in_place(self.boxed.as_ptr());
83 }));
84
85 // Overwrite the future behind the pointer. This is safe because the
86 // allocation was allocated with the same size and alignment as the type F.
87 let self_ptr: *mut F = self.boxed.as_ptr() as *mut F;
88 ptr::write(self_ptr, future);
89
90 // Update the vtable of self.boxed. The pointer is not null because we
91 // just got it from self.boxed, which is not null.
92 self.boxed = NonNull::new_unchecked(self_ptr);
93
94 // If the old future's destructor panicked, resume unwinding.
95 match result {
96 Ok(()) => {}
97 Err(payload) => {
98 panic::resume_unwind(payload);
99 }
100 }
101 }
102
103 /// Get a pinned reference to the underlying future.
104 pub fn get_pin(&mut self) -> Pin<&mut (dyn Future<Output = T> + Send)> {
105 // SAFETY: The user of this box cannot move the box, and we do not move it
106 // either.
107 unsafe { Pin::new_unchecked(self.boxed.as_mut()) }
108 }
109
110 /// Poll the future stored inside this box.
111 pub fn poll(&mut self, cx: &mut Context<'_>) -> Poll<T> {
112 self.get_pin().poll(cx)
113 }
114}
115
116impl<T> Future for ReusableBoxFuture<'_, T> {
117 type Output = T;
118
119 /// Poll the future stored inside this box.
120 fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T> {
121 Pin::into_inner(self).get_pin().poll(cx)
122 }
123}
124
125// The future stored inside ReusableBoxFuture<'_, T> must be Send.
126unsafe impl<T> Send for ReusableBoxFuture<'_, T> {}
127
128// The only method called on self.boxed is poll, which takes &mut self, so this
129// struct being Sync does not permit any invalid access to the Future, even if
130// the future is not Sync.
131unsafe impl<T> Sync for ReusableBoxFuture<'_, T> {}
132
133// Just like a Pin<Box<dyn Future>> is always Unpin, so is this type.
134impl<T> Unpin for ReusableBoxFuture<'_, T> {}
135
136impl<T> Drop for ReusableBoxFuture<'_, T> {
137 fn drop(&mut self) {
138 unsafe {
139 drop(Box::from_raw(self.boxed.as_ptr()));
140 }
141 }
142}
143
144impl<T> fmt::Debug for ReusableBoxFuture<'_, T> {
145 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
146 f.debug_struct(name:"ReusableBoxFuture").finish()
147 }
148}
149