From 0fe2b2c5688533e11b0a5e139f56ac3b17fb5d4c Mon Sep 17 00:00:00 2001
From: Waffle Maybe <waffle.lapkin@gmail.com>
Date: Sun, 24 Jul 2022 07:09:32 +0400
Subject: [PATCH] Add `{wrap,peel}_box` functions to `TransparentWrapperAlloc`
 (#119)

* Add `{wrap,peel}_box` functions to `TransparentWrapperAlloc`

* Add tests for `{wrap,peel}_box`
---
 src/allocation.rs    | 48 ++++++++++++++++++++++++++++++++++++++++++++
 tests/transparent.rs | 27 +++++++++++++++++--------
 2 files changed, 67 insertions(+), 8 deletions(-)

diff --git a/src/allocation.rs b/src/allocation.rs
index 735f6db..888d3c1 100644
--- a/src/allocation.rs
+++ b/src/allocation.rs
@@ -339,6 +339,30 @@ pub trait TransparentWrapperAlloc<Inner: ?Sized>: TransparentWrapper<Inner> {
     }
   }
 
+  /// Convert a box to the inner type into a box to the wrapper
+  /// type.
+  #[inline]
+  fn wrap_box(s: Box<Inner>) -> Box<Self> {
+    assert!(size_of::<*mut Inner>() == size_of::<*mut Self>());
+
+    unsafe {
+      // A pointer cast doesn't work here because rustc can't tell that
+      // the vtables match (because of the `?Sized` restriction relaxation).
+      // A `transmute` doesn't work because the sizes are unspecified.
+      //
+      // SAFETY: 
+      // * The unsafe contract requires that pointers to Inner and Self
+      //   have identical representations 
+      // * Box is guaranteed to have representation identical to a 
+      //   (non-null) pointer
+      // * The pointer comes from a box (and thus satisfies all safety 
+      //   requirements of Box)
+      let inner_ptr: *mut Inner = Box::into_raw(s);
+      let wrapper_ptr: *mut Self = transmute!(inner_ptr);
+      Box::from_raw(wrapper_ptr)
+    }
+  }
+
   /// Convert a vec of the wrapper type into a vec of the inner type.
   fn peel_vec(s: Vec<Self>) -> Vec<Inner>
   where
@@ -363,5 +387,29 @@ pub trait TransparentWrapperAlloc<Inner: ?Sized>: TransparentWrapper<Inner> {
       )
     }
   }
+
+  /// Convert a box to the wrapper type into a box to the inner
+  /// type.
+  #[inline]
+  fn peel_box(s: Box<Self>) -> Box<Inner> {
+    assert!(size_of::<*mut Inner>() == size_of::<*mut Self>());
+
+    unsafe {
+      // A pointer cast doesn't work here because rustc can't tell that
+      // the vtables match (because of the `?Sized` restriction relaxation).
+      // A `transmute` doesn't work because the sizes are unspecified.
+      //
+      // SAFETY: 
+      // * The unsafe contract requires that pointers to Inner and Self
+      //   have identical representations 
+      // * Box is guaranteed to have representation identical to a 
+      //   (non-null) pointer
+      // * The pointer comes from a box (and thus satisfies all safety 
+      //   requirements of Box)
+      let wrapper_ptr: *mut Self = Box::into_raw(s);
+      let inner_ptr: *mut Inner = transmute!(wrapper_ptr);
+      Box::from_raw(inner_ptr)
+    }
+  }
 }
 impl<I: ?Sized, T: TransparentWrapper<I>> TransparentWrapperAlloc<I> for T {}
diff --git a/tests/transparent.rs b/tests/transparent.rs
index fc4ec1b..ad3435f 100644
--- a/tests/transparent.rs
+++ b/tests/transparent.rs
@@ -5,12 +5,12 @@
 #[test]
 fn test_transparent_wrapper() {
   // An external type defined in a different crate.
-  #[derive(Copy, Clone, Default)]
+  #[derive(Debug, Copy, Clone, Default)]
   struct Foreign(u8);
 
   use bytemuck::TransparentWrapper;
 
-  #[derive(Copy, Clone)]
+  #[derive(Debug, Copy, Clone)]
   #[repr(transparent)]
   struct Wrapper(Foreign);
 
@@ -20,6 +20,14 @@ fn test_transparent_wrapper() {
   unsafe impl bytemuck::Zeroable for Wrapper {}
   unsafe impl bytemuck::Pod for Wrapper {}
 
+  impl PartialEq<u8> for Foreign {
+    fn eq(&self, &other: &u8) -> bool { self.0 == other }
+  }
+  
+  impl PartialEq<u8> for Wrapper {
+    fn eq(&self, &other: &u8) -> bool { self.0 == other }
+  }
+
   let _: u8 = bytemuck::cast(Wrapper::wrap(Foreign::default()));
   let _: Foreign = Wrapper::peel(bytemuck::cast(u8::default()));
 
@@ -69,13 +77,16 @@ fn test_transparent_wrapper() {
     let a: Vec<Foreign> = vec![Foreign::default(); 2];
     
     let b: Vec<Wrapper> = Wrapper::wrap_vec(a);
-    assert_eq!(
-      bytemuck::cast_slice::<Wrapper, u8>(b.as_slice()), &[0u8, 0]
-    );
+    assert_eq!(b, [0, 0]);
 
     let c: Vec<Foreign> = Wrapper::peel_vec(b);
-    assert_eq!(
-      bytemuck::cast_slice::<Wrapper, u8>(Wrapper::wrap_slice(c.as_slice())), &[0, 0]
-    );
+    assert_eq!(c, [0, 0]);
+
+    let d: Box<Foreign> = Box::new(Foreign::default());
+    
+    let e: Box<Wrapper> = Wrapper::wrap_box(d);
+    assert_eq!(&*e, &0);
+    let f: Box<Foreign> = Wrapper::peel_box(e);
+    assert_eq!(&*f, &0);
   }
 }