1 // Copyright (c) 2023 Huawei Device Co., Ltd.
2 // Licensed under the Apache License, Version 2.0 (the "License");
3 // you may not use this file except in compliance with the License.
4 // You may obtain a copy of the License at
5 //
6 //     http://www.apache.org/licenses/LICENSE-2.0
7 //
8 // Unless required by applicable law or agreed to in writing, software
9 // distributed under the License is distributed on an "AS IS" BASIS,
10 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11 // See the License for the specific language governing permissions and
12 // limitations under the License.
13 
14 use std::mem::MaybeUninit;
15 
16 /// This buf comes from std::io::ReadBuf, an unstable std lib. This buffer is a
17 /// wrapper around byte buffer and it allows users to read data into an
18 /// uninitialized memory. It tracks three regions in the buffer: a region at the
19 /// beginning of the buffer that has been logically filled with data,
20 /// a region that has been initialized at some point but not yet logically
21 /// filled, and a region at the end that is fully uninitialized. The filled
22 /// region is guaranteed to be a subset of the initialized region.
23 ///
24 /// In summary, the contents of the buffer can be visualized as:
25 /// ```not_rust
26 /// [             capacity              ]
27 /// [ filled |         unfilled         ]
28 /// [    initialized    | uninitialized ]
29 /// ```
30 pub struct ReadBuf<'a> {
31     pub(crate) buf: &'a mut [MaybeUninit<u8>],
32     filled: usize,
33     initialized: usize,
34 }
35 
36 impl<'a> ReadBuf<'a> {
37     /// Creates a `ReadBuf` from a fully initialized byte buffer.
38     #[inline]
new(buf: &'a mut [u8]) -> ReadBuf<'a>39     pub fn new(buf: &'a mut [u8]) -> ReadBuf<'a> {
40         ReadBuf {
41             buf: unsafe { &mut *(buf as *mut [u8] as *mut [MaybeUninit<u8>]) },
42             filled: 0,
43             initialized: buf.len(),
44         }
45     }
46 
47     /// Creates a `ReadBuf` from an uninitialized byte buffer.
48     #[inline]
create( buf: &'a mut [MaybeUninit<u8>], filled: usize, initialized: usize, ) -> ReadBuf<'a>49     pub fn create(
50         buf: &'a mut [MaybeUninit<u8>],
51         filled: usize,
52         initialized: usize,
53     ) -> ReadBuf<'a> {
54         ReadBuf {
55             buf,
56             filled,
57             initialized,
58         }
59     }
60 
61     /// Creates a `ReadBuf` from a fully uninitialized byte buffer.
62     #[inline]
uninit(buf: &mut [MaybeUninit<u8>]) -> ReadBuf<'_>63     pub fn uninit(buf: &mut [MaybeUninit<u8>]) -> ReadBuf<'_> {
64         ReadBuf {
65             buf,
66             filled: 0,
67             initialized: 0,
68         }
69     }
70 
71     /// Returns the total buffer capacity.
72     #[inline]
capacity(&self) -> usize73     pub fn capacity(&self) -> usize {
74         self.buf.len()
75     }
76 
77     /// Returns the size of the filled portion of the buffer
78     #[inline]
filled_len(&self) -> usize79     pub fn filled_len(&self) -> usize {
80         self.filled
81     }
82 
83     /// Returns the length of the initialized portion of the buffer.
84     #[inline]
initialized_len(&self) -> usize85     pub fn initialized_len(&self) -> usize {
86         self.initialized
87     }
88 
89     /// Returns a new ReadBuf that uses the first `n` unfilled bytes of the
90     /// buffer.
91     #[inline]
take(&mut self, n: usize) -> ReadBuf<'_>92     pub fn take(&mut self, n: usize) -> ReadBuf<'_> {
93         let rsize = std::cmp::min(n, self.remaining());
94         ReadBuf::uninit(&mut self.unfilled_mut()[..rsize])
95     }
96 
97     /// Returns a reference to the filled portion of the `ReadBuf`.
98     #[inline]
filled(&self) -> &[u8]99     pub fn filled(&self) -> &[u8] {
100         unsafe { &*(&self.buf[..self.filled] as *const [MaybeUninit<u8>] as *const [u8]) }
101     }
102 
103     /// Returns a mutable reference to the filled portion of the `ReadBuf`.
104     #[inline]
filled_mut(&mut self) -> &mut [u8]105     pub fn filled_mut(&mut self) -> &mut [u8] {
106         unsafe { &mut *(&mut self.buf[..self.filled] as *mut [MaybeUninit<u8>] as *mut [u8]) }
107     }
108 
109     /// Returns a mutable reference to the unfilled portion of the `ReadBuf`.
110     #[inline]
unfilled_mut(&mut self) -> &mut [MaybeUninit<u8>]111     pub fn unfilled_mut(&mut self) -> &mut [MaybeUninit<u8>] {
112         &mut self.buf[self.filled..]
113     }
114 
115     /// Returns a reference to the initialized portion of the `ReadBuf`.
116     #[inline]
initialized(&self) -> &[u8]117     pub fn initialized(&self) -> &[u8] {
118         unsafe { &*(&self.buf[..self.initialized] as *const [MaybeUninit<u8>] as *const [u8]) }
119     }
120 
121     /// Returns a mutable reference to the initialized portion of the `ReadBuf`.
122     #[inline]
initialized_mut(&mut self) -> &mut [u8]123     pub fn initialized_mut(&mut self) -> &mut [u8] {
124         unsafe { &mut *(&mut self.buf[..self.initialized] as *mut [MaybeUninit<u8>] as *mut [u8]) }
125     }
126 
127     /// Returns a mutable reference to the entire buffer, including the
128     /// initialized and uninitialized portion. If the buffer is partially
129     /// initialized, the caller must call [`ReadBuf::assume_init`] with
130     /// the number of bytes initialized.
131     #[inline]
inner_mut(&mut self) -> &mut [MaybeUninit<u8>]132     pub fn inner_mut(&mut self) -> &mut [MaybeUninit<u8>] {
133         self.buf
134     }
135 
136     /// Returns the remaining unfilled space of the `ReadBuf`.
137     #[inline]
remaining(&self) -> usize138     pub fn remaining(&self) -> usize {
139         self.buf.len() - self.filled
140     }
141 
142     /// Returns a mutable reference to the first n bytes of the unfilled portion
143     /// of the `ReadBuf`. This method guarantees the returned buffer is
144     /// fully initialized.
145     ///
146     /// # Panics
147     /// Panics if n is bigger than the remaining capacity of the buf.
148     #[inline]
initialize_unfilled_to(&mut self, n: usize) -> &mut [u8]149     pub fn initialize_unfilled_to(&mut self, n: usize) -> &mut [u8] {
150         let remaining = self.remaining();
151         assert!(n <= remaining, "overflowed: try to initialize more bytes than the buffer's capacity, n: {}, remaining: {}", n, remaining);
152         let end = self.filled + n;
153 
154         if self.initialized < end {
155             unsafe {
156                 self.buf[self.initialized..self.filled + n]
157                     .as_mut_ptr()
158                     .write_bytes(0, end - self.initialized);
159             }
160             self.initialized = end;
161         }
162         unsafe { &mut *(&mut self.buf[self.filled..end] as *mut [MaybeUninit<u8>] as *mut [u8]) }
163     }
164 
165     /// Returns a mutable reference to the unfilled portion of the `ReadBuf`.
166     /// This method guarantees the the buffer is fully initialized.
167     #[inline]
initialize_unfilled(&mut self) -> &mut [u8]168     pub fn initialize_unfilled(&mut self) -> &mut [u8] {
169         self.initialize_unfilled_to(self.remaining())
170     }
171 
172     /// Clears the `ReadBuf`. The filled size turns to zero while the
173     /// initialized size is unchanged.
174     #[inline]
clear(&mut self)175     pub fn clear(&mut self) {
176         self.filled = 0;
177     }
178 
179     /// Sets the filled size of the buffer.
180     ///
181     /// # Panics
182     /// Panics if the filled portion is bigger than the initialized portion of
183     /// the buffer.
184     #[inline]
set_filled(&mut self, n: usize)185     pub fn set_filled(&mut self, n: usize) {
186         let initialized_len = self.initialized;
187         assert!(
188             n <= initialized_len,
189             "buf's filled size becomes larger than the initialized size, n: {}, initialized: {}",
190             n,
191             initialized_len
192         );
193         self.filled = n;
194     }
195 
196     /// Advances the filled portion of the buffer by n bytes.
197     ///
198     /// # Panics
199     /// 1. Panics if the filled size is overflowed after adding the advance
200     ///    size.
201     /// 2. Panics if the filled portion becomes larger than the initialized
202     ///    portion of the buffer.
203     #[inline]
advance(&mut self, n: usize)204     pub fn advance(&mut self, n: usize) {
205         let filled = self
206             .filled
207             .checked_add(n)
208             .expect("buf filled size overflow");
209         self.set_filled(filled);
210     }
211 
212     /// Makes the n bytes after the filled portion of the buffer become
213     /// initialized. If adding n bytes exceeds the capacity, the initialized
214     /// size will be set to the capacity.
215     #[inline]
assume_init(&mut self, n: usize)216     pub fn assume_init(&mut self, n: usize) {
217         let end = std::cmp::min(self.filled + n, self.capacity());
218         if end > self.initialized {
219             self.initialized = end;
220         }
221     }
222 
223     /// Appends the input data into the `BufRead`. Advances the filled size and
224     /// initialized size accordingly.
225     ///
226     /// # Panics
227     /// Panics if the size of the appending buffer is greater than the remaining
228     /// size of the `ReadBuf`
229     #[inline]
append(&mut self, buf: &[u8])230     pub fn append(&mut self, buf: &[u8]) {
231         let remaining = self.remaining();
232         assert!(
233             buf.len() <= remaining,
234             "slice size is larger than the buf's remaining size, buf len: {}, remaining: {}",
235             buf.len(),
236             remaining
237         );
238 
239         let end = self.filled + buf.len();
240         unsafe {
241             self.buf[self.filled..end]
242                 .as_mut_ptr()
243                 .cast::<u8>()
244                 .copy_from_nonoverlapping(buf.as_ptr(), buf.len());
245         }
246 
247         if self.initialized < end {
248             self.initialized = end;
249         }
250         self.filled = end;
251     }
252 }
253 
254 #[cfg(test)]
255 mod test {
256     use std::mem::MaybeUninit;
257 
258     use crate::io::ReadBuf;
259 
260     /// UT test cases for `ReadBuf`.
261     ///
262     /// # Brief
263     /// 1. Create ReadBuf.
264     /// 2. Calls functions to get parameters
265     /// 3. Check if the test results are correct.
266     /// 4. Create ReadBuf using uninitialized buffer.
267     /// 5. Calls functions to get parameters.
268     /// 6. Check if the test results are correct.
269     #[test]
ut_test_readbuf_new()270     fn ut_test_readbuf_new() {
271         let mut buf = [0; 16];
272         let buf_len = buf.len();
273         let mut read_buf: ReadBuf<'_> = ReadBuf::new(&mut buf);
274 
275         assert_eq!(read_buf.filled().len(), 0);
276         assert_eq!(read_buf.initialized_len(), 16);
277         assert_eq!(read_buf.capacity(), 16);
278         assert_eq!(read_buf.initialized().len(), 16);
279 
280         read_buf.set_filled(9);
281         assert_eq!(read_buf.filled().len(), 9);
282 
283         {
284             let borrow = read_buf.filled_mut();
285             borrow[0] = 1;
286         }
287         assert_eq!(read_buf.filled()[0], 1);
288 
289         {
290             let borrow = read_buf.initialized();
291             assert_eq!(borrow[0], 1);
292         }
293 
294         {
295             let borrow = read_buf.initialized_mut();
296             borrow[1] = 2;
297         }
298         assert_eq!(read_buf.filled()[1], 2);
299 
300         let buf_b = read_buf.initialize_unfilled_to(3);
301         assert_eq!(buf_b.len(), 3);
302 
303         let buf_b = read_buf.initialize_unfilled();
304         assert_eq!(buf_b.len(), buf_len - 9);
305 
306         let buf_b = read_buf.inner_mut();
307         assert_eq!(buf_b.len(), buf_len);
308 
309         read_buf.clear();
310         assert_eq!(read_buf.filled().len(), 0);
311 
312         let new_read_buf = read_buf.take(8);
313         assert_eq!(new_read_buf.initialized_len(), 0);
314 
315         let mut uninit_buf: [MaybeUninit<u8>; 16] = unsafe { MaybeUninit::zeroed().assume_init() };
316         let mut len: usize = 0;
317         let filled = 4;
318         for var in &mut uninit_buf[0..8] {
319             if len < filled {
320                 var.write(3);
321             } else {
322                 var.write(0);
323             }
324             len += 1;
325         }
326         let mut uninit_read_buf = ReadBuf::create(&mut uninit_buf, filled, len);
327         assert_eq!(uninit_read_buf.filled().len(), filled);
328         assert_eq!(uninit_read_buf.initialized_len(), len);
329         assert_eq!(uninit_read_buf.capacity(), 16);
330         assert_eq!(uninit_read_buf.initialized().len(), len);
331 
332         let buf_b = uninit_read_buf.initialize_unfilled_to(6);
333         assert_eq!(buf_b.len(), 6);
334     }
335 }
336