1 | use std::fmt; |
2 | use 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. |
23 | pub struct ReadBuf<'a> { |
24 | buf: &'a mut [MaybeUninit<u8>], |
25 | filled: usize, |
26 | initialized: usize, |
27 | } |
28 | |
29 | impl<'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" )))] |
275 | unsafe 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 | |
300 | impl 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 | |
310 | unsafe 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. |
315 | unsafe 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. |
320 | unsafe fn slice_assume_init_mut(slice: &mut [MaybeUninit<u8>]) -> &mut [u8] { |
321 | &mut *(slice as *mut [MaybeUninit<u8>] as *mut [u8]) |
322 | } |
323 | |