Handle drop

tailcall/src/lib.rs

@@ -10,5 +10,6 @@
 
 pub use tailcall_impl::tailcall;
 
+pub(crate) mod slot;
 pub mod thunk;
 pub mod trampoline;

tailcall/src/slot.rs

@@ -0,0 +1,69 @@
+use core::mem::{align_of, size_of, ManuallyDrop, MaybeUninit};
+
+#[repr(C, align(16))]
+pub struct Slot<const SIZE: usize> {
+    bytes: MaybeUninit<[u8; SIZE]>,
+}
+
+#[repr(C)]
+union SlotView<T, const SIZE: usize> {
+    value: ManuallyDrop<T>,
+    slot: ManuallyDrop<Slot<SIZE>>,
+}
+
+impl<const SIZE: usize> Slot<SIZE> {
+    pub const fn uninit() -> Self {
+        Self {
+            bytes: MaybeUninit::uninit(),
+        }
+    }
+
+    pub const fn new<T>(value: T) -> Self {
+        assert!(
+            align_of::<T>() <= align_of::<Self>(),
+            "unsupport value alignment",
+        );
+
+        assert!(
+            size_of::<T>() <= size_of::<Self>(),
+            "value size exceeds slot capacity",
+        );
+
+        SlotView::of_value(value).into_slot()
+    }
+
+    // SAFETY: The caller must ensure that `self` contains a valid `T`.
+    pub const unsafe fn into_value<T>(self) -> T {
+        unsafe { SlotView::of_slot(self).into_value() }
+    }
+}
+
+impl<const SIZE: usize> Default for Slot<SIZE> {
+    fn default() -> Self {
+        Self::uninit()
+    }
+}
+
+impl<T, const SIZE: usize> SlotView<T, SIZE> {
+    const fn of_value(value: T) -> Self {
+        Self {
+            value: ManuallyDrop::new(value),
+        }
+    }
+
+    const fn of_slot(slot: Slot<SIZE>) -> Self {
+        Self {
+            slot: ManuallyDrop::new(slot),
+        }
+    }
+
+    const fn into_slot(self) -> Slot<SIZE> {
+        // SAFETY: `Slot<SIZE>` is valid at all bit patterns.
+        ManuallyDrop::into_inner(unsafe { self.slot })
+    }
+
+    // SAFETY: The caller must ensure that `self` contains a valid `T`.
+    const unsafe fn into_value(self) -> T {
+        ManuallyDrop::into_inner(unsafe { self.value })
+    }
+}

tailcall/src/thunk.rs

@@ -1,122 +1,81 @@
-use core::marker::PhantomData;
+use crate::slot::Slot;
+use core::{marker::PhantomData, mem::transmute, ptr::drop_in_place};
 
-mod slot {
-    use core::mem::{align_of, size_of, ManuallyDrop, MaybeUninit};
+const MAX_THUNK_DATA_SIZE: usize = 48;
 
-    #[repr(C, align(16))]
-    pub struct Slot<const SIZE: usize> {
-        bytes: MaybeUninit<[u8; SIZE]>,
-    }
-
-    union SlotView<T, const SIZE: usize> {
-        value: ManuallyDrop<T>,
-        slot: ManuallyDrop<Slot<SIZE>>,
-    }
-
-    impl<const SIZE: usize> Slot<SIZE> {
-        pub const fn new<T>(value: T) -> Self {
-            Self::assert_valid_at::<T>();
-
-            SlotView::of_value(value).into_slot()
-        }
-
-        // SAFETY: The caller must ensure that `self` contains a valid `T`.
-        pub const unsafe fn into_inner<T>(self) -> T {
-            Self::assert_valid_at::<T>();
-
-            unsafe { SlotView::of_slot(self).into_value() }
-        }
-
-        const fn assert_valid_at<T>() {
-            assert!(size_of::<T>() <= size_of::<Self>());
-            assert!(align_of::<T>() <= align_of::<Self>());
-        }
-    }
-
-    impl<T, const SIZE: usize> SlotView<T, SIZE> {
-        const fn of_value(value: T) -> Self {
-            Self {
-                value: ManuallyDrop::new(value),
-            }
-        }
+type ThunkSlot = Slot<MAX_THUNK_DATA_SIZE>;
+type CallFn<T> = fn(ThunkSlot) -> T;
+type DropInPlaceFn = unsafe fn(*mut ThunkSlot);
 
-        const fn into_slot(self) -> Slot<SIZE> {
-            // SAFETY: `Slot<SIZE>` is valid at all bit patterns.
-            ManuallyDrop::into_inner(unsafe { self.slot })
-        }
-
-        const fn of_slot(slot: Slot<SIZE>) -> Self {
-            Self {
-                slot: ManuallyDrop::new(slot),
-            }
-        }
-
-        // SAFETY: The caller must ensure that `self` contains a valid `T`.
-        const unsafe fn into_value(self) -> T {
-            ManuallyDrop::into_inner(unsafe { self.value })
-        }
-    }
+#[repr(transparent)]
+pub struct Thunk<'a, T = ()> {
+    inner: ThunkInner<'a, T>,
 }
 
