From abe55e525ee8b4f17dfd05ba31732b6a3580fe07 Mon Sep 17 00:00:00 2001
From: Edward <theedward162@gmail.com>
Date: Wed, 29 Jun 2022 01:55:31 +0200
Subject: [PATCH] feat: add TransparentWrapperAlloc trait (#111)

* feat: add TransparentWrapperAlloc trait

* implements wrap_vec and peel_vec

* fix: review comments

* remote asserts as they don't do anything
* avoid issue with stacked borrows
---
 src/allocation.rs    | 54 ++++++++++++++++++++++++++++++++++++++++++++
 tests/transparent.rs | 17 ++++++++++++++
 2 files changed, 71 insertions(+)

diff --git a/src/allocation.rs b/src/allocation.rs
index 72f3f16..50be3b3 100644
--- a/src/allocation.rs
+++ b/src/allocation.rs
@@ -297,3 +297,57 @@ pub fn pod_collect_to_vec<
   dst_bytes[..src_size].copy_from_slice(src_bytes);
   dst
 }
+
+/// An extension trait for `TransparentWrapper` and alloc types.
+pub trait TransparentWrapperAlloc<Inner: ?Sized>: TransparentWrapper<Inner> {
+  /// Convert a vec of the inner type into a vec of the wrapper type.
+  fn wrap_vec(s: Vec<Inner>) -> Vec<Self>
+  where
+    Self: Sized,
+    Inner: Sized
+  {
+    let mut s = core::mem::ManuallyDrop::new(s);
+
+    let length = s.len();
+    let capacity = s.capacity();
+    let ptr = s.as_mut_ptr();
+
+    unsafe {
+      // SAFETY:
+      // * ptr comes from Vec (and will not be double-dropped)
+      // * the two types have the identical representation
+      // * the len and capacity fields are valid
+      Vec::from_raw_parts(
+        ptr as *mut Self,
+        length,
+        capacity
+      )
+    }
+  }
+
+  /// Convert a vec of the wrapper type into a vec of the inner type.
+  fn peel_vec(s: Vec<Self>) -> Vec<Inner>
+  where
+    Self: Sized,
+    Inner: Sized
+  {
+    let mut s = core::mem::ManuallyDrop::new(s);
+
+    let length = s.len();
+    let capacity = s.capacity();
+    let ptr = s.as_mut_ptr();
+
+    unsafe {
+      // SAFETY:
+      // * ptr comes from Vec (and will not be double-dropped)
+      // * the two types have the identical representation
+      // * the len and capacity fields are valid
+      Vec::from_raw_parts(
+        ptr as *mut Inner,
+        length,
+        capacity
+      )
+    }
+  }
+}
+impl<I: ?Sized, T: TransparentWrapper<I>> TransparentWrapperAlloc<I> for T {}
diff --git a/tests/transparent.rs b/tests/transparent.rs
index b1ce847..fc4ec1b 100644
--- a/tests/transparent.rs
+++ b/tests/transparent.rs
@@ -61,4 +61,21 @@ fn test_transparent_wrapper() {
     Foreign::default(),
   ]));
   // counterpart?
+
+  #[cfg(feature = "extern_crate_alloc")]
+  {
+    use bytemuck::allocation::TransparentWrapperAlloc;
+
+    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]
+    );
+
+    let c: Vec<Foreign> = Wrapper::peel_vec(b);
+    assert_eq!(
+      bytemuck::cast_slice::<Wrapper, u8>(Wrapper::wrap_slice(c.as_slice())), &[0, 0]
+    );
+  }
 }