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