1use std::fmt;
2use std::mem::MaybeUninit;
3
4/// A wrapper around a byte buffer that is incrementally filled and initialized.
5///
6/// This type is a sort of "double cursor". It tracks three regions in the
7/// buffer: a region at the beginning of the buffer that has been logically
8/// filled with data, a region that has been initialized at some point but not
9/// yet logically filled, and a region at the end that may be uninitialized.
10/// The filled region is guaranteed to be a subset of the initialized region.
11///
12/// In summary, the contents of the buffer can be visualized as:
13///
14/// ```not_rust
15/// [ capacity ]
16/// [ filled | unfilled ]
17/// [ initialized | uninitialized ]
18/// ```
19///
20/// It is undefined behavior to de-initialize any bytes from the uninitialized
21/// region, since it is merely unknown whether this region is uninitialized or
22/// not, and if part of it turns out to be initialized, it must stay initialized.
23pub struct ReadBuf<'a> {
24 buf: &'a mut [MaybeUninit<u8>],
25 filled: usize,
26 initialized: usize,
27}
28
29impl<'a> ReadBuf<'a> {
30 /// Creates a new `ReadBuf` from a fully initialized buffer.
31 #[inline]
32 pub fn new(buf: &'a mut [u8]) -> ReadBuf<'a> {
33 let initialized = buf.len();
34 let buf = unsafe { slice_to_uninit_mut(buf) };
35 ReadBuf {
36 buf,
37 filled: 0,
38 initialized,
39 }
40 }
41
42 /// Creates a new `ReadBuf` from a fully uninitialized buffer.
43 ///
44 /// Use `assume_init` if part of the buffer is known to be already initialized.
45 #[inline]
46 pub fn uninit(buf: &'a mut [MaybeUninit<u8>]) -> ReadBuf<'a> {
47 ReadBuf {
48 buf,
49 filled: 0,
50 initialized: 0,
51 }
52 }
53
54 /// Returns the total capacity of the buffer.
55 #[inline]
56 pub fn capacity(&self) -> usize {
57 self.buf.len()
58 }
59
60 /// Returns a shared reference to the filled portion of the buffer.
61 #[inline]
62 pub fn filled(&self) -> &[u8] {
63 let slice = &self.buf[..self.filled];
64 // safety: filled describes how far into the buffer that the
65 // user has filled with bytes, so it's been initialized.
66 unsafe { slice_assume_init(slice) }
67 }
68
69 /// Returns a mutable reference to the filled portion of the buffer.
70 #[inline]
71 pub fn filled_mut(&mut self) -> &mut [u8] {
72 let slice = &mut self.buf[..self.filled];
73 // safety: filled describes how far into the buffer that the
74 // user has filled with bytes, so it's been initialized.
75 unsafe { slice_assume_init_mut(slice) }
76 }
77
78 /// Returns a new `ReadBuf` comprised of the unfilled section up to `n`.
79 #[inline]
80 pub fn take(&mut self, n: usize) -> ReadBuf<'_> {
81 let max = std::cmp::min(self.remaining(), n);
82 // Safety: We don't set any of the `unfilled_mut` with `MaybeUninit::uninit`.
83 unsafe { ReadBuf::uninit(&mut self.unfilled_mut()[..max]) }
84 }
85
86 /// Returns a shared reference to the initialized portion of the buffer.
87 ///
88 /// This includes the filled portion.
89 #[inline]
90 pub fn initialized(&self) -> &[u8] {
91 let slice = &self.buf[..self.initialized];
92 // safety: initialized describes how far into the buffer that the
93 // user has at some point initialized with bytes.
94 unsafe { slice_assume_init(slice) }
95 }
96
97 /// Returns a mutable reference to the initialized portion of the buffer.
98 ///
99 /// This includes the filled portion.
100 #[inline]
101 pub fn initialized_mut(&mut self) -> &mut [u8] {
102 let slice = &mut self.buf[..self.initialized];
103 // safety: initialized describes how far into the buffer that the
104 // user has at some point initialized with bytes.
105 unsafe { slice_assume_init_mut(slice) }
106 }
107
108 /// Returns a mutable reference to the entire buffer, without ensuring that it has been fully
109 /// initialized.
110 ///
111 /// The elements between 0 and `self.filled().len()` are filled, and those between 0 and
112 /// `self.initialized().len()` are initialized (and so can be converted to a `&mut [u8]`).
113 ///
114 /// The caller of this method must ensure that these invariants are upheld. For example, if the
115 /// caller initializes some of the uninitialized section of the buffer, it must call
116 /// [`assume_init`](Self::assume_init) with the number of bytes initialized.
117 ///
118 /// # Safety
119 ///
120 /// The caller must not de-initialize portions of the buffer that have already been initialized.
121 /// This includes any bytes in the region marked as uninitialized by `ReadBuf`.
122 #[inline]
123 pub unsafe fn inner_mut(&mut self) -> &mut [MaybeUninit<u8>] {
124 self.buf
125 }
126
127 /// Returns a mutable reference to the unfilled part of the buffer without ensuring that it has been fully
128 /// initialized.
129 ///
130 /// # Safety
131 ///
132 /// The caller must not de-initialize portions of the buffer that have already been initialized.
133 /// This includes any bytes in the region marked as uninitialized by `ReadBuf`.
134 #[inline]
135 pub unsafe fn unfilled_mut(&mut self) -> &mut [MaybeUninit<u8>] {
136 &mut self.buf[self.filled..]
137 }
138
139 /// Returns a mutable reference to the unfilled part of the buffer, ensuring it is fully initialized.
140 ///
141 /// Since `ReadBuf` tracks the region of the buffer that has been initialized, this is effectively "free" after
142 /// the first use.
143 #[inline]
144 pub fn initialize_unfilled(&mut self) -> &mut [u8] {
145 self.initialize_unfilled_to(self.remaining())
146 }
147
148 /// Returns a mutable reference to the first `n` bytes of the unfilled part of the buffer, ensuring it is
149 /// fully initialized.
150 ///
151 /// # Panics
152 ///
153 /// Panics if `self.remaining()` is less than `n`.
154 #[inline]
155 #[track_caller]
156 pub fn initialize_unfilled_to(&mut self, n: usize) -> &mut [u8] {
157 assert!(self.remaining() >= n, "n overflows remaining");
158
159 // This can't overflow, otherwise the assert above would have failed.
160 let end = self.filled + n;
161
162 if self.initialized < end {
163 unsafe {
164 self.buf[self.initialized..end]
165 .as_mut_ptr()
166 .write_bytes(0, end - self.initialized);
167 }
168 self.initialized = end;
169 }
170
171 let slice = &mut self.buf[self.filled..end];
172 // safety: just above, we checked that the end of the buf has
173 // been initialized to some value.
174 unsafe { slice_assume_init_mut(slice) }
175 }
176
177 /// Returns the number of bytes at the end of the slice that have not yet been filled.
178 #[inline]
179 pub fn remaining(&self) -> usize {
180 self.capacity() - self.filled
181 }
182
183 /// Clears the buffer, resetting the filled region to empty.
184 ///
185 /// The number of initialized bytes is not changed, and the contents of the buffer are not modified.
186 #[inline]
187 pub fn clear(&mut self) {
188 self.filled = 0;
189 }
190
191 /// Advances the size of the filled region of the buffer.
192 ///
193 /// The number of initialized bytes is not changed.
194 ///
195 /// # Panics
196 ///
197 /// Panics if the filled region of the buffer would become larger than the initialized region.
198 #[inline]
199 #[track_caller]
200 pub fn advance(&mut self, n: usize) {
201 let new = self.filled.checked_add(n).expect("filled overflow");
202 self.set_filled(new);
203 }
204
205 /// Sets the size of the filled region of the buffer.
206 ///
207 /// The number of initialized bytes is not changed.
208 ///
209 /// Note that this can be used to *shrink* the filled region of the buffer in addition to growing it (for
210 /// example, by a `AsyncRead` implementation that compresses data in-place).
211 ///
212 /// # Panics
213 ///
214 /// Panics if the filled region of the buffer would become larger than the initialized region.
215 #[inline]
216 #[track_caller]
217 pub fn set_filled(&mut self, n: usize) {
218 assert!(
219 n <= self.initialized,
220 "filled must not become larger than initialized"
221 );
222 self.filled = n;
223 }
224
225 /// Asserts that the first `n` unfilled bytes of the buffer are initialized.
226 ///
227 /// `ReadBuf` assumes that bytes are never de-initialized, so this method does nothing when called with fewer
228 /// bytes than are already known to be initialized.
229 ///
230 /// # Safety
231 ///
232 /// The caller must ensure that `n` unfilled bytes of the buffer have already been initialized.
233 #[inline]
234 pub unsafe fn assume_init(&mut self, n: usize) {
235 let new = self.filled + n;
236 if new > self.initialized {
237 self.initialized = new;
238 }
239 }
240
241 /// Appends data to the buffer, advancing the written position and possibly also the initialized position.
242 ///
243 /// # Panics
244 ///
245 /// Panics if `self.remaining()` is less than `buf.len()`.
246 #[inline]
247 #[track_caller]
248 pub fn put_slice(&mut self, buf: &[u8]) {
249 assert!(
250 self.remaining() >= buf.len(),
251 "buf.len() must fit in remaining()"
252 );
253
254 let amt = buf.len();
255 // Cannot overflow, asserted above
256 let end = self.filled + amt;
257
258 // Safety: the length is asserted above
259 unsafe {
260 self.buf[self.filled..end]
261 .as_mut_ptr()
262 .cast::<u8>()
263 .copy_from_nonoverlapping(buf.as_ptr(), amt);
264 }
265
266 if self.initialized < end {
267 self.initialized = end;
268 }
269 self.filled = end;
270 }
271}
272
273#[cfg(feature = "io-util")]
274#[cfg_attr(docsrs, doc(cfg(feature = "io-util")))]
275unsafe impl<'a> bytes::BufMut for ReadBuf<'a> {
276 fn remaining_mut(&self) -> usize {
277 self.remaining()
278 }
279
280 // SAFETY: The caller guarantees that at least `cnt` unfilled bytes have been initialized.
281 unsafe fn advance_mut(&mut self, cnt: usize) {
282 self.assume_init(cnt);
283 self.advance(cnt);
284 }
285
286 fn chunk_mut(&mut self) -> &mut bytes::buf::UninitSlice {
287 // SAFETY: No region of `unfilled` will be deinitialized because it is
288 // exposed as an `UninitSlice`, whose API guarantees that the memory is
289 // never deinitialized.
290 let unfilled = unsafe { self.unfilled_mut() };
291 let len = unfilled.len();
292 let ptr = unfilled.as_mut_ptr() as *mut u8;
293
294 // SAFETY: The pointer is valid for `len` bytes because it comes from a
295 // slice of that length.
296 unsafe { bytes::buf::UninitSlice::from_raw_parts_mut(ptr, len) }
297 }
298}
299
300impl fmt::Debug for ReadBuf<'_> {
301 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
302 f.debug_struct("ReadBuf")
303 .field("filled", &self.filled)
304 .field("initialized", &self.initialized)
305 .field("capacity", &self.capacity())
306 .finish()
307 }
308}
309
310unsafe fn slice_to_uninit_mut(slice: &mut [u8]) -> &mut [MaybeUninit<u8>] {
311 &mut *(slice as *mut [u8] as *mut [MaybeUninit<u8>])
312}
313
314// TODO: This could use `MaybeUninit::slice_assume_init` when it is stable.
315unsafe fn slice_assume_init(slice: &[MaybeUninit<u8>]) -> &[u8] {
316 &*(slice as *const [MaybeUninit<u8>] as *const [u8])
317}
318
319// TODO: This could use `MaybeUninit::slice_assume_init_mut` when it is stable.
320unsafe fn slice_assume_init_mut(slice: &mut [MaybeUninit<u8>]) -> &mut [u8] {
321 &mut *(slice as *mut [MaybeUninit<u8>] as *mut [u8])
322}
323