1 // This lint claims ugly casting is somehow safer than transmute, but there's
2 // no evidence that is the case. Shush.
3 #![allow(clippy::transmute_ptr_to_ptr)]
4 
5 use std::fmt;
6 use std::mem::{self, MaybeUninit};
7 
8 /// A wrapper around a byte buffer that is incrementally filled and initialized.
9 ///
10 /// This type is a sort of "double cursor". It tracks three regions in the
11 /// buffer: a region at the beginning of the buffer that has been logically
12 /// filled with data, a region that has been initialized at some point but not
13 /// yet logically filled, and a region at the end that may be uninitialized.
14 /// The filled region is guaranteed to be a subset of the initialized region.
15 ///
16 /// In summary, the contents of the buffer can be visualized as:
17 ///
18 /// ```not_rust
19 /// [             capacity              ]
20 /// [ filled |         unfilled         ]
21 /// [    initialized    | uninitialized ]
22 /// ```
23 ///
24 /// It is undefined behavior to de-initialize any bytes from the uninitialized
25 /// region, since it is merely unknown whether this region is uninitialized or
26 /// not, and if part of it turns out to be initialized, it must stay initialized.
27 pub struct ReadBuf<'a> {
28     buf: &'a mut [MaybeUninit<u8>],
29     filled: usize,
30     initialized: usize,
31 }
32 
33 impl<'a> ReadBuf<'a> {
34     /// Creates a new `ReadBuf` from a fully initialized buffer.
35     #[inline]
new(buf: &'a mut [u8]) -> ReadBuf<'a>36     pub fn new(buf: &'a mut [u8]) -> ReadBuf<'a> {
37         let initialized = buf.len();
38         let buf = unsafe { mem::transmute::<&mut [u8], &mut [MaybeUninit<u8>]>(buf) };
39         ReadBuf {
40             buf,
41             filled: 0,
42             initialized,
43         }
44     }
45 
46     /// Creates a new `ReadBuf` from a fully uninitialized buffer.
47     ///
48     /// Use `assume_init` if part of the buffer is known to be already inintialized.
49     #[inline]
uninit(buf: &'a mut [MaybeUninit<u8>]) -> ReadBuf<'a>50     pub fn uninit(buf: &'a mut [MaybeUninit<u8>]) -> ReadBuf<'a> {
51         ReadBuf {
52             buf,
53             filled: 0,
54             initialized: 0,
55         }
56     }
57 
58     /// Returns the total capacity of the buffer.
59     #[inline]
capacity(&self) -> usize60     pub fn capacity(&self) -> usize {
61         self.buf.len()
62     }
63 
64     /// Returns a shared reference to the filled portion of the buffer.
65     #[inline]
filled(&self) -> &[u8]66     pub fn filled(&self) -> &[u8] {
67         let slice = &self.buf[..self.filled];
68         // safety: filled describes how far into the buffer that the
69         // user has filled with bytes, so it's been initialized.
70         // TODO: This could use `MaybeUninit::slice_get_ref` when it is stable.
71         unsafe { mem::transmute::<&[MaybeUninit<u8>], &[u8]>(slice) }
72     }
73 
74     /// Returns a mutable reference to the filled portion of the buffer.
75     #[inline]
filled_mut(&mut self) -> &mut [u8]76     pub fn filled_mut(&mut self) -> &mut [u8] {
77         let slice = &mut self.buf[..self.filled];
78         // safety: filled describes how far into the buffer that the
79         // user has filled with bytes, so it's been initialized.
80         // TODO: This could use `MaybeUninit::slice_get_mut` when it is stable.
81         unsafe { mem::transmute::<&mut [MaybeUninit<u8>], &mut [u8]>(slice) }
82     }
83 
84     /// Returns a new `ReadBuf` comprised of the unfilled section up to `n`.
85     #[inline]
take(&mut self, n: usize) -> ReadBuf<'_>86     pub fn take(&mut self, n: usize) -> ReadBuf<'_> {
87         let max = std::cmp::min(self.remaining(), n);
88         // Saftey: We don't set any of the `unfilled_mut` with `MaybeUninit::uninit`.
89         unsafe { ReadBuf::uninit(&mut self.unfilled_mut()[..max]) }
90     }
91 
92     /// Returns a shared reference to the initialized portion of the buffer.
93     ///
94     /// This includes the filled portion.
95     #[inline]
initialized(&self) -> &[u8]96     pub fn initialized(&self) -> &[u8] {
97         let slice = &self.buf[..self.initialized];
98         // safety: initialized describes how far into the buffer that the
99         // user has at some point initialized with bytes.
100         // TODO: This could use `MaybeUninit::slice_get_ref` when it is stable.
101         unsafe { mem::transmute::<&[MaybeUninit<u8>], &[u8]>(slice) }
102     }
103 
104     /// Returns a mutable reference to the initialized portion of the buffer.
105     ///
106     /// This includes the filled portion.
107     #[inline]
initialized_mut(&mut self) -> &mut [u8]108     pub fn initialized_mut(&mut self) -> &mut [u8] {
109         let slice = &mut self.buf[..self.initialized];
110         // safety: initialized describes how far into the buffer that the
111         // user has at some point initialized with bytes.
112         // TODO: This could use `MaybeUninit::slice_get_mut` when it is stable.
113         unsafe { mem::transmute::<&mut [MaybeUninit<u8>], &mut [u8]>(slice) }
114     }
115 
116     /// Returns a mutable reference to the entire buffer, without ensuring that it has been fully
117     /// initialized.
118     ///
119     /// The elements between 0 and `self.filled().len()` are filled, and those between 0 and
120     /// `self.initialized().len()` are initialized (and so can be transmuted to a `&mut [u8]`).
121     ///
122     /// The caller of this method must ensure that these invariants are upheld. For example, if the
123     /// caller initializes some of the uninitialized section of the buffer, it must call
124     /// [`assume_init`](Self::assume_init) with the number of bytes initialized.
125     ///
126     /// # Safety
127     ///
128     /// The caller must not de-initialize portions of the buffer that have already been initialized.
129     /// This includes any bytes in the region marked as uninitialized by `ReadBuf`.
130     #[inline]
inner_mut(&mut self) -> &mut [MaybeUninit<u8>]131     pub unsafe fn inner_mut(&mut self) -> &mut [MaybeUninit<u8>] {
132         self.buf
133     }
134 
135     /// Returns a mutable reference to the unfilled part of the buffer without ensuring that it has been fully
136     /// initialized.
137     ///
138     /// # Safety
139     ///
140     /// The caller must not de-initialize portions of the buffer that have already been initialized.
141     /// This includes any bytes in the region marked as uninitialized by `ReadBuf`.
142     #[inline]
unfilled_mut(&mut self) -> &mut [MaybeUninit<u8>]143     pub unsafe fn unfilled_mut(&mut self) -> &mut [MaybeUninit<u8>] {
144         &mut self.buf[self.filled..]
145     }
146 
147     /// Returns a mutable reference to the unfilled part of the buffer, ensuring it is fully initialized.
148     ///
149     /// Since `ReadBuf` tracks the region of the buffer that has been initialized, this is effectively "free" after
150     /// the first use.
151     #[inline]
initialize_unfilled(&mut self) -> &mut [u8]152     pub fn initialize_unfilled(&mut self) -> &mut [u8] {
153         self.initialize_unfilled_to(self.remaining())
154     }
155 
156     /// Returns a mutable reference to the first `n` bytes of the unfilled part of the buffer, ensuring it is
157     /// fully initialized.
158     ///
159     /// # Panics
160     ///
161     /// Panics if `self.remaining()` is less than `n`.
162     #[inline]
initialize_unfilled_to(&mut self, n: usize) -> &mut [u8]163     pub fn initialize_unfilled_to(&mut self, n: usize) -> &mut [u8] {
164         assert!(self.remaining() >= n, "n overflows remaining");
165 
166         // This can't overflow, otherwise the assert above would have failed.
167         let end = self.filled + n;
168 
169         if self.initialized < end {
170             unsafe {
171                 self.buf[self.initialized..end]
172                     .as_mut_ptr()
173                     .write_bytes(0, end - self.initialized);
174             }
175             self.initialized = end;
176         }
177 
178         let slice = &mut self.buf[self.filled..end];
179         // safety: just above, we checked that the end of the buf has
180         // been initialized to some value.
181         unsafe { mem::transmute::<&mut [MaybeUninit<u8>], &mut [u8]>(slice) }
182     }
183 
184     /// Returns the number of bytes at the end of the slice that have not yet been filled.
185     #[inline]
remaining(&self) -> usize186     pub fn remaining(&self) -> usize {
187         self.capacity() - self.filled
188     }
189 
190     /// Clears the buffer, resetting the filled region to empty.
191     ///
192     /// The number of initialized bytes is not changed, and the contents of the buffer are not modified.
193     #[inline]
clear(&mut self)194     pub fn clear(&mut self) {
195         self.filled = 0;
196     }
197 
198     /// Advances the size of the filled region of the buffer.
199     ///
200     /// The number of initialized bytes is not changed.
201     ///
202     /// # Panics
203     ///
204     /// Panics if the filled region of the buffer would become larger than the initialized region.
205     #[inline]
advance(&mut self, n: usize)206     pub fn advance(&mut self, n: usize) {
207         let new = self.filled.checked_add(n).expect("filled overflow");
208         self.set_filled(new);
209     }
210 
211     /// Sets the size of the filled region of the buffer.
212     ///
213     /// The number of initialized bytes is not changed.
214     ///
215     /// Note that this can be used to *shrink* the filled region of the buffer in addition to growing it (for
216     /// example, by a `AsyncRead` implementation that compresses data in-place).
217     ///
218     /// # Panics
219     ///
220     /// Panics if the filled region of the buffer would become larger than the intialized region.
221     #[inline]
set_filled(&mut self, n: usize)222     pub fn set_filled(&mut self, n: usize) {
223         assert!(
224             n <= self.initialized,
225             "filled must not become larger than initialized"
226         );
227         self.filled = n;
228     }
229 
230     /// Asserts that the first `n` unfilled bytes of the buffer are initialized.
231     ///
232     /// `ReadBuf` assumes that bytes are never de-initialized, so this method does nothing when called with fewer
233     /// bytes than are already known to be initialized.
234     ///
235     /// # Safety
236     ///
237     /// The caller must ensure that `n` unfilled bytes of the buffer have already been initialized.
238     #[inline]
assume_init(&mut self, n: usize)239     pub unsafe fn assume_init(&mut self, n: usize) {
240         let new = self.filled + n;
241         if new > self.initialized {
242             self.initialized = new;
243         }
244     }
245 
246     /// Appends data to the buffer, advancing the written position and possibly also the initialized position.
247     ///
248     /// # Panics
249     ///
250     /// Panics if `self.remaining()` is less than `buf.len()`.
251     #[inline]
put_slice(&mut self, buf: &[u8])252     pub fn put_slice(&mut self, buf: &[u8]) {
253         assert!(
254             self.remaining() >= buf.len(),
255             "buf.len() must fit in 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 impl fmt::Debug for ReadBuf<'_> {
fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result278     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
279         f.debug_struct("ReadBuf")
280             .field("filled", &self.filled)
281             .field("initialized", &self.initialized)
282             .field("capacity", &self.capacity())
283             .finish()
284     }
285 }
286