1use crate::io::ReadBuf;
2use std::mem::MaybeUninit;
3
4/// Something that looks like a `Vec<u8>`.
5///
6/// # Safety
7///
8/// The implementor must guarantee that the vector returned by the
9/// `as_mut` and `as_mut` methods do not change from one call to
10/// another.
11pub(crate) unsafe trait VecU8: AsRef<Vec<u8>> + AsMut<Vec<u8>> {}
12
13unsafe impl VecU8 for Vec<u8> {}
14unsafe impl VecU8 for &mut Vec<u8> {}
15
16/// This struct wraps a `Vec<u8>` or `&mut Vec<u8>`, combining it with a
17/// `num_initialized`, which keeps track of the number of initialized bytes
18/// in the unused capacity.
19///
20/// The purpose of this struct is to remember how many bytes were initialized
21/// through a `ReadBuf` from call to call.
22///
23/// This struct has the safety invariant that the first `num_initialized` of the
24/// vector's allocation must be initialized at any time.
25#[derive(Debug)]
26pub(crate) struct VecWithInitialized<V> {
27 vec: V,
28 // The number of initialized bytes in the vector.
29 // Always between `vec.len()` and `vec.capacity()`.
30 num_initialized: usize,
31 starting_capacity: usize,
32}
33
34impl VecWithInitialized<Vec<u8>> {
35 #[cfg(feature = "io-util")]
36 pub(crate) fn take(&mut self) -> Vec<u8> {
37 self.num_initialized = 0;
38 std::mem::take(&mut self.vec)
39 }
40}
41
42impl<V> VecWithInitialized<V>
43where
44 V: VecU8,
45{
46 pub(crate) fn new(mut vec: V) -> Self {
47 // SAFETY: The safety invariants of vector guarantee that the bytes up
48 // to its length are initialized.
49 Self {
50 num_initialized: vec.as_mut().len(),
51 starting_capacity: vec.as_ref().capacity(),
52 vec,
53 }
54 }
55
56 pub(crate) fn reserve(&mut self, num_bytes: usize) {
57 let vec = self.vec.as_mut();
58 if vec.capacity() - vec.len() >= num_bytes {
59 return;
60 }
61 // SAFETY: Setting num_initialized to `vec.len()` is correct as
62 // `reserve` does not change the length of the vector.
63 self.num_initialized = vec.len();
64 vec.reserve(num_bytes);
65 }
66
67 #[cfg(feature = "io-util")]
68 pub(crate) fn is_empty(&self) -> bool {
69 self.vec.as_ref().is_empty()
70 }
71
72 pub(crate) fn get_read_buf<'a>(&'a mut self) -> ReadBuf<'a> {
73 let num_initialized = self.num_initialized;
74
75 // SAFETY: Creating the slice is safe because of the safety invariants
76 // on Vec<u8>. The safety invariants of `ReadBuf` will further guarantee
77 // that no bytes in the slice are de-initialized.
78 let vec = self.vec.as_mut();
79 let len = vec.len();
80 let cap = vec.capacity();
81 let ptr = vec.as_mut_ptr().cast::<MaybeUninit<u8>>();
82 let slice = unsafe { std::slice::from_raw_parts_mut::<'a, MaybeUninit<u8>>(ptr, cap) };
83
84 // SAFETY: This is safe because the safety invariants of
85 // VecWithInitialized say that the first num_initialized bytes must be
86 // initialized.
87 let mut read_buf = ReadBuf::uninit(slice);
88 unsafe {
89 read_buf.assume_init(num_initialized);
90 }
91 read_buf.set_filled(len);
92
93 read_buf
94 }
95
96 pub(crate) fn apply_read_buf(&mut self, parts: ReadBufParts) {
97 let vec = self.vec.as_mut();
98 assert_eq!(vec.as_ptr(), parts.ptr);
99
100 // SAFETY:
101 // The ReadBufParts really does point inside `self.vec` due to the above
102 // check, and the safety invariants of `ReadBuf` guarantee that the
103 // first `parts.initialized` bytes of `self.vec` really have been
104 // initialized. Additionally, `ReadBuf` guarantees that `parts.len` is
105 // at most `parts.initialized`, so the first `parts.len` bytes are also
106 // initialized.
107 //
108 // Note that this relies on the fact that `V` is either `Vec<u8>` or
109 // `&mut Vec<u8>`, so the vector returned by `self.vec.as_mut()` cannot
110 // change from call to call.
111 unsafe {
112 self.num_initialized = parts.initialized;
113 vec.set_len(parts.len);
114 }
115 }
116
117 // Returns a boolean telling the caller to try reading into a small local buffer first if true.
118 // Doing so would avoid overallocating when vec is filled to capacity and we reached EOF.
119 pub(crate) fn try_small_read_first(&self, num_bytes: usize) -> bool {
120 let vec = self.vec.as_ref();
121 vec.capacity() - vec.len() < num_bytes
122 && self.starting_capacity == vec.capacity()
123 && self.starting_capacity >= num_bytes
124 }
125}
126
127pub(crate) struct ReadBufParts {
128 // Pointer is only used to check that the ReadBuf actually came from the
129 // right VecWithInitialized.
130 ptr: *const u8,
131 len: usize,
132 initialized: usize,
133}
134
135// This is needed to release the borrow on `VecWithInitialized<V>`.
136pub(crate) fn into_read_buf_parts(rb: ReadBuf<'_>) -> ReadBufParts {
137 ReadBufParts {
138 ptr: rb.filled().as_ptr(),
139 len: rb.filled().len(),
140 initialized: rb.initialized().len(),
141 }
142}
143