Try dropchk patch again

Author: juntyr

src/lib.rs

@@ -28,6 +28,7 @@
 #![feature(sync_unsafe_cell)]
 #![feature(never_type)]
 #![feature(layout_for_ptr)]
+#![feature(dropck_eyepatch)]
 #![feature(cfg_version)]
 #![cfg_attr(not(version("1.76.0")), feature(c_str_literals))]
 #![cfg_attr(not(version("1.76.0")), feature(ptr_from_ref))]

src/utils/async.rs

@@ -1,5 +1,8 @@
 #[cfg(feature = "host")]
-use std::{borrow::BorrowMut, future::Future, future::IntoFuture, marker::PhantomData, task::Poll};
+use std::{
+    borrow::BorrowMut, future::Future, future::IntoFuture, marker::PhantomData, ptr::NonNull,
+    task::Poll,
+};
 
 #[cfg(feature = "host")]
 use rustacuda::{
@@ -272,7 +275,7 @@ impl<'a, 'stream, T: BorrowMut<C::Completed>, C: Completion<T>> Async<'a, 'strea
         }
     }
 
-    pub const fn as_ref(&self) -> AsyncProj<'_, 'stream, &T> {
+    pub fn as_ref(&self) -> AsyncProj<'_, 'stream, &T> {
         // Safety: this projection captures this async
         unsafe { AsyncProj::new(&self.value, None) }
     }
@@ -327,7 +330,7 @@ impl<'a, 'stream, T: BorrowMut<C::Completed>, C: Completion<T>> Async<'a, 'strea
             (
                 this.stream,
                 std::ptr::read(&this.value),
-                (std::ptr::read(&this.status)),
+                std::ptr::read(&this.status),
             )
         }
     }
@@ -478,7 +481,19 @@ pub struct AsyncProj<'a, 'stream, T: 'a> {
     _capture: PhantomData<&'a ()>,
     _stream: PhantomData<Stream<'stream>>,
     value: T,
-    use_callback: Option<Box<dyn FnMut() -> CudaResult<()> + 'a>>,
+    use_callback: Option<NonNull<dyn FnMut() -> CudaResult<()> + 'a>>,
+}
+
+#[cfg(feature = "host")]
+unsafe impl<#[may_dangle] 'a, 'stream, T: 'a> Drop for AsyncProj<'a, 'stream, T> {
+    fn drop(&mut self) {
+        if let Some(use_callback) = self.use_callback.take() {
+            let layout = unsafe { std::alloc::Layout::for_value_raw(use_callback.as_ptr()) };
+            if layout.size() != 0 {
+                unsafe { std::alloc::dealloc(use_callback.as_ptr().cast(), layout) };
+            }
+        }
+    }
 }
 
 #[cfg(feature = "host")]
@@ -489,15 +504,16 @@ impl<'a, 'stream, T: 'a> AsyncProj<'a, 'stream, T> {
     /// This projection must either capture an existing [`Async`] or come from
     /// a source that ensures that the projected value can never (async) move
     /// to a different [`Stream`].
-    pub(crate) const unsafe fn new(
+    pub(crate) unsafe fn new(
         value: T,
         use_callback: Option<Box<dyn FnMut() -> CudaResult<()> + 'a>>,
     ) -> Self {
         Self {
             _capture: PhantomData::<&'a ()>,
             _stream: PhantomData::<Stream<'stream>>,
             value,
-            use_callback,
+            use_callback: use_callback
+                .map(|use_callback| unsafe { NonNull::new_unchecked(Box::into_raw(use_callback)) }),
         }
     }
 
@@ -511,7 +527,8 @@ impl<'a, 'stream, T: 'a> AsyncProj<'a, 'stream, T> {
     /// computation out of smaller ones that have all been submitted to the
     /// same [`Stream`].
     pub(crate) unsafe fn unwrap_unchecked(self) -> T {
-        self.value
+        let (value, _use_callback) = self.unwrap_unchecked_with_use();
+        value
     }
 
     #[allow(clippy::type_complexity)]
@@ -527,7 +544,21 @@ impl<'a, 'stream, T: 'a> AsyncProj<'a, 'stream, T> {
     pub(crate) unsafe fn unwrap_unchecked_with_use(
         self,
     ) -> (T, Option<Box<dyn FnMut() -> CudaResult<()> + 'a>>) {
-        (self.value, self.use_callback)
+        let this = std::mem::ManuallyDrop::new(self);
+
+        // Safety: we destructure self into its droppable components,
+        //         value and use_callback, without dropping self itself
+        let (value, use_callback) = unsafe {
+            (
+                std::ptr::read(&this.value),
+                std::ptr::read(&this.use_callback),
+            )
+        };
+
+        (
+            value,
+            use_callback.map(|use_callback| Box::from_raw(use_callback.as_ptr())),
+        )
     }
 }
 
@@ -556,16 +587,20 @@ impl<'a, 'stream, T: 'a> AsyncProj<'a, 'stream, T> {
             _stream: PhantomData::<Stream<'stream>>,
             value: &mut self.value,
             use_callback: self.use_callback.as_mut().map(|use_callback| {
+                let use_callback: &mut dyn FnMut() -> CudaResult<()> =
+                    unsafe { &mut *use_callback.as_ptr() };
                 let use_callback: Box<dyn FnMut() -> CudaResult<()>> = Box::new(use_callback);
-                use_callback
+                unsafe { NonNull::new_unchecked(Box::into_raw(use_callback)) }
             }),
         }
     }
 
     pub(crate) fn record_mut_use(&mut self) -> CudaResult<()> {
-        self.use_callback
-            .as_mut()
-            .map_or(Ok(()), |use_callback| use_callback())
+        self.use_callback.as_mut().map_or(Ok(()), |use_callback| {
+            let use_callback: &mut dyn FnMut() -> CudaResult<()> =
+                unsafe { &mut *use_callback.as_ptr() };
+            use_callback()
+        })
     }
 }
 
@@ -623,8 +658,10 @@ impl<'a, 'stream, T: 'a> AsyncProj<'a, 'stream, &'a mut T> {
             _stream: PhantomData::<Stream<'stream>>,
             value: self.value,
             use_callback: self.use_callback.as_mut().map(|use_callback| {
+                let use_callback: &mut dyn FnMut() -> CudaResult<()> =
+                    unsafe { &mut *use_callback.as_ptr() };
                 let use_callback: Box<dyn FnMut() -> CudaResult<()>> = Box::new(use_callback);
-                use_callback
+                unsafe { NonNull::new_unchecked(Box::into_raw(use_callback)) }
             }),
         }
     }