diff --git a/tailcall/src/slot.rs b/tailcall/src/slot.rs index 7141baa..469e8a5 100644 --- a/tailcall/src/slot.rs +++ b/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 { - bytes: MaybeUninit<[u8; SIZE]>, +pub struct Slot { + bytes: MaybeUninit<[u8; SLOT_SIZE]>, } -impl Default for Slot { +impl Default for Slot { fn default() -> Self { Self::new() } } -impl Slot { +impl Slot { pub const fn new() -> Self { Self { bytes: MaybeUninit::uninit(), @@ -23,7 +29,7 @@ impl Slot { let slot_ptr = self as *mut _; // Verify the size and alignment of T. - assert!(size_of::() <= SIZE); + assert!(size_of::() <= size_of::()); assert!(align_of::() <= align_of::()); // SAFETY: We just checked the size and alignment of T. Since we are @@ -39,3 +45,87 @@ impl Slot { casted } } + +pub struct SlotBox<'slot, T: ?Sized + 'slot> { + pointer: NonNull, + _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(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 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) } + } +} diff --git a/tailcall/src/thunk.rs b/tailcall/src/thunk.rs index 832e039..a983db9 100644 --- a/tailcall/src/thunk.rs +++ b/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 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)