util: simplify ReusableBoxFuture::try_set logic

tokio-util/src/sync/reusable_box.rs

@@ -1,9 +1,7 @@
-use std::alloc::Layout;
 use std::fmt;
-use std::future::{self, Future};
-use std::mem::{self, ManuallyDrop};
+use std::future::Future;
+use std::mem::{align_of_val, size_of_val};
 use std::pin::Pin;
-use std::ptr;
 use std::task::{Context, Poll};
 
 /// A reusable `Pin<Box<dyn Future<Output = T> + Send + 'a>>`.
@@ -47,24 +45,14 @@ impl<'a, T> ReusableBoxFuture<'a, T> {
     where
         F: Future<Output = T> + Send + 'a,
     {
-        // If we try to inline the contents of this function, the type checker complains because
-        // the bound `T: 'a` is not satisfied in the call to `pending()`. But by putting it in an
-        // inner function that doesn't have `T` as a generic parameter, we implicitly get the bound
-        // `F::Output: 'a` transitively through `F: 'a`, allowing us to call `pending()`.
-        #[inline(always)]
-        fn real_try_set<'a, F>(
-            this: &mut ReusableBoxFuture<'a, F::Output>,
-            future: F,
-        ) -> Result<(), F>
-        where
-            F: Future + Send + 'a,
+        if size_of_val(&*self.boxed) == size_of_val(&future)
+            && align_of_val(&*self.boxed) == align_of_val(&future)
         {
-            // future::Pending<T> is a ZST so this never allocates.
-            let boxed = mem::replace(&mut this.boxed, Box::pin(future::pending()));
-            reuse_pin_box(boxed, future, |boxed| this.boxed = Pin::from(boxed))
+            self.boxed = Box::pin(future);
+            Ok(())
+        } else {
+            Err(future)
         }
-
-        real_try_set(self, future)
     }
 
     /// Get a pinned reference to the underlying future.
@@ -97,61 +85,3 @@ impl<T> fmt::Debug for ReusableBoxFuture<'_, T> {
         f.debug_struct("ReusableBoxFuture").finish()
     }
 }
-
-fn reuse_pin_box<T: ?Sized, U, O, F>(boxed: Pin<Box<T>>, new_value: U, callback: F) -> Result<O, U>
-where
-    F: FnOnce(Box<U>) -> O,
-{
-    let layout = Layout::for_value::<T>(&*boxed);
-    if layout != Layout::new::<U>() {
-        return Err(new_value);
-    }
-
-    // SAFETY: We don't ever construct a non-pinned reference to the old `T` from now on, and we
-    // always drop the `T`.
-    let raw: *mut T = Box::into_raw(unsafe { Pin::into_inner_unchecked(boxed) });
-
-    // When dropping the old value panics, we still want to call `callback` — so move the rest of
-    // the code into a guard type.
-    let guard = CallOnDrop::new(|| {
-        let raw: *mut U = raw.cast::<U>();
-        unsafe { raw.write(new_value) };
-
-        // SAFETY:
-        // - `T` and `U` have the same layout.
-        // - `raw` comes from a `Box` that uses the same allocator as this one.
-        // - `raw` points to a valid instance of `U` (we just wrote it in).
-        let boxed = unsafe { Box::from_raw(raw) };
-
-        callback(boxed)
-    });
-
-    // Drop the old value.
-    unsafe { ptr::drop_in_place(raw) };
-
-    // Run the rest of the code.
-    Ok(guard.call())
-}
-
-struct CallOnDrop<O, F: FnOnce() -> O> {
-    f: ManuallyDrop<F>,
-}
-
-impl<O, F: FnOnce() -> O> CallOnDrop<O, F> {
-    fn new(f: F) -> Self {
-        let f = ManuallyDrop::new(f);
-        Self { f }
-    }
-    fn call(self) -> O {
-        let mut this = ManuallyDrop::new(self);
-        let f = unsafe { ManuallyDrop::take(&mut this.f) };
-        f()
-    }
-}
-
-impl<O, F: FnOnce() -> O> Drop for CallOnDrop<O, F> {
-    fn drop(&mut self) {
-        let f = unsafe { ManuallyDrop::take(&mut self.f) };
-        f();
-    }
-}