thin trait object pointers in Rust

thin_trait_object_pointers.rs

#![feature(thin_box, ptr_metadata, unsize, extern_types, const_mut_refs, const_refs_to_cell)]

use std::{
  borrow::{Borrow, BorrowMut},
  boxed::ThinBox,
  marker::{PhantomData, Unsize},
  mem::MaybeUninit,
  ops::{Deref, DerefMut},
  ptr::Pointee,
};

/// `&Thin<dyn Foo>` is a thin pointer (has a size equal to `usize`), but can be
/// dereferenced to get a `&dyn Foo`. This is somewhat like a `&Box<dyn Foo>`,
/// except it avoids the additional allocation.
///
/// This type is unsized, so values of this type cannot be owned, but references
/// to this type can be created by dereferencing a `KnownThin` or `ThinBox`.
#[repr(transparent)]
pub struct Thin<U: ?Sized + Pointee>(PhantomData<U::Metadata>, Unsized);

extern "C" {
  type Unsized;
}

impl<U: ?Sized + Pointee> Thin<U> {
  fn metadata(&self) -> U::Metadata {
    unsafe { *(self as *const _ as *const U::Metadata).offset(-1) }
  }
}

impl<U: ?Sized + Pointee> Deref for Thin<U> {
  type Target = U;
  #[inline(always)]
  fn deref(&self) -> &Self::Target {
    unsafe { &*std::ptr::from_raw_parts(self as *const _ as *const _, self.metadata()) }
  }
}

impl<U: ?Sized + Pointee> DerefMut for Thin<U> {
  #[inline(always)]
  fn deref_mut(&mut self) -> &mut Self::Target {
    unsafe { &mut *std::ptr::from_raw_parts_mut(self as *const _ as *mut _, self.metadata()) }
  }
}

#[repr(C)]
struct ThinHeaderLayout<U: ?Sized + Pointee, T>(U::Metadata, [T; 0]);

/// An owned reference to a `T` with a `U::Metadata` header such that references
/// to this type can be dereferenced into a `Thin` reference.
#[repr(C)]
struct KnownThin<U: ?Sized + Pointee, T> {
  header: MaybeUninit<ThinHeaderLayout<U, T>>,
  value: T,
}

impl<U: ?Sized + Pointee, T> KnownThin<U, T> {
  pub const fn new(value: T) -> Self
  where
    T: Unsize<U>,
  {
    let meta = std::ptr::metadata::<U>(&value);
    let mut thin = Self { header: MaybeUninit::uninit(), value };
    unsafe { *thin.metadata_mut() = meta };
    thin
  }
  const unsafe fn metadata_mut(&mut self) -> *mut U::Metadata {
    unsafe { (self as *mut _ as *mut MaybeUninit<ThinHeaderLayout<U, T>>).offset(1).cast::<U::Metadata>().offset(-1) }
  }
}

impl<U: ?Sized + Pointee, T> Deref for KnownThin<U, T> {
  type Target = Thin<U>;

  fn deref(&self) -> &Self::Target {
    unsafe { &*((self as *const _ as *const MaybeUninit<ThinHeaderLayout<U, T>>).offset(1) as *const _) }
  }
}

impl<U: ?Sized + Pointee, T> DerefMut for KnownThin<U, T> {
  fn deref_mut(&mut self) -> &mut Self::Target {
    unsafe { &mut *((self as *mut _ as *mut MaybeUninit<ThinHeaderLayout<U, T>>).offset(1) as *mut _) }
  }
}

impl<U: ?Sized + Pointee> Borrow<Thin<U>> for ThinBox<U> {
  fn borrow(&self) -> &Thin<U> {
    unsafe { std::mem::transmute_copy(self) }
  }
}

impl<U: ?Sized + Pointee> BorrowMut<Thin<U>> for ThinBox<U> {
  fn borrow_mut(&mut self) -> &mut Thin<U> {
    unsafe { std::mem::transmute_copy(self) }
  }
}

#[cfg(test)]
mod tests {
  use super::*;

  #[repr(align(256))]
  struct Big(u8);

  fn call_thin_fn_u8(f: &Thin<dyn Fn() -> u8>) -> u8 {
    f()
  }

  fn call_thin_fn_mut_u8(f: &mut Thin<dyn FnMut() -> u8>) -> u8 {
    f()
  }

  impl Big {
    fn inner(&self) -> u8 {
      self.0
    }
  }

  #[test]
  fn thin_ref() {
    assert_eq!(call_thin_fn_u8(&*KnownThin::new(|| 0)), 0);
    let y = 2u8;
    assert_eq!(call_thin_fn_u8(&*KnownThin::new(move || y)), 2);
  }

  #[test]
  fn thin_big() {
    assert_eq!(
      call_thin_fn_u8(&*{
        let x = Big(1);
        KnownThin::new(move || x.inner())
      }),
      1
    );
  }

  #[test]
  fn thin_mut() {
    let mut i = 2;
    let f = &mut *KnownThin::new(move || {
      i += 1;
      i
    });
    assert_eq!(call_thin_fn_mut_u8(f), 3);
    assert_eq!(call_thin_fn_mut_u8(f), 4);
    assert_eq!(call_thin_fn_mut_u8(f), 5);
  }

  #[test]
  fn thin_box() {
    let x = 123;
    let b = ThinBox::<dyn Fn() -> _>::new_unsize(move || x);
    assert_eq!(call_thin_fn_u8(b.borrow()), 123);
  }
}