alloc: add ExternalBytes

TimLuq · Aug 3, 2024 · 388eca2 · 388eca2
1 parent 2f1e72c
commit 388eca2
Show file tree

Hide file tree

Showing 9 changed files with 481 additions and 32 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -25,12 +25,13 @@ default = ["macros", "chunk"]
 alloc = ["serde_1?/alloc"]
 chunk = []
 macros = []
-nightly = []
 queue = []
-core_io_borrowed_buf = []
 std = ["alloc"]
-bytes_1 = ["dep:bytes_1"]
 
+nightly = []
+core_io_borrowed_buf = []
+
+bytes_1 = ["dep:bytes_1"]
 bytes_1_safe = ["bytes_1"]
 http-body_04 = ["dep:http-body_04", "dep:http_02", "bytes_1"]
 http-body_1 = ["dep:http-body_1", "dep:http_1", "bytes_1"]

diff --git a/src/byte_chunk.rs b/src/byte_chunk.rs
@@ -31,7 +31,7 @@ impl ByteChunk {
 
     /// Create a `ByteChunk` from a fixed-size array.
     pub const fn from_array<const L: usize>(data: &[u8; L]) -> Self {
-        core::assert!(L <= 14, "chunk data too large");
+        core::assert!(L <= Self::LEN, "chunk data too large");
         let mut chunk = ByteChunk {
             len: L as u8,
             data: unsafe { core::mem::zeroed() },
@@ -156,9 +156,6 @@ impl ByteChunk {
 impl Default for ByteChunk {
     #[inline]
     fn default() -> Self {
-        ByteChunk {
-            len: 0,
-            data: [0; 14],
-        }
+        unsafe { core::mem::zeroed() }
     }
 }
diff --git a/src/bytedata.rs b/src/bytedata.rs
@@ -92,14 +92,24 @@ impl<T: Copy> DataKind<T> {
             return Kind::Slice;
         }
 
+        #[cfg(feature = "alloc")]
+        if base_kind == crate::external::KIND_EXT_BYTES {
+            return Kind::External;
+        }
+
         #[cfg(feature = "alloc")]
         if (base_kind & 0b0000_0011) == 0b0000_0000 {
             return Kind::Shared;
         }
+
         #[cfg(not(feature = "alloc"))]
         if (base_kind & 0b0000_0011) == 0b0000_0000 {
             panic!("alloc feature is not enabled, so no path should trigger this");
         }
+        #[cfg(not(feature = "alloc"))]
+        if (base_kind & 0b0000_0011) == 0b0000_0001 {
+            panic!("alloc feature is not enabled, so no path should trigger this");
+        }
 
         panic!("no path should trigger this");
     }
@@ -110,6 +120,8 @@ pub(crate) enum Kind {
     Slice,
     #[cfg(feature = "alloc")]
     Shared,
+    #[cfg(feature = "alloc")]
+    External,
 }
 
 const KIND_CHUNK_MASK: u8 = 0b0000_0111;
@@ -126,6 +138,10 @@ pub union ByteData<'a> {
     /// A shared byte slice.
     #[cfg_attr(docsrs, doc(cfg(feature = "alloc")))]
     pub(crate) shared: core::mem::ManuallyDrop<SharedBytes>,
+    #[cfg(feature = "alloc")]
+    /// A shared byte slice.
+    #[cfg_attr(docsrs, doc(cfg(feature = "alloc")))]
+    pub(crate) external: core::mem::ManuallyDrop<crate::external::ExtBytes>,
 }
 
 impl<'a> Clone for ByteData<'a> {
@@ -141,6 +157,12 @@ impl<'a> Clone for ByteData<'a> {
             Kind::Shared => Self {
                 shared: core::mem::ManuallyDrop::new(unsafe { SharedBytes::clone(&self.shared) }),
             },
+            #[cfg(feature = "alloc")]
+            Kind::External => Self {
+                external: core::mem::ManuallyDrop::new(unsafe {
+                    crate::external::ExtBytes::clone(&self.external)
+                }),
+            },
         }
     }
 }
@@ -161,6 +183,11 @@ impl<'a> Drop for ByteData<'a> {
                 core::ptr::drop_in_place(&mut self.shared);
                 self.chunk = empty_chunk();
             },
+            #[cfg(feature = "alloc")]
+            Kind::External => unsafe {
+                core::ptr::drop_in_place(&mut self.external);
+                self.chunk = empty_chunk();
+            },
         }
     }
 }
