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);
+}