tokio/io/util/
vec_with_initialized.rs

Help
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}
tokio/io/util/vec_with_initialized.rs

tokio/io/util/
vec_with_initialized.rs
