Extract lifetime management from Thunk

tailcall/src/slot.rs

@@ -1,17 +1,23 @@
-use core::mem::{align_of, size_of, MaybeUninit};
+use core::{
+    marker::PhantomData,
+    mem::{align_of, forget, size_of, MaybeUninit},
+    ptr::{drop_in_place, NonNull},
+};
+
+const SLOT_SIZE: usize = 128;
 
 #[repr(C, align(128))]
-pub struct Slot<const SIZE: usize = 128> {
-    bytes: MaybeUninit<[u8; SIZE]>,
+pub struct Slot {
+    bytes: MaybeUninit<[u8; SLOT_SIZE]>,
 }
 
-impl<const SIZE: usize> Default for Slot<SIZE> {
+impl Default for Slot {
     fn default() -> Self {
         Self::new()
     }
 }
 
-impl<const SIZE: usize> Slot<SIZE> {
+impl Slot {
     pub const fn new() -> Self {
         Self {
             bytes: MaybeUninit::uninit(),
@@ -23,7 +29,7 @@ impl<const SIZE: usize> Slot<SIZE> {
         let slot_ptr = self as *mut _;
 
         // Verify the size and alignment of T.
-        assert!(size_of::<T>() <= SIZE);
+        assert!(size_of::<T>() <= size_of::<Self>());
         assert!(align_of::<T>() <= align_of::<Self>());
 
         // SAFETY: We just checked the size and alignment of T. Since we are
@@ -39,3 +45,87 @@ impl<const SIZE: usize> Slot<SIZE> {
         casted
     }
 }
+
+pub struct SlotBox<'slot, T: ?Sized + 'slot> {
+    pointer: NonNull<T>,
+    _marker: PhantomData<(&'slot mut Slot, T)>,
+}
+
+impl<'slot, T: ?Sized> SlotBox<'slot, T> {
+    /// # Safety
+    ///
+    /// The caller must ensure that `value` is stored in a `Slot`.
+    pub unsafe fn adopt(value: &'slot mut T) -> Self {
+        Self {
+            pointer: value.into(),
+            _marker: PhantomData,
+        }
+    }
+
+    pub fn coerce<U, F>(slot_box: Self, coerce_fn: F) -> SlotBox<'slot, U>
+    where
+        U: ?Sized,
+        F: FnOnce(&mut T) -> &mut U,
+    {
+        let leaked = Self::leak(slot_box);
+        let leaked_ptr = leaked as *mut _ as *mut u8;
+
+        let coerced = coerce_fn(leaked);
+        let coerced_ptr = coerced as *mut _ as *mut u8;
+
+        assert!(leaked_ptr as usize == coerced_ptr as usize);
+
+        // SAFETY: Since the addresss of the pointer did not change, we know
+        // that the value is still in a slot and only the type has changed.
+        unsafe { SlotBox::adopt(coerced) }
+    }
+
+    pub fn leak(slot_box: Self) -> &'slot mut T {
+        let value_ptr = slot_box.pointer.as_ptr();
+        forget(slot_box);
+
+        // SAFETY: We know that the value is in the `Slot` because we placed it
+        // there in `SlotBox::new_in`. Since the value cannot otherwise be
+        // dropped, the reference is valid for the lifetime of the `Slot`.
+        unsafe { &mut *value_ptr }
+    }
+
+    fn leak_as_slot(slot_box: Self) -> &'slot mut Slot {
+        let slot_ptr: *mut Slot = slot_box.pointer.as_ptr().cast();
+        forget(slot_box);
+
+        // SAFETY: We checked in `Slot::cast` that the address of the value is
+        // also the address of the slot.
+        unsafe { &mut *slot_ptr }
+    }
+}
+
+impl<'slot, T> SlotBox<'slot, T> {
+    pub fn new_in(slot: &'slot mut Slot, value: T) -> Self {
+        let value_ptr = slot.cast().write(value);
+
+        Self {
+            pointer: value_ptr.into(),
+            _marker: PhantomData,
+        }
+    }
+
+    pub fn unwrap(slot_box: Self) -> (&'slot mut Slot, T) {
+        let slot = Self::leak_as_slot(slot_box);
+
+        // SAFETY: We know there is a `T` in the `Slot` because we placed it
+        // there in `SlotBox::new_in`.
+        let value: T = unsafe { slot.cast().assume_init_read() };
+
+        (slot, value)
+    }
+}
+
+impl<T: ?Sized> Drop for SlotBox<'_, T> {
+    fn drop(&mut self) {
+        let value_ptr = self.pointer.as_ptr();
+
+        // SAFETY: The `SlotBox` logically owns this pointer.
+        unsafe { drop_in_place(value_ptr) }
+    }
+}

tailcall/src/thunk.rs

@@ -1,7 +1,7 @@
-use crate::slot::Slot;
+use crate::slot::{Slot, SlotBox};
 
 pub struct Thunk<'slot, T> {
-    ptr: &'slot mut dyn ThunkFn<'slot, T>,
+    thunk_fn: SlotBox<'slot, dyn ThunkFn<'slot, T>>,
 }
 
 impl<'slot, T> Thunk<'slot, T> {
@@ -10,28 +10,23 @@ impl<'slot, T> Thunk<'slot, T> {
     where
         F: FnOnce(&'slot mut Slot) -> T + 'slot,
     {
+        let fn_once = SlotBox::new_in(slot, fn_once);
+
         Self {
-            ptr: slot.cast().write(fn_once),
+            // Convert the thin pointer to `F` into a fat pointer to a
+            // `dyn ThunkFn`. This is required since stable Rust does not yet
+            // support "unsized coercions" on user-defined pointer types.
+            thunk_fn: SlotBox::coerce(fn_once, |p| p as _),
         }
     }
 
     #[inline(always)]
     pub fn call(self) -> T {
-        let ptr: *mut dyn ThunkFn<'slot, T> = self.ptr;
-        core::mem::forget(self);
-
-        // SAFETY: The only way to create a `Thunk` is through `Thunk::new_in`
-        // which stores the value in a `Slot`. Additionally, we just forgot
-        // `self`, so we know that it is impossible to call this method again.
-        unsafe { (*ptr).call_once_in_slot() }
-    }
-}
+        let thunk_fn = SlotBox::leak(self.thunk_fn);
 
-impl<T> Drop for Thunk<'_, T> {
-    fn drop(&mut self) {
-        // SAFETY: The owned value was stored in a `Slot` which does not drop,
-        // and this struct has a unique pointer to the value there.
-        unsafe { core::ptr::drop_in_place(self.ptr) }
+        // SAFETY: `thunk_fn` is in a `Slot` and since this function takes
+        // ownership of self, it cannot be called again.
+        unsafe { thunk_fn.call_once_in_slot() }
     }
 }
 
@@ -46,15 +41,9 @@ where
     F: FnOnce(&'slot mut Slot) -> T,
 {
     unsafe fn call_once_in_slot(&'slot mut self) -> T {
-        // SAFETY: Our caller guarantees that `self` is currently in a `Slot`,
-        // and `Slot` guarantees that it is safe to transmute between `&mut F`
-        // and `&mut Slot`.
-        let slot: &'slot mut Slot = unsafe { core::mem::transmute(self) };
-
-        // SAFETY: We know that there is a `F` in the slot because this method
-        // was just called on it. Although the bits remain the same, logically,
-        // `fn_once` has been moved *out* of the slot beyond this point.
-        let fn_once: F = unsafe { slot.cast().assume_init_read() };
+        // SAFETY: Our caller garentees that `self` is stored in a `Slot`.
+        let in_slot = unsafe { SlotBox::adopt(self) };
+        let (slot, fn_once) = SlotBox::unwrap(in_slot);
 
         // Call the underlying function with the now empty slot.
         fn_once(slot)