Skip to main content

tl/
bytes.rs

1#[cfg(feature = "std")]
2use core::mem::ManuallyDrop;
3use core::{
4    fmt,
5    fmt::Debug,
6    hash::{Hash, Hasher},
7    marker::PhantomData,
8};
9#[cfg(feature = "std")]
10use std::borrow::Cow;
11
12#[cfg(feature = "std")]
13use crate::errors::SetBytesError;
14
15/// A storage type for raw bytes, used by the parser
16#[derive(Eq, PartialOrd, Ord)]
17pub struct Bytes<'a> {
18    /// The inner data
19    data: BytesInner,
20    /// Enforce the lifetime of the referenced data
21    _lt: PhantomData<&'a [u8]>,
22}
23
24/// The inner data of [`Bytes`]
25///
26/// Instead of using `&[u8]` and `Vec<u8>` for the variants,
27/// we use raw pointers and a `u32` for the length.
28/// This is to keep the size of the enum to 16 (on 64-bit machines),
29/// which is the same as if this was just `struct Bytes<'a>(&'a [u8])`
30#[derive(PartialEq, Eq, PartialOrd, Ord)]
31enum BytesInner {
32    /// Borrowed bytes
33    Borrowed(*const u8, u32),
34    /// Owned bytes
35    ///
36    /// This pointer is managed and will be freed when dropped
37    #[cfg(feature = "std")]
38    Owned(*mut u8, u32),
39}
40
41impl<'a> PartialEq<str> for Bytes<'a> {
42    #[inline]
43    fn eq(&self, other: &str) -> bool {
44        self == other.as_bytes()
45    }
46}
47
48impl<'a> PartialEq<[u8]> for Bytes<'a> {
49    #[inline]
50    fn eq(&self, other: &[u8]) -> bool {
51        self.as_bytes() == other
52    }
53}
54
55impl<'a> PartialEq for Bytes<'a> {
56    #[inline]
57    fn eq(&self, other: &Self) -> bool {
58        let this = self.as_bytes();
59        let that = other.as_bytes();
60        this == that
61    }
62}
63
64impl<'a> Hash for Bytes<'a> {
65    #[inline]
66    fn hash<H: Hasher>(&self, state: &mut H) {
67        // Hash must be implemented manually for Bytes, otherwise it would only hash a pointer
68        let this = self.as_bytes();
69        this.hash(state);
70    }
71}
72
73impl<'a> Clone for Bytes<'a> {
74    fn clone(&self) -> Self {
75        // It is important to manually implement Clone for Bytes,
76        // because if `self` was owned, then the default clone
77        // implementation would only clone the pointer
78        // which leads to aliasing boxes, and later, when `Bytes` is dropped,
79        // the box is freed twice!
80        match &self.data {
81            BytesInner::Borrowed(data, len) => {
82                Bytes::from(unsafe { compact_bytes_to_slice(*data, *len) })
83            }
84            #[cfg(feature = "std")]
85            BytesInner::Owned(data, len) => {
86                let (ptr, len) = unsafe { clone_compact_bytes_parts(*data, *len) };
87                Bytes {
88                    data: BytesInner::Owned(ptr, len),
89                    _lt: PhantomData,
90                }
91            }
92        }
93    }
94}
95
96impl<'a> From<&'a str> for Bytes<'a> {
97    #[inline]
98    fn from(s: &'a str) -> Self {
99        <Self as From<&'a [u8]>>::from(s.as_bytes())
100    }
101}
102
103impl<'a> From<&'a [u8]> for Bytes<'a> {
104    #[inline]
105    fn from(s: &'a [u8]) -> Self {
106        Bytes {
107            data: BytesInner::Borrowed(s.as_ptr(), s.len() as u32),
108            _lt: PhantomData,
109        }
110    }
111}
112
113#[cfg(feature = "std")]
114impl TryFrom<String> for Bytes<'static> {
115    type Error = SetBytesError;
116
117    #[inline]
118    fn try_from(s: String) -> Result<Self, Self::Error> {
119        let mut bytes = Bytes::new();
120        bytes.set(s)?;
121        Ok(bytes)
122    }
123}
124
125/// Converts `Bytes` raw parts to a slice
126#[inline]
127unsafe fn compact_bytes_to_slice<'a>(ptr: *const u8, l: u32) -> &'a [u8] {
128    unsafe { core::slice::from_raw_parts(ptr, l as usize) }
129}
130
131/// Converts a boxed byte slice to compact raw parts
132///
133/// The caller is responsible for freeing the returned pointer and that the length of the slice does not overflow a u32!
134#[cfg(feature = "std")]
135unsafe fn boxed_slice_into_compact_parts(slice: Box<[u8]>) -> (*mut u8, u32) {
136    // wrap box in `ManuallyDrop` so it's not dropped at the end of the scope
137    let mut slice = ManuallyDrop::new(slice);
138    let len = slice.len();
139    let ptr = slice.as_mut_ptr();
140
141    (ptr, len as u32)
142}
143
144/// Clones a slice given its raw parts and returns the new, cloned parts
145#[inline]
146#[cfg(feature = "std")]
147unsafe fn clone_compact_bytes_parts(ptr: *mut u8, len: u32) -> (*mut u8, u32) {
148    let slice = unsafe { compact_bytes_to_slice(ptr, len) }
149        .to_vec()
150        .into_boxed_slice();
151    unsafe { boxed_slice_into_compact_parts(slice) }
152}
153
154// Custom `Debug` trait is implemented which displays the data as a UTF8 string,
155// to make it easier to read for humans when logging
156impl<'a> Debug for Bytes<'a> {
157    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
158        f.debug_tuple("Bytes")
159            .field(&self.try_as_utf8_str().unwrap_or("<non-utf8>"))
160            .finish()
161    }
162}
163
164impl<'a> Default for Bytes<'a> {
165    fn default() -> Self {
166        Self::new()
167    }
168}
169
170impl<'a> Bytes<'a> {
171    /// Creates an empty `Bytes`
172    #[inline]
173    pub fn new() -> Self {
174        Self {
175            data: BytesInner::Borrowed("".as_bytes().as_ptr(), 0),
176            _lt: PhantomData,
177        }
178    }
179
180    /// Convenient method for lossy-encoding the data as UTF8
181    #[inline]
182    #[cfg(feature = "std")]
183    pub fn as_utf8_str(&self) -> Cow<'_, str> {
184        String::from_utf8_lossy(self.as_bytes())
185    }
186
187    /// Tries to convert the inner data to a `&str`, without allocating in the case
188    /// that the inner data is not valid UTF8
189    #[inline]
190    pub fn try_as_utf8_str(&self) -> Option<&str> {
191        core::str::from_utf8(self.as_bytes()).ok()
192    }
193
194    /// Returns the raw data wrapped by this struct
195    #[inline]
196    pub fn as_bytes(&self) -> &[u8] {
197        match &self.data {
198            BytesInner::Borrowed(b, l) => unsafe { compact_bytes_to_slice(*b, *l) },
199            #[cfg(feature = "std")]
200            BytesInner::Owned(o, l) => unsafe { compact_bytes_to_slice(*o, *l) },
201        }
202    }
203
204    /// Returns the raw data referenced by this struct
205    ///
206    /// The lifetime of the returned data is tied to 'a, unlike `Bytes::as_bytes`
207    /// which has a lifetime of '_ (self) in case it is owned
208    #[inline]
209    pub fn as_bytes_borrowed(&self) -> Option<&'a [u8]> {
210        match &self.data {
211            BytesInner::Borrowed(b, l) => Some(unsafe { compact_bytes_to_slice(*b, *l) }),
212            #[cfg(feature = "std")]
213            _ => None,
214        }
215    }
216
217    /// Returns a read-only raw pointer to the inner data
218    #[inline]
219    pub fn as_ptr(&self) -> *const u8 {
220        match &self.data {
221            BytesInner::Borrowed(b, _) => *b,
222            #[cfg(feature = "std")]
223            BytesInner::Owned(o, _) => *o,
224        }
225    }
226
227    /// Sets the inner data to the given data and returns the old bytes
228    #[cfg(feature = "std")]
229    pub fn set<B: IntoOwnedBytes>(&mut self, data: B) -> Result<Option<Box<[u8]>>, SetBytesError> {
230        const MAX: usize = u32::MAX as usize;
231
232        let data = <B as IntoOwnedBytes>::into_bytes(data);
233
234        if data.len() > MAX {
235            return Err(SetBytesError::LengthOverflow);
236        }
237
238        // SAFETY: All invariants are checked
239        Ok(unsafe { self.set_unchecked(data) })
240    }
241
242    /// Sets the inner data to the given data without checking for validity of the data
243    ///
244    /// ## Safety
245    /// - Once `data` is converted to a `Box<[u8]>`, its length must not be greater than u32::MAX
246    #[inline]
247    #[cfg(feature = "std")]
248    pub unsafe fn set_unchecked<B: IntoOwnedBytes>(&mut self, data: B) -> Option<Box<[u8]>> {
249        let data = <B as IntoOwnedBytes>::into_bytes(data);
250
251        let (ptr, len) = unsafe { boxed_slice_into_compact_parts(data) };
252
253        let bytes = BytesInner::Owned(ptr, len);
254        let old = core::mem::replace(&mut self.data, bytes);
255
256        // we cannot let Drop code run because that would deallocate `old`
257        let old = ManuallyDrop::new(old);
258
259        match &*old {
260            BytesInner::Borrowed(_, _) => None,
261            BytesInner::Owned(ptr, len) => {
262                let len = *len as usize;
263                Some(unsafe { Vec::from_raw_parts(*ptr, len, len) }.into_boxed_slice())
264            }
265        }
266    }
267}
268
269#[cfg(feature = "std")]
270mod private {
271    pub trait Sealed {}
272}
273
274/// A trait implemented on types that can be used for `Bytes::set`.
275///
276/// This trait is sealed and cannot be implemented outside of this crate.
277#[cfg(feature = "std")]
278pub trait IntoOwnedBytes: private::Sealed {
279    fn into_bytes(self) -> Box<[u8]>;
280}
281
282#[cfg(feature = "std")]
283macro_rules! impl_into_owned_bytes_trivial {
284    ($($t:ty),*) => {
285        $(
286            impl private::Sealed for $t {}
287            impl IntoOwnedBytes for $t {
288                #[inline]
289                fn into_bytes(self) -> Box<[u8]> {
290                    self.into()
291                }
292            }
293        )*
294    };
295}
296
297#[cfg(feature = "std")]
298impl_into_owned_bytes_trivial!(Box<[u8]>, &[u8], Vec<u8>);
299
300#[cfg(feature = "std")]
301impl private::Sealed for &str {}
302#[cfg(feature = "std")]
303impl IntoOwnedBytes for &str {
304    #[inline]
305    fn into_bytes(self) -> Box<[u8]> {
306        self.as_bytes().into()
307    }
308}
309
310#[cfg(feature = "std")]
311impl private::Sealed for String {}
312#[cfg(feature = "std")]
313impl IntoOwnedBytes for String {
314    #[inline]
315    fn into_bytes(self) -> Box<[u8]> {
316        self.into_bytes().into()
317    }
318}
319
320#[cfg(feature = "std")]
321impl Drop for BytesInner {
322    fn drop(&mut self) {
323        // we only need to deallocate if we own the data
324        // if we don't, just do nothing
325        if let BytesInner::Owned(ptr, len) = self {
326            let ptr = *ptr;
327            let len = *len as usize;
328
329            // carefully reconstruct a `Box<[u8]>` from the raw pointer and length
330            // and immediately drop it to free memory
331            unsafe { drop(Vec::from_raw_parts(ptr, len, len).into_boxed_slice()) };
332        }
333    }
334}