1 use crate::io::ReadBuf;
2 use 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.
11 pub(crate) unsafe trait VecU8: AsRef<Vec<u8>> + AsMut<Vec<u8>> {}
12
13 unsafe impl VecU8 for Vec<u8> {}
14 unsafe 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)]
26 pub(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 }
32
33 impl VecWithInitialized<Vec<u8>> {
34 #[cfg(feature = "io-util")]
take(&mut self) -> Vec<u8>35 pub(crate) fn take(&mut self) -> Vec<u8> {
36 self.num_initialized = 0;
37 std::mem::take(&mut self.vec)
38 }
39 }
40
41 impl<V> VecWithInitialized<V>
42 where
43 V: VecU8,
44 {
new(mut vec: V) -> Self45 pub(crate) fn new(mut vec: V) -> Self {
46 // SAFETY: The safety invariants of vector guarantee that the bytes up
47 // to its length are initialized.
48 Self {
49 num_initialized: vec.as_mut().len(),
50 vec,
51 }
52 }
53
reserve(&mut self, num_bytes: usize)54 pub(crate) fn reserve(&mut self, num_bytes: usize) {
55 let vec = self.vec.as_mut();
56 if vec.capacity() - vec.len() >= num_bytes {
57 return;
58 }
59 // SAFETY: Setting num_initialized to `vec.len()` is correct as
60 // `reserve` does not change the length of the vector.
61 self.num_initialized = vec.len();
62 vec.reserve(num_bytes);
63 }
64
65 #[cfg(feature = "io-util")]
is_empty(&self) -> bool66 pub(crate) fn is_empty(&self) -> bool {
67 self.vec.as_ref().is_empty()
68 }
69
get_read_buf<'a>(&'a mut self) -> ReadBuf<'a>70 pub(crate) fn get_read_buf<'a>(&'a mut self) -> ReadBuf<'a> {
71 let num_initialized = self.num_initialized;
72
73 // SAFETY: Creating the slice is safe because of the safety invariants
74 // on Vec<u8>. The safety invariants of `ReadBuf` will further guarantee
75 // that no bytes in the slice are de-initialized.
76 let vec = self.vec.as_mut();
77 let len = vec.len();
78 let cap = vec.capacity();
79 let ptr = vec.as_mut_ptr().cast::<MaybeUninit<u8>>();
80 let slice = unsafe { std::slice::from_raw_parts_mut::<'a, MaybeUninit<u8>>(ptr, cap) };
81
82 // SAFETY: This is safe because the safety invariants of
83 // VecWithInitialized say that the first num_initialized bytes must be
84 // initialized.
85 let mut read_buf = ReadBuf::uninit(slice);
86 unsafe {
87 read_buf.assume_init(num_initialized);
88 }
89 read_buf.set_filled(len);
90
91 read_buf
92 }
93
apply_read_buf(&mut self, parts: ReadBufParts)94 pub(crate) fn apply_read_buf(&mut self, parts: ReadBufParts) {
95 let vec = self.vec.as_mut();
96 assert_eq!(vec.as_ptr(), parts.ptr);
97
98 // SAFETY:
99 // The ReadBufParts really does point inside `self.vec` due to the above
100 // check, and the safety invariants of `ReadBuf` guarantee that the
101 // first `parts.initialized` bytes of `self.vec` really have been
102 // initialized. Additionally, `ReadBuf` guarantees that `parts.len` is
103 // at most `parts.initialized`, so the first `parts.len` bytes are also
104 // initialized.
105 //
106 // Note that this relies on the fact that `V` is either `Vec<u8>` or
107 // `&mut Vec<u8>`, so the vector returned by `self.vec.as_mut()` cannot
108 // change from call to call.
109 unsafe {
110 self.num_initialized = parts.initialized;
111 vec.set_len(parts.len);
112 }
113 }
114 }
115
116 pub(crate) struct ReadBufParts {
117 // Pointer is only used to check that the ReadBuf actually came from the
118 // right VecWithInitialized.
119 ptr: *const u8,
120 len: usize,
121 initialized: usize,
122 }
123
124 // This is needed to release the borrow on `VecWithInitialized<V>`.
into_read_buf_parts(rb: ReadBuf<'_>) -> ReadBufParts125 pub(crate) fn into_read_buf_parts(rb: ReadBuf<'_>) -> ReadBufParts {
126 ReadBufParts {
127 ptr: rb.filled().as_ptr(),
128 len: rb.filled().len(),
129 initialized: rb.initialized().len(),
130 }
131 }
132