| use crate::io::ReadBuf; |
| use std::mem::MaybeUninit; |
| |
| /// Something that looks like a `Vec<u8>`. |
| /// |
| /// # Safety |
| /// |
| /// The implementor must guarantee that the vector returned by the |
| /// `as_mut` and `as_mut` methods do not change from one call to |
| /// another. |
| pub(crate) unsafe trait VecU8: AsRef<Vec<u8>> + AsMut<Vec<u8>> {} |
| |
| unsafe impl VecU8 for Vec<u8> {} |
| unsafe impl VecU8 for &mut Vec<u8> {} |
| |
| /// This struct wraps a `Vec<u8>` or `&mut Vec<u8>`, combining it with a |
| /// `num_initialized`, which keeps track of the number of initialized bytes |
| /// in the unused capacity. |
| /// |
| /// The purpose of this struct is to remember how many bytes were initialized |
| /// through a `ReadBuf` from call to call. |
| /// |
| /// This struct has the safety invariant that the first `num_initialized` of the |
| /// vector's allocation must be initialized at any time. |
| #[derive(Debug)] |
| pub(crate) struct VecWithInitialized<V> { |
| vec: V, |
| // The number of initialized bytes in the vector. |
| // Always between `vec.len()` and `vec.capacity()`. |
| num_initialized: usize, |
| } |
| |
| impl VecWithInitialized<Vec<u8>> { |
| #[cfg(feature = "io-util")] |
| pub(crate) fn take(&mut self) -> Vec<u8> { |
| self.num_initialized = 0; |
| std::mem::take(&mut self.vec) |
| } |
| } |
| |
| impl<V> VecWithInitialized<V> |
| where |
| V: VecU8, |
| { |
| pub(crate) fn new(mut vec: V) -> Self { |
| // SAFETY: The safety invariants of vector guarantee that the bytes up |
| // to its length are initialized. |
| Self { |
| num_initialized: vec.as_mut().len(), |
| vec, |
| } |
| } |
| |
| pub(crate) fn reserve(&mut self, num_bytes: usize) { |
| let vec = self.vec.as_mut(); |
| if vec.capacity() - vec.len() >= num_bytes { |
| return; |
| } |
| // SAFETY: Setting num_initialized to `vec.len()` is correct as |
| // `reserve` does not change the length of the vector. |
| self.num_initialized = vec.len(); |
| vec.reserve(num_bytes); |
| } |
| |
| #[cfg(feature = "io-util")] |
| pub(crate) fn is_empty(&self) -> bool { |
| self.vec.as_ref().is_empty() |
| } |
| |
| pub(crate) fn get_read_buf<'a>(&'a mut self) -> ReadBuf<'a> { |
| let num_initialized = self.num_initialized; |
| |
| // SAFETY: Creating the slice is safe because of the safety invariants |
| // on Vec<u8>. The safety invariants of `ReadBuf` will further guarantee |
| // that no bytes in the slice are de-initialized. |
| let vec = self.vec.as_mut(); |
| let len = vec.len(); |
| let cap = vec.capacity(); |
| let ptr = vec.as_mut_ptr().cast::<MaybeUninit<u8>>(); |
| let slice = unsafe { std::slice::from_raw_parts_mut::<'a, MaybeUninit<u8>>(ptr, cap) }; |
| |
| // SAFETY: This is safe because the safety invariants of |
| // VecWithInitialized say that the first num_initialized bytes must be |
| // initialized. |
| let mut read_buf = ReadBuf::uninit(slice); |
| unsafe { |
| read_buf.assume_init(num_initialized); |
| } |
| read_buf.set_filled(len); |
| |
| read_buf |
| } |
| |
| pub(crate) fn apply_read_buf(&mut self, parts: ReadBufParts) { |
| let vec = self.vec.as_mut(); |
| assert_eq!(vec.as_ptr(), parts.ptr); |
| |
| // SAFETY: |
| // The ReadBufParts really does point inside `self.vec` due to the above |
| // check, and the safety invariants of `ReadBuf` guarantee that the |
| // first `parts.initialized` bytes of `self.vec` really have been |
| // initialized. Additionally, `ReadBuf` guarantees that `parts.len` is |
| // at most `parts.initialized`, so the first `parts.len` bytes are also |
| // initialized. |
| // |
| // Note that this relies on the fact that `V` is either `Vec<u8>` or |
| // `&mut Vec<u8>`, so the vector returned by `self.vec.as_mut()` cannot |
| // change from call to call. |
| unsafe { |
| self.num_initialized = parts.initialized; |
| vec.set_len(parts.len); |
| } |
| } |
| } |
| |
| pub(crate) struct ReadBufParts { |
| // Pointer is only used to check that the ReadBuf actually came from the |
| // right VecWithInitialized. |
| ptr: *const u8, |
| len: usize, |
| initialized: usize, |
| } |
| |
| // This is needed to release the borrow on `VecWithInitialized<V>`. |
| pub(crate) fn into_read_buf_parts(rb: ReadBuf<'_>) -> ReadBufParts { |
| ReadBufParts { |
| ptr: rb.filled().as_ptr(), |
| len: rb.filled().len(), |
| initialized: rb.initialized().len(), |
| } |
| } |