Add free conversion from Bytes back to Arc if possible (#217)

getditto · May 30, 2024 · 8b7a8e0 · 8b7a8e0
1 parent da13fcc
commit 8b7a8e0
Showing 1 changed file with 33 additions and 17 deletions.
diff --git a/src/bytes.rs b/src/bytes.rs
@@ -224,31 +224,28 @@ impl<'a> From<&'a [u8]> for Bytes<'a> {
     }
 }
 #[cfg(feature = "alloc")]
+unsafe extern "C" fn retain_arc_bytes(this: *const (), capacity: usize) {
+    Arc::increment_strong_count(core::ptr::slice_from_raw_parts(this.cast::<u8>(), capacity))
+}
+#[cfg(feature = "alloc")]
+unsafe extern "C" fn release_arc_bytes(this: *const (), capacity: usize) {
+    Arc::decrement_strong_count(core::ptr::slice_from_raw_parts(this.cast::<u8>(), capacity))
+}
+#[cfg(feature = "alloc")]
+static ARC_BYTES_VT: BytesVt = BytesVt {
+    release: Some(release_arc_bytes),
+    retain: Some(retain_arc_bytes),
+};
+#[cfg(feature = "alloc")]
 impl From<Arc<[u8]>> for Bytes<'static> {
     fn from(data: Arc<[u8]>) -> Self {
-        unsafe extern "C" fn retain(this: *const (), capacity: usize) {
-            Arc::increment_strong_count(core::ptr::slice_from_raw_parts(
-                this.cast::<u8>(),
-                capacity,
-            ))
-        }
-        unsafe extern "C" fn release(this: *const (), capacity: usize) {
-            Arc::decrement_strong_count(core::ptr::slice_from_raw_parts(
-                this.cast::<u8>(),
-                capacity,
-            ))
-        }
-        static VT: BytesVt = BytesVt {
-            release: Some(release),
-            retain: Some(retain),
-        };
         let capacity = data.len();
         Bytes {
             start: core::ptr::NonNull::<[u8]>::from(data.as_ref()).cast(),
             len: data.len(),
             data: Arc::into_raw(data) as *const (),
             capacity,
-            vtable: &VT,
+            vtable: &ARC_BYTES_VT,
         }
     }
 }
@@ -356,6 +353,25 @@ impl Drop for Bytes<'_> {
     }
 }
 
+#[cfg(feature = "alloc")]
+impl<'a> TryFrom<Bytes<'a>> for Arc<[u8]> {
+    type Error = Bytes<'a>;
+    fn try_from(value: Bytes<'a>) -> Result<Self, Self::Error> {
+        let data = value.data.cast();
+        match core::ptr::eq(value.vtable, &ARC_BYTES_VT)
+            && core::ptr::eq(value.start.as_ptr(), data)
+            && value.len == value.capacity
+        {
+            true => unsafe {
+                let arc = Arc::from_raw(core::ptr::slice_from_raw_parts(data, value.capacity));
+                core::mem::forget(value);
+                Ok(arc)
+            },
+            false => Err(value),
+        }
+    }
+}
+
 #[cfg(feature = "alloc")]
 #[test]
 fn fuzz() {