Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
io: keep track of initialized bytes in read_to_end (#3426)
- Loading branch information
Showing
5 changed files
with
227 additions
and
48 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,120 @@ | ||
| use crate::io::ReadBuf; | ||
| use std::mem::MaybeUninit; | ||
|
|
||
| /// 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: AsMut<Vec<u8>>, | ||
| { | ||
| /// Safety: The generic parameter `V` must be either `Vec<u8>` or `&mut Vec<u8>`. | ||
| pub(crate) unsafe 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(&mut self) -> bool { | ||
| self.vec.as_mut().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(), | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters