Jason Macnak | ea22e81 | 2020-03-18 01:29:19 +0000 | [diff] [blame] | 1 | use core::{ |
Jason Macnak | ea22e81 | 2020-03-18 01:29:19 +0000 | [diff] [blame] | 2 | fmt, |
| 3 | future::Future, |
| 4 | marker::PhantomData, |
Joel Galenson | 9ceddf6 | 2021-05-19 15:31:02 -0700 | [diff] [blame] | 5 | mem, |
Jason Macnak | ea22e81 | 2020-03-18 01:29:19 +0000 | [diff] [blame] | 6 | pin::Pin, |
| 7 | task::{Context, Poll}, |
| 8 | }; |
| 9 | |
| 10 | /// A custom trait object for polling futures, roughly akin to |
| 11 | /// `Box<dyn Future<Output = T> + 'a>`. |
| 12 | /// |
| 13 | /// This custom trait object was introduced as currently it is not possible to |
| 14 | /// take `dyn Trait` by value and `Box<dyn Trait>` is not available in no_std |
| 15 | /// contexts. |
| 16 | pub struct LocalFutureObj<'a, T> { |
| 17 | future: *mut (dyn Future<Output = T> + 'static), |
| 18 | drop_fn: unsafe fn(*mut (dyn Future<Output = T> + 'static)), |
| 19 | _marker: PhantomData<&'a ()>, |
| 20 | } |
| 21 | |
| 22 | // As LocalFutureObj only holds pointers, even if we move it, the pointed to values won't move, |
| 23 | // so this is safe as long as we don't provide any way for a user to directly access the pointers |
| 24 | // and move their values. |
| 25 | impl<T> Unpin for LocalFutureObj<'_, T> {} |
| 26 | |
| 27 | #[allow(single_use_lifetimes)] |
| 28 | #[allow(clippy::transmute_ptr_to_ptr)] |
Joel Galenson | 9ceddf6 | 2021-05-19 15:31:02 -0700 | [diff] [blame] | 29 | unsafe fn remove_future_lifetime<'a, T>( |
| 30 | ptr: *mut (dyn Future<Output = T> + 'a), |
| 31 | ) -> *mut (dyn Future<Output = T> + 'static) { |
Jason Macnak | ea22e81 | 2020-03-18 01:29:19 +0000 | [diff] [blame] | 32 | mem::transmute(ptr) |
| 33 | } |
| 34 | |
| 35 | #[allow(single_use_lifetimes)] |
Joel Galenson | 9ceddf6 | 2021-05-19 15:31:02 -0700 | [diff] [blame] | 36 | unsafe fn remove_drop_lifetime<'a, T>( |
| 37 | ptr: unsafe fn(*mut (dyn Future<Output = T> + 'a)), |
| 38 | ) -> unsafe fn(*mut (dyn Future<Output = T> + 'static)) { |
Jason Macnak | ea22e81 | 2020-03-18 01:29:19 +0000 | [diff] [blame] | 39 | mem::transmute(ptr) |
| 40 | } |
| 41 | |
| 42 | impl<'a, T> LocalFutureObj<'a, T> { |
| 43 | /// Create a `LocalFutureObj` from a custom trait object representation. |
| 44 | #[inline] |
Haibo Huang | 0ea7c87 | 2021-02-09 17:15:58 -0800 | [diff] [blame] | 45 | pub fn new<F: UnsafeFutureObj<'a, T> + 'a>(f: F) -> Self { |
| 46 | Self { |
Jason Macnak | ea22e81 | 2020-03-18 01:29:19 +0000 | [diff] [blame] | 47 | future: unsafe { remove_future_lifetime(f.into_raw()) }, |
| 48 | drop_fn: unsafe { remove_drop_lifetime(F::drop) }, |
| 49 | _marker: PhantomData, |
| 50 | } |
| 51 | } |
| 52 | |
| 53 | /// Converts the `LocalFutureObj` into a `FutureObj`. |
| 54 | /// |
| 55 | /// # Safety |
| 56 | /// |
| 57 | /// To make this operation safe one has to ensure that the `UnsafeFutureObj` |
| 58 | /// instance from which this `LocalFutureObj` was created actually |
| 59 | /// implements `Send`. |
| 60 | #[inline] |
| 61 | pub unsafe fn into_future_obj(self) -> FutureObj<'a, T> { |
| 62 | FutureObj(self) |
| 63 | } |
| 64 | } |
| 65 | |
| 66 | impl<T> fmt::Debug for LocalFutureObj<'_, T> { |
| 67 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { |
Joel Galenson | 9ceddf6 | 2021-05-19 15:31:02 -0700 | [diff] [blame] | 68 | f.debug_struct("LocalFutureObj").finish() |
Jason Macnak | ea22e81 | 2020-03-18 01:29:19 +0000 | [diff] [blame] | 69 | } |
| 70 | } |
| 71 | |
| 72 | impl<'a, T> From<FutureObj<'a, T>> for LocalFutureObj<'a, T> { |
| 73 | #[inline] |
Haibo Huang | 0ea7c87 | 2021-02-09 17:15:58 -0800 | [diff] [blame] | 74 | fn from(f: FutureObj<'a, T>) -> Self { |
Jason Macnak | ea22e81 | 2020-03-18 01:29:19 +0000 | [diff] [blame] | 75 | f.0 |
| 76 | } |
| 77 | } |
| 78 | |
| 79 | impl<T> Future for LocalFutureObj<'_, T> { |
| 80 | type Output = T; |
| 81 | |
| 82 | #[inline] |
| 83 | fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T> { |
Joel Galenson | 9ceddf6 | 2021-05-19 15:31:02 -0700 | [diff] [blame] | 84 | unsafe { Pin::new_unchecked(&mut *self.future).poll(cx) } |
Jason Macnak | ea22e81 | 2020-03-18 01:29:19 +0000 | [diff] [blame] | 85 | } |
| 86 | } |
| 87 | |
| 88 | impl<T> Drop for LocalFutureObj<'_, T> { |
| 89 | fn drop(&mut self) { |
Joel Galenson | 9ceddf6 | 2021-05-19 15:31:02 -0700 | [diff] [blame] | 90 | unsafe { (self.drop_fn)(self.future) } |
Jason Macnak | ea22e81 | 2020-03-18 01:29:19 +0000 | [diff] [blame] | 91 | } |
| 92 | } |
| 93 | |
| 94 | /// A custom trait object for polling futures, roughly akin to |
| 95 | /// `Box<dyn Future<Output = T> + Send + 'a>`. |
| 96 | /// |
| 97 | /// This custom trait object was introduced as currently it is not possible to |
| 98 | /// take `dyn Trait` by value and `Box<dyn Trait>` is not available in no_std |
| 99 | /// contexts. |
| 100 | /// |
| 101 | /// You should generally not need to use this type outside of `no_std` or when |
| 102 | /// implementing `Spawn`, consider using `BoxFuture` instead. |
| 103 | pub struct FutureObj<'a, T>(LocalFutureObj<'a, T>); |
| 104 | |
| 105 | impl<T> Unpin for FutureObj<'_, T> {} |
| 106 | unsafe impl<T> Send for FutureObj<'_, T> {} |
| 107 | |
| 108 | impl<'a, T> FutureObj<'a, T> { |
| 109 | /// Create a `FutureObj` from a custom trait object representation. |
| 110 | #[inline] |
Haibo Huang | 0ea7c87 | 2021-02-09 17:15:58 -0800 | [diff] [blame] | 111 | pub fn new<F: UnsafeFutureObj<'a, T> + Send>(f: F) -> Self { |
| 112 | Self(LocalFutureObj::new(f)) |
Jason Macnak | ea22e81 | 2020-03-18 01:29:19 +0000 | [diff] [blame] | 113 | } |
| 114 | } |
| 115 | |
| 116 | impl<T> fmt::Debug for FutureObj<'_, T> { |
| 117 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { |
Joel Galenson | 9ceddf6 | 2021-05-19 15:31:02 -0700 | [diff] [blame] | 118 | f.debug_struct("FutureObj").finish() |
Jason Macnak | ea22e81 | 2020-03-18 01:29:19 +0000 | [diff] [blame] | 119 | } |
| 120 | } |
| 121 | |
| 122 | impl<T> Future for FutureObj<'_, T> { |
| 123 | type Output = T; |
| 124 | |
| 125 | #[inline] |
| 126 | fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T> { |
Joel Galenson | 9ceddf6 | 2021-05-19 15:31:02 -0700 | [diff] [blame] | 127 | Pin::new(&mut self.0).poll(cx) |
Jason Macnak | ea22e81 | 2020-03-18 01:29:19 +0000 | [diff] [blame] | 128 | } |
| 129 | } |
| 130 | |
| 131 | /// A custom implementation of a future trait object for `FutureObj`, providing |
| 132 | /// a vtable with drop support. |
| 133 | /// |
| 134 | /// This custom representation is typically used only in `no_std` contexts, |
| 135 | /// where the default `Box`-based implementation is not available. |
| 136 | /// |
| 137 | /// # Safety |
| 138 | /// |
| 139 | /// See the safety notes on individual methods for what guarantees an |
| 140 | /// implementor must provide. |
| 141 | pub unsafe trait UnsafeFutureObj<'a, T>: 'a { |
| 142 | /// Convert an owned instance into a (conceptually owned) fat pointer. |
| 143 | /// |
| 144 | /// # Safety |
| 145 | /// |
| 146 | /// ## Implementor |
| 147 | /// |
| 148 | /// The trait implementor must guarantee that it is safe to convert the |
| 149 | /// provided `*mut (dyn Future<Output = T> + 'a)` into a `Pin<&mut (dyn |
| 150 | /// Future<Output = T> + 'a)>` and call methods on it, non-reentrantly, |
| 151 | /// until `UnsafeFutureObj::drop` is called with it. |
| 152 | fn into_raw(self) -> *mut (dyn Future<Output = T> + 'a); |
| 153 | |
| 154 | /// Drops the future represented by the given fat pointer. |
| 155 | /// |
| 156 | /// # Safety |
| 157 | /// |
| 158 | /// ## Implementor |
| 159 | /// |
| 160 | /// The trait implementor must guarantee that it is safe to call this |
| 161 | /// function once per `into_raw` invocation. |
| 162 | /// |
| 163 | /// ## Caller |
| 164 | /// |
| 165 | /// The caller must ensure: |
| 166 | /// |
| 167 | /// * the pointer passed was obtained from an `into_raw` invocation from |
| 168 | /// this same trait object |
| 169 | /// * the pointer is not currently in use as a `Pin<&mut (dyn Future<Output |
| 170 | /// = T> + 'a)>` |
| 171 | /// * the pointer must not be used again after this function is called |
| 172 | unsafe fn drop(ptr: *mut (dyn Future<Output = T> + 'a)); |
| 173 | } |
| 174 | |
| 175 | unsafe impl<'a, T, F> UnsafeFutureObj<'a, T> for &'a mut F |
| 176 | where |
Joel Galenson | 9ceddf6 | 2021-05-19 15:31:02 -0700 | [diff] [blame] | 177 | F: Future<Output = T> + Unpin + 'a, |
Jason Macnak | ea22e81 | 2020-03-18 01:29:19 +0000 | [diff] [blame] | 178 | { |
| 179 | fn into_raw(self) -> *mut (dyn Future<Output = T> + 'a) { |
| 180 | self as *mut dyn Future<Output = T> |
| 181 | } |
| 182 | |
| 183 | unsafe fn drop(_ptr: *mut (dyn Future<Output = T> + 'a)) {} |
| 184 | } |
| 185 | |
Joel Galenson | 9ceddf6 | 2021-05-19 15:31:02 -0700 | [diff] [blame] | 186 | unsafe impl<'a, T> UnsafeFutureObj<'a, T> for &'a mut (dyn Future<Output = T> + Unpin + 'a) { |
Jason Macnak | ea22e81 | 2020-03-18 01:29:19 +0000 | [diff] [blame] | 187 | fn into_raw(self) -> *mut (dyn Future<Output = T> + 'a) { |
| 188 | self as *mut dyn Future<Output = T> |
| 189 | } |
| 190 | |
| 191 | unsafe fn drop(_ptr: *mut (dyn Future<Output = T> + 'a)) {} |
| 192 | } |
| 193 | |
| 194 | unsafe impl<'a, T, F> UnsafeFutureObj<'a, T> for Pin<&'a mut F> |
| 195 | where |
Joel Galenson | 9ceddf6 | 2021-05-19 15:31:02 -0700 | [diff] [blame] | 196 | F: Future<Output = T> + 'a, |
Jason Macnak | ea22e81 | 2020-03-18 01:29:19 +0000 | [diff] [blame] | 197 | { |
| 198 | fn into_raw(self) -> *mut (dyn Future<Output = T> + 'a) { |
| 199 | unsafe { self.get_unchecked_mut() as *mut dyn Future<Output = T> } |
| 200 | } |
| 201 | |
| 202 | unsafe fn drop(_ptr: *mut (dyn Future<Output = T> + 'a)) {} |
| 203 | } |
| 204 | |
Joel Galenson | 9ceddf6 | 2021-05-19 15:31:02 -0700 | [diff] [blame] | 205 | unsafe impl<'a, T> UnsafeFutureObj<'a, T> for Pin<&'a mut (dyn Future<Output = T> + 'a)> { |
Jason Macnak | ea22e81 | 2020-03-18 01:29:19 +0000 | [diff] [blame] | 206 | fn into_raw(self) -> *mut (dyn Future<Output = T> + 'a) { |
| 207 | unsafe { self.get_unchecked_mut() as *mut dyn Future<Output = T> } |
| 208 | } |
| 209 | |
| 210 | unsafe fn drop(_ptr: *mut (dyn Future<Output = T> + 'a)) {} |
| 211 | } |
| 212 | |
| 213 | #[cfg(feature = "alloc")] |
| 214 | mod if_alloc { |
| 215 | use super::*; |
| 216 | use alloc::boxed::Box; |
| 217 | |
| 218 | unsafe impl<'a, T, F> UnsafeFutureObj<'a, T> for Box<F> |
Joel Galenson | 9ceddf6 | 2021-05-19 15:31:02 -0700 | [diff] [blame] | 219 | where |
| 220 | F: Future<Output = T> + 'a, |
Jason Macnak | ea22e81 | 2020-03-18 01:29:19 +0000 | [diff] [blame] | 221 | { |
| 222 | fn into_raw(self) -> *mut (dyn Future<Output = T> + 'a) { |
| 223 | Box::into_raw(self) |
| 224 | } |
| 225 | |
| 226 | unsafe fn drop(ptr: *mut (dyn Future<Output = T> + 'a)) { |
| 227 | drop(Box::from_raw(ptr as *mut F)) |
| 228 | } |
| 229 | } |
| 230 | |
| 231 | unsafe impl<'a, T: 'a> UnsafeFutureObj<'a, T> for Box<dyn Future<Output = T> + 'a> { |
| 232 | fn into_raw(self) -> *mut (dyn Future<Output = T> + 'a) { |
| 233 | Box::into_raw(self) |
| 234 | } |
| 235 | |
| 236 | unsafe fn drop(ptr: *mut (dyn Future<Output = T> + 'a)) { |
| 237 | drop(Box::from_raw(ptr)) |
| 238 | } |
| 239 | } |
| 240 | |
| 241 | unsafe impl<'a, T: 'a> UnsafeFutureObj<'a, T> for Box<dyn Future<Output = T> + Send + 'a> { |
| 242 | fn into_raw(self) -> *mut (dyn Future<Output = T> + 'a) { |
| 243 | Box::into_raw(self) |
| 244 | } |
| 245 | |
| 246 | unsafe fn drop(ptr: *mut (dyn Future<Output = T> + 'a)) { |
| 247 | drop(Box::from_raw(ptr)) |
| 248 | } |
| 249 | } |
| 250 | |
| 251 | unsafe impl<'a, T, F> UnsafeFutureObj<'a, T> for Pin<Box<F>> |
| 252 | where |
Joel Galenson | 9ceddf6 | 2021-05-19 15:31:02 -0700 | [diff] [blame] | 253 | F: Future<Output = T> + 'a, |
Jason Macnak | ea22e81 | 2020-03-18 01:29:19 +0000 | [diff] [blame] | 254 | { |
| 255 | fn into_raw(mut self) -> *mut (dyn Future<Output = T> + 'a) { |
| 256 | let ptr = unsafe { self.as_mut().get_unchecked_mut() as *mut _ }; |
| 257 | mem::forget(self); |
| 258 | ptr |
| 259 | } |
| 260 | |
| 261 | unsafe fn drop(ptr: *mut (dyn Future<Output = T> + 'a)) { |
| 262 | drop(Pin::from(Box::from_raw(ptr))) |
| 263 | } |
| 264 | } |
| 265 | |
| 266 | unsafe impl<'a, T: 'a> UnsafeFutureObj<'a, T> for Pin<Box<dyn Future<Output = T> + 'a>> { |
| 267 | fn into_raw(mut self) -> *mut (dyn Future<Output = T> + 'a) { |
| 268 | let ptr = unsafe { self.as_mut().get_unchecked_mut() as *mut _ }; |
| 269 | mem::forget(self); |
| 270 | ptr |
| 271 | } |
| 272 | |
| 273 | unsafe fn drop(ptr: *mut (dyn Future<Output = T> + 'a)) { |
| 274 | drop(Pin::from(Box::from_raw(ptr))) |
| 275 | } |
| 276 | } |
| 277 | |
| 278 | unsafe impl<'a, T: 'a> UnsafeFutureObj<'a, T> for Pin<Box<dyn Future<Output = T> + Send + 'a>> { |
| 279 | fn into_raw(mut self) -> *mut (dyn Future<Output = T> + 'a) { |
| 280 | let ptr = unsafe { self.as_mut().get_unchecked_mut() as *mut _ }; |
| 281 | mem::forget(self); |
| 282 | ptr |
| 283 | } |
| 284 | |
| 285 | unsafe fn drop(ptr: *mut (dyn Future<Output = T> + 'a)) { |
| 286 | drop(Pin::from(Box::from_raw(ptr))) |
| 287 | } |
| 288 | } |
| 289 | |
| 290 | impl<'a, F: Future<Output = ()> + Send + 'a> From<Box<F>> for FutureObj<'a, ()> { |
| 291 | fn from(boxed: Box<F>) -> Self { |
Haibo Huang | 0ea7c87 | 2021-02-09 17:15:58 -0800 | [diff] [blame] | 292 | Self::new(boxed) |
Jason Macnak | ea22e81 | 2020-03-18 01:29:19 +0000 | [diff] [blame] | 293 | } |
| 294 | } |
| 295 | |
| 296 | impl<'a> From<Box<dyn Future<Output = ()> + Send + 'a>> for FutureObj<'a, ()> { |
| 297 | fn from(boxed: Box<dyn Future<Output = ()> + Send + 'a>) -> Self { |
Haibo Huang | 0ea7c87 | 2021-02-09 17:15:58 -0800 | [diff] [blame] | 298 | Self::new(boxed) |
Jason Macnak | ea22e81 | 2020-03-18 01:29:19 +0000 | [diff] [blame] | 299 | } |
| 300 | } |
| 301 | |
| 302 | impl<'a, F: Future<Output = ()> + Send + 'a> From<Pin<Box<F>>> for FutureObj<'a, ()> { |
| 303 | fn from(boxed: Pin<Box<F>>) -> Self { |
Haibo Huang | 0ea7c87 | 2021-02-09 17:15:58 -0800 | [diff] [blame] | 304 | Self::new(boxed) |
Jason Macnak | ea22e81 | 2020-03-18 01:29:19 +0000 | [diff] [blame] | 305 | } |
| 306 | } |
| 307 | |
| 308 | impl<'a> From<Pin<Box<dyn Future<Output = ()> + Send + 'a>>> for FutureObj<'a, ()> { |
| 309 | fn from(boxed: Pin<Box<dyn Future<Output = ()> + Send + 'a>>) -> Self { |
Haibo Huang | 0ea7c87 | 2021-02-09 17:15:58 -0800 | [diff] [blame] | 310 | Self::new(boxed) |
Jason Macnak | ea22e81 | 2020-03-18 01:29:19 +0000 | [diff] [blame] | 311 | } |
| 312 | } |
| 313 | |
| 314 | impl<'a, F: Future<Output = ()> + 'a> From<Box<F>> for LocalFutureObj<'a, ()> { |
| 315 | fn from(boxed: Box<F>) -> Self { |
Haibo Huang | 0ea7c87 | 2021-02-09 17:15:58 -0800 | [diff] [blame] | 316 | Self::new(boxed) |
Jason Macnak | ea22e81 | 2020-03-18 01:29:19 +0000 | [diff] [blame] | 317 | } |
| 318 | } |
| 319 | |
| 320 | impl<'a> From<Box<dyn Future<Output = ()> + 'a>> for LocalFutureObj<'a, ()> { |
| 321 | fn from(boxed: Box<dyn Future<Output = ()> + 'a>) -> Self { |
Haibo Huang | 0ea7c87 | 2021-02-09 17:15:58 -0800 | [diff] [blame] | 322 | Self::new(boxed) |
Jason Macnak | ea22e81 | 2020-03-18 01:29:19 +0000 | [diff] [blame] | 323 | } |
| 324 | } |
| 325 | |
| 326 | impl<'a, F: Future<Output = ()> + 'a> From<Pin<Box<F>>> for LocalFutureObj<'a, ()> { |
| 327 | fn from(boxed: Pin<Box<F>>) -> Self { |
Haibo Huang | 0ea7c87 | 2021-02-09 17:15:58 -0800 | [diff] [blame] | 328 | Self::new(boxed) |
Jason Macnak | ea22e81 | 2020-03-18 01:29:19 +0000 | [diff] [blame] | 329 | } |
| 330 | } |
| 331 | |
| 332 | impl<'a> From<Pin<Box<dyn Future<Output = ()> + 'a>>> for LocalFutureObj<'a, ()> { |
| 333 | fn from(boxed: Pin<Box<dyn Future<Output = ()> + 'a>>) -> Self { |
Haibo Huang | 0ea7c87 | 2021-02-09 17:15:58 -0800 | [diff] [blame] | 334 | Self::new(boxed) |
Jason Macnak | ea22e81 | 2020-03-18 01:29:19 +0000 | [diff] [blame] | 335 | } |
| 336 | } |
| 337 | } |