@@ -253,17 +280,23 @@ impl<'a> ByteData<'a> {
         if dat.is_empty() {
             return Self::empty();
         }
-        if dat.len() <= 14 {
+        if dat.len() <= crate::byte_chunk::ByteChunk::LEN {
             return Self {
                 chunk: WrappedChunk {
                     kind: KIND_CHUNK_MASK,
                     data: crate::byte_chunk::ByteChunk::from_slice(dat.as_slice()),
                 },
             };
         }
-        Self {
-            shared: core::mem::ManuallyDrop::new(dat.into()),
-        }
+        crate::external::ExtBytes::create(dat)
+    }
+
+    #[cfg(feature = "alloc")]
+    /// Creates a `ByteData` from an externally kept byte sequence.
+    #[cfg_attr(docsrs, doc(cfg(feature = "alloc")))]
+    #[inline]
+    pub fn from_external<E: crate::external::ExternalBytes>(dat: E) -> Self {
+        crate::external::ExtBytes::create(dat)
     }
 
     #[cfg(feature = "alloc")]
@@ -298,6 +331,14 @@ impl<'a> ByteData<'a> {
                 )
             }
             .as_slice(),
+            #[cfg(feature = "alloc")]
+            Kind::External => unsafe {
+                core::mem::transmute::<
+                    &core::mem::ManuallyDrop<crate::external::ExtBytes>,
+                    &crate::external::ExtBytes,
+                >(&self.external)
+            }
+            .as_slice(),
         }
     }
 
@@ -313,6 +354,14 @@ impl<'a> ByteData<'a> {
                 )
             }
             .len(),
+            #[cfg(feature = "alloc")]
+            Kind::External => unsafe {
+                core::mem::transmute::<
+                    &core::mem::ManuallyDrop<crate::external::ExtBytes>,
+                    &crate::external::ExtBytes,
+                >(&self.external)
+            }
+            .len(),
         }
     }
 
@@ -328,6 +377,8 @@ impl<'a> ByteData<'a> {
                 )
             }
             .is_empty(),
+            #[cfg(feature = "alloc")]
+            Kind::External => false,
         }
     }
 
@@ -376,12 +427,22 @@ impl<'a> ByteData<'a> {
                     )
                 };
                 let dat = dat.sliced_range(range);
-                if dat.len() <= 14 {
+                if dat.len() <= crate::byte_chunk::ByteChunk::LEN {
                     Self::from_chunk_slice(dat.as_slice())
                 } else {
                     Self::from_shared(dat)
                 }
             }
+            #[cfg(feature = "alloc")]
+            Kind::External => {
+                let dat = unsafe {
+                    core::mem::transmute::<
+                        &core::mem::ManuallyDrop<crate::external::ExtBytes>,
+                        &crate::external::ExtBytes,
+                    >(&self.external)
+                };
+                dat.sliced_range(range)
+            }
         }
     }
 
@@ -409,7 +470,25 @@ impl<'a> ByteData<'a> {
                     >(&mut self.shared)
                 };
                 dat.make_sliced_range(range);