-mod guard {
-    use core::mem::forget;
-
-    pub struct Gurad();
+struct ThunkInner<'a, T> {
+    slot: ThunkSlot,
+    call_fn: CallFn<T>,
+    drop_in_place_fn: DropInPlaceFn,
+    _marker: PhantomData<dyn FnOnce() -> T + 'a>,
+}
 
-    impl Gurad {
-        pub const fn new() -> Self {
-            Self()
+impl<'a, T> Thunk<'a, T> {
+    pub const fn new<F>(fn_once: F) -> Self
+    where
+        F: FnOnce() -> T + 'a,
+    {
+        Self {
+            inner: ThunkInner::new(fn_once),
         }
+    }
 
-        pub const fn disarm(self) {
-            forget(self)
-        }
+    #[inline(always)]
+    pub fn call(self) -> T {
+        self.into_inner().call()
     }
 
-    impl Drop for Gurad {
-        fn drop(&mut self) {
-            unreachable!()
-        }
+    const fn into_inner(self) -> ThunkInner<'a, T> {
+        // SAFETY: `Thunk` is a transparent wrapper around `ThunkInner`.
+        unsafe { transmute(self) }
     }
 }
 
-use guard::Gurad;
-use slot::Slot;
-
-const SLOT_SIZE: usize = 48;
-
-#[must_use]
-pub struct Thunk<'a, T = ()> {
-    guard: Gurad,
-    slot: Slot<SLOT_SIZE>,
-    call_impl: fn(Slot<SLOT_SIZE>) -> T,
-    _marker: PhantomData<dyn FnOnce() -> T + 'a>,
+impl<'a, T> Drop for Thunk<'a, T> {
+    fn drop(&mut self) {
+        // SAFETY: We own `inner`, and it cannot be used after dropping.
+        unsafe { self.inner.drop_in_place() }
+    }
 }
 
-impl<'a, T> Thunk<'a, T> {
+impl<'a, T> ThunkInner<'a, T> {
     pub const fn new<F>(fn_once: F) -> Self
     where
         F: FnOnce() -> T + 'a,
     {
         Self {
-            guard: Gurad::new(),
             slot: Slot::new(fn_once),
-            call_impl: |slot| unsafe { slot.into_inner::<F>()() },
+            call_fn: |slot| {
+                // SAFETY: `slot` is initialized above with `F`.
+                unsafe { slot.into_value::<F>()() }
+            },
+            drop_in_place_fn: |slot_ptr| {
+                // SAFETY: `slot` is initialized above with `F`.
+                unsafe { drop_in_place(slot_ptr.cast::<F>()) };
+            },
             _marker: PhantomData,
         }
     }
 
     #[inline(always)]
     pub fn call(self) -> T {
-        let Self {
-            call_impl,
-            slot,
-            guard,
-            _marker,
-        } = self;
+        let Self { slot, call_fn, .. } = self;
 
-        guard.disarm();
+        call_fn(slot)
+    }
 
-        call_impl(slot)
+    // SAFETY: `Self::call` cannot be called after dropping in place.
+    #[inline(always)]
+    pub unsafe fn drop_in_place(&mut self) {
+        unsafe { (self.drop_in_place_fn)(&mut self.slot) }
     }
 }

tailcall/tests/thunk.rs

@@ -16,17 +16,17 @@ fn with_captures() {
     assert_eq!(thunk.call(), x + y);
 }
 
-// #[test]
-// fn with_too_many_captures() {
-//     let a = 1;
-//     let b = 2;
-//     let c = 3;
-//     let d = 4;
-//     let e = 5;
-//     let f = 6;
-//     let g = 7;
-
-//     let thunk = Thunk::new(move || a + b + c + d + e + f + g);
+#[test]
+#[should_panic]
+fn with_too_many_captures() {
+    let a = 1usize;
+    let b = 2usize;
+    let c = 3usize;
+    let d = 4usize;
+    let e = 5usize;
+    let f = 6usize;
+    let g = 7usize;
+    let h = 8usize;
 
-//     assert_eq!(thunk.call(), a + b + c + d + e + f + g,);
-// }
+    Thunk::new(move || a + b + c + d + e + f + g + h);
+}