Refactor and add safety comments

tailcall/src/lib.rs

@@ -1,4 +1,5 @@
 #![no_std]
+#![deny(unsafe_op_in_unsafe_fn)]
 
 pub use tailcall_impl::tailcall;
 

tailcall/src/slot.rs

@@ -5,32 +5,37 @@ pub struct Slot<const SIZE: usize = 128> {
  bytes: MaybeUninit<[u8; SIZE]>,
 }
 
+impl<const SIZE: usize> Default for Slot<SIZE> {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
 impl<const SIZE: usize> Slot<SIZE> {
  pub const fn new() -> Self {
  Self {
  bytes: MaybeUninit::uninit(),
  }
  }
 
- pub unsafe fn take<T>(in_slot: &mut T) -> (T, &mut Self) {
-  let in_slot: *mut T = in_slot;
- debug_assert!((in_slot as usize) % align_of::<Self>() == 0);
+ #[inline(always)]
+ pub fn cast<T>(&mut self) -> &mut MaybeUninit<T> {
+ let slot_ptr = self as *mut _;
 
- let slot: &mut Self = &mut *in_slot.cast();
- let value = slot.cast().assume_init_read();
+ // Verify the size and alignment of T.
+ assert!(size_of::<T>() <= SIZE);
+ assert!(align_of::<T>() <= align_of::<Self>());
 
- (value, slot)
- }
-
- pub fn put<T>(&mut self, value: T) -> &mut T {
- self.cast().write(value)
- }
+ // SAFETY: We just checked the size and alignment of T. Since we are
+ // only returning `MaybeUninit<T>`, we need not worry about the bits.
+ let casted = unsafe { &mut *self.bytes.as_mut_ptr().cast() };
+ let casted_ptr = casted as *mut _;
 
- fn cast<T>(&mut self) -> &mut MaybeUninit<T> {
- debug_assert!(size_of::<T>() <= SIZE);
- debug_assert!(align_of::<T>() <= align_of::<Self>());
+ // Verify that the address of the pointer has not actually changed. This
+ // ensures that it is safe to transume between `&mut Slot` and `&mut T`
+ // (provided that there is a `T` in the slot).
+ assert!(casted_ptr as usize == slot_ptr as usize);
 
- // SAFETY: We just checked the size and alignment of T.
- unsafe { &mut *self.bytes.as_mut_ptr().cast() }
+ casted
  }
 }

tailcall/src/thunk.rs

@@ -10,27 +10,34 @@ impl<'slot, T> Thunk<'slot, T> {
  where
  F: FnOnce(&'slot mut Slot) -> T + 'slot,
  {
- let ptr = slot.put(fn_once);
-
- Self { ptr }
+ Self {
+ ptr: slot.cast().write(fn_once),
+ }
  }
 
  #[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() }
  }
 }
 
 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) }
  }
 }
 
 trait ThunkFn<'slot, T>: FnOnce(&'slot mut Slot) -> T {
+ // SAFETY: This method may only be called once and `self` must be stored in
+ // a `Slot`.
  unsafe fn call_once_in_slot(&'slot mut self) -> T;
 }
 
@@ -39,8 +46,17 @@ where
  F: FnOnce(&'slot mut Slot) -> T,
 {
  unsafe fn call_once_in_slot(&'slot mut self) -> T {
- let (fn_once, slot) = Slot::take(self);
+ // 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() };
 
+ // Call the underlying function with the now empty slot.
  fn_once(slot)
  }
 }