-                if dat.len() <= 14 {
+                if dat.len() <= crate::byte_chunk::ByteChunk::LEN {
+                    let r = Self::from_chunk_slice(dat.as_slice());
+                    unsafe { core::ptr::drop_in_place(dat) };
+                    core::mem::forget(self);
+                    r
+                } else {
+                    self
+                }
+            }
+            #[cfg(feature = "alloc")]
+            Kind::External => {
+                let dat = unsafe {
+                    core::mem::transmute::<
+                        &mut core::mem::ManuallyDrop<crate::external::ExtBytes>,
+                        &mut crate::external::ExtBytes,
+                    >(&mut self.external)
+                };
+                dat.make_sliced_range(range);
+                if dat.len() <= crate::byte_chunk::ByteChunk::LEN {
                     let r = Self::from_chunk_slice(dat.as_slice());
                     unsafe { core::ptr::drop_in_place(dat) };
                     core::mem::forget(self);
@@ -443,7 +522,27 @@ impl<'a> ByteData<'a> {
                     >(&mut self.shared)
                 };
                 dat.make_sliced_range(range);
-                if dat.len() <= 14 {
+                if dat.len() <= crate::byte_chunk::ByteChunk::LEN {
+                    let r = crate::ByteChunk::from_slice(dat.as_slice());
+                    unsafe {
+                        core::ptr::drop_in_place(dat);
+                        self.chunk = DataKind {
+                            kind: KIND_CHUNK_MASK,
+                            data: r,
+                        };
+                    }
+                }
+            }
+            #[cfg(feature = "alloc")]
+            Kind::External => {
+                let dat = unsafe {
+                    core::mem::transmute::<
+                        &mut core::mem::ManuallyDrop<crate::external::ExtBytes>,
+                        &mut crate::external::ExtBytes,
+                    >(&mut self.external)
+                };
+                dat.make_sliced_range(range);
+                if dat.len() <= crate::byte_chunk::ByteChunk::LEN {
                     let r = crate::ByteChunk::from_slice(dat.as_slice());
                     unsafe {
                         core::ptr::drop_in_place(dat);
@@ -462,14 +561,14 @@ impl<'a> ByteData<'a> {
     #[cfg_attr(docsrs, doc(cfg(feature = "alloc")))]
     pub fn into_shared<'s>(mut self) -> ByteData<'s> {
         match unsafe { self.chunk.kind() } {
-            Kind::Chunk | Kind::Shared => unsafe {
+            Kind::Chunk | Kind::Shared | Kind::External => unsafe {
                 core::mem::transmute::<ByteData, ByteData>(self)
             },
             Kind::Slice => {
                 let a = unsafe { &self.slice };
                 if a.is_static() {
                     unsafe { core::mem::transmute::<ByteData, ByteData>(self) }
-                } else if a.len() <= 14 {
+                } else if a.len() <= crate::byte_chunk::ByteChunk::LEN {
                     let r = crate::byte_chunk::ByteChunk::from_slice(a.as_slice());
                     core::mem::forget(self);
                     ByteData {
@@ -505,10 +604,14 @@ impl<'a> ByteData<'a> {
                 (*self.shared).make_sliced_range(range);
                 core::mem::transmute::<ByteData, ByteData>(self)
             },
+            Kind::External => unsafe {
+                (*self.external).make_sliced_range(range);
+                core::mem::transmute::<ByteData, ByteData>(self)
+            },
             Kind::Slice => {
                 let a = unsafe { &self.slice };
                 let r = &a.as_slice()[range];
-                if r.len() <= 14 {
+                if r.len() <= crate::byte_chunk::ByteChunk::LEN {
                     let r = crate::byte_chunk::ByteChunk::from_slice(r);
                     core::mem::forget(self);
                     return ByteData {
@@ -664,7 +767,7 @@ impl<'a> From<&'a [u8]> for ByteData<'a> {
 impl<'a> From<SharedBytes> for ByteData<'a> {
     #[inline]
     fn from(dat: SharedBytes) -> Self {
-        if dat.len() <= 14 {
+        if dat.len() <= crate::byte_chunk::ByteChunk::LEN {
             Self::from_chunk_slice(&dat)
         } else {
             Self::from_shared(dat)

diff --git a/src/bytes_1/bytedata.rs b/src/bytes_1/bytedata.rs
@@ -36,6 +36,22 @@ impl From<ByteData<'_>> for bytes::Bytes {
                 let s = unsafe { core::mem::transmute::<ByteData, crate::SharedBytes>(dat) };
                 s.into()
             }
+            #[cfg(feature = "alloc")]
+            crate::bytedata::Kind::External => {
+                let s = unsafe { core::mem::transmute::<ByteData, crate::external::ExtBytes>(dat) };
+                let r = s.with_inner(|t: &bytes::Bytes, s: &[u8]| {
+                    let start = {
+                        let sp: *const u8 = s.as_ptr();
+                        let tp: *const u8 = t.as_ptr();
+                        sp as usize - tp as usize
+                    };
+                    t.slice(start..(start + s.len()))
+                });
+                match r {
+                    Some(r) => r,
+                    None => bytes::Bytes::copy_from_slice(s.as_slice()),
+                }
+            }
         }
     }
 }

diff --git a/src/bytes_1/external.rs b/src/bytes_1/external.rs
@@ -0,0 +1,4 @@
+impl crate::external::ExternalBytes for ::bytes_1::Bytes {
+    const OPS: crate::external::ExternalOps<Self> =
+        crate::external::ExternalOps::new(::bytes_1::Bytes::as_ref);
+}
diff --git a/src/bytes_1/mod.rs b/src/bytes_1/mod.rs
@@ -6,6 +6,8 @@ use ::bytes_1 as bytes;
 
 mod bytedata;
 
+#[cfg(feature = "alloc")]
+mod external;
 #[cfg(feature = "queue")]
 mod queue;
 #[cfg(feature = "alloc")]