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 if n > self.remaining() { 151 panic!("overflowed: try to initialize more bytes than the buffer's capacity") 152 } 153 let end = self.filled + n; 154 155 if self.initialized < end { 156 unsafe { 157 self.buf[self.initialized..self.filled + n] 158 .as_mut_ptr() 159 .write_bytes(0, end - self.initialized); 160 } 161 self.initialized = end; 162 } 163 unsafe { &mut *(&mut self.buf[self.filled..end] as *mut [MaybeUninit<u8>] as *mut [u8]) } 164 } 165 166 /// Returns a mutable reference to the unfilled portion of the `ReadBuf`. 167 /// This method guarantees the the buffer is fully initialized. 168 #[inline] initialize_unfilled(&mut self) -> &mut [u8]169 pub fn initialize_unfilled(&mut self) -> &mut [u8] { 170 self.initialize_unfilled_to(self.remaining()) 171 } 172 173 /// Clears the `ReadBuf`. The filled size turns to zero while the 174 /// initialized size is unchanged. 175 #[inline] clear(&mut self)176 pub fn clear(&mut self) { 177 self.filled = 0; 178 } 179 180 /// Sets the filled size of the buffer. 181 /// 182 /// # Panics 183 /// Panics if the filled portion is bigger than the initialized portion of 184 /// the buffer. 185 #[inline] set_filled(&mut self, n: usize)186 pub fn set_filled(&mut self, n: usize) { 187 if n > self.initialized { 188 panic!("buf's filled size becomes larger than the initialized size") 189 } 190 self.filled = n; 191 } 192 193 /// Advances the filled portion of the buffer by n bytes. 194 /// 195 /// # Panics 196 /// 1. Panics if the filled size is overflowed after adding the advance 197 /// size. 198 /// 2. Panics if the filled portion becomes larger than the initialized 199 /// portion of the buffer. 200 #[inline] advance(&mut self, n: usize)201 pub fn advance(&mut self, n: usize) { 202 let filled = self 203 .filled 204 .checked_add(n) 205 .expect("buf filled size overflow"); 206 self.set_filled(filled); 207 } 208 209 /// Makes the n bytes after the filled portion of the buffer become 210 /// initialized. If adding n bytes exceeds the capacity, the initialized 211 /// size will be set to the capacity. 212 #[inline] assume_init(&mut self, n: usize)213 pub fn assume_init(&mut self, n: usize) { 214 let end = std::cmp::min(self.filled + n, self.capacity()); 215 if end > self.initialized { 216 self.initialized = end; 217 } 218 } 219 220 /// Appends the input data into the `BufRead`. Advances the filled size and 221 /// initialized size accordingly. 222 /// 223 /// # Panics 224 /// Panics if the size of the appending buffer is greater than the remaining 225 /// size of the `ReadBuf` 226 #[inline] append(&mut self, buf: &[u8])227 pub fn append(&mut self, buf: &[u8]) { 228 if buf.len() > self.remaining() { 229 panic!("slice size is larger than the buf's remaining size"); 230 } 231 let end = self.filled + buf.len(); 232 unsafe { 233 self.buf[self.filled..end] 234 .as_mut_ptr() 235 .cast::<u8>() 236 .copy_from_nonoverlapping(buf.as_ptr(), buf.len()); 237 } 238 239 if self.initialized < end { 240 self.initialized = end; 241 } 242 self.filled = end; 243 } 244 } 245 246 #[cfg(test)] 247 mod test { 248 use crate::io::ReadBuf; 249 250 /// UT test cases for `ReadBuf`. 251 /// 252 /// # Brief 253 /// 1. Create ReadBuf. 254 /// 2. Calls functions to get parameters 255 /// 3. Check if the test results are correct. 256 #[test] ut_test_readbuf_new()257 fn ut_test_readbuf_new() { 258 let mut buf = [0; 16]; 259 let buf_len = buf.len(); 260 let mut read_buf: ReadBuf<'_> = ReadBuf::new(&mut buf); 261 262 assert_eq!(read_buf.filled().len(), 0); 263 assert_eq!(read_buf.initialized_len(), 16); 264 assert_eq!(read_buf.capacity(), 16); 265 assert_eq!(read_buf.initialized().len(), 16); 266 267 read_buf.set_filled(9); 268 assert_eq!(read_buf.filled().len(), 9); 269 270 { 271 let borrow = read_buf.filled_mut(); 272 borrow[0] = 1; 273 } 274 assert_eq!(read_buf.filled()[0], 1); 275 276 { 277 let borrow = read_buf.initialized(); 278 assert_eq!(borrow[0], 1); 279 } 280 281 { 282 let borrow = read_buf.initialized_mut(); 283 borrow[1] = 2; 284 } 285 assert_eq!(read_buf.filled()[1], 2); 286 287 let buf_b = read_buf.initialize_unfilled_to(3); 288 assert_eq!(buf_b.len(), 3); 289 290 let buf_b = read_buf.initialize_unfilled(); 291 assert_eq!(buf_b.len(), buf_len - 9); 292 293 let buf_b = read_buf.inner_mut(); 294 assert_eq!(buf_b.len(), buf_len); 295 296 read_buf.clear(); 297 assert_eq!(read_buf.filled().len(), 0); 298 299 let new_read_buf = read_buf.take(8); 300 assert_eq!(new_read_buf.initialized_len(), 0); 301 } 302 } 303