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