library/std/src/io/readbuf.rs

   1 #![unstable(feature = "read_buf", issue = "78485")]
   2
   3 #[cfg(test)]
   4 mod tests;
   5
   6 use crate::cmp;
   7 use crate::fmt::{self, Debug, Formatter};
   8 use crate::mem::MaybeUninit;
   9
  10 /// A wrapper around a byte buffer that is incrementally filled and initialized.
  11 ///
  12 /// This type is a sort of "double cursor". It tracks three regions in the buffer: a region at the beginning of the
  13 /// buffer that has been logically filled with data, a region that has been initialized at some point but not yet
  14 /// logically filled, and a region at the end that is fully uninitialized. The filled region is guaranteed to be a
  15 /// subset of the initialized region.
  16 ///
  17 /// In summary, the contents of the buffer can be visualized as:
  18 /// ```not_rust
  19 /// [             capacity              ]
  20 /// [ filled |         unfilled         ]
  21 /// [    initialized    | uninitialized ]
  22 /// ```
  23 pub struct ReadBuf<'a> {
  24     buf: &'a mut [MaybeUninit<u8>],
  25     filled: usize,
  26     initialized: usize,
  27 }
  28
  29 impl Debug for ReadBuf<'_> {
  30     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
  31         f.debug_struct("ReadBuf")
  32             .field("init", &self.initialized())
  33             .field("filled", &self.filled)
  34             .field("capacity", &self.capacity())
  35             .finish()
  36     }
  37 }
  38
  39 impl<'a> ReadBuf<'a> {
  40     /// Creates a new `ReadBuf` from a fully initialized buffer.
  41     #[inline]
  42     pub fn new(buf: &'a mut [u8]) -> ReadBuf<'a> {
  43         let len = buf.len();
  44
  45         ReadBuf {
  46             //SAFETY: initialized data never becoming uninitialized is an invariant of ReadBuf
  47             buf: unsafe { (buf as *mut [u8]).as_uninit_slice_mut().unwrap() },
  48             filled: 0,
  49             initialized: len,
  50         }
  51     }
  52
  53     /// Creates a new `ReadBuf` from a fully uninitialized buffer.
  54     ///
  55     /// Use `assume_init` if part of the buffer is known to be already initialized.
  56     #[inline]
  57     pub fn uninit(buf: &'a mut [MaybeUninit<u8>]) -> ReadBuf<'a> {
  58         ReadBuf { buf, filled: 0, initialized: 0 }
  59     }
  60
  61     /// Returns the total capacity of the buffer.
  62     #[inline]
  63     pub fn capacity(&self) -> usize {
  64         self.buf.len()
  65     }
  66
  67     /// Returns a shared reference to the filled portion of the buffer.
  68     #[inline]
  69     pub fn filled(&self) -> &[u8] {
  70         //SAFETY: We only slice the filled part of the buffer, which is always valid
  71         unsafe { MaybeUninit::slice_assume_init_ref(&self.buf[0..self.filled]) }
  72     }
  73
  74     /// Returns a mutable reference to the filled portion of the buffer.
  75     #[inline]
  76     pub fn filled_mut(&mut self) -> &mut [u8] {
  77         //SAFETY: We only slice the filled part of the buffer, which is always valid
  78         unsafe { MaybeUninit::slice_assume_init_mut(&mut self.buf[0..self.filled]) }
  79     }
  80
  81     /// Returns a shared reference to the initialized portion of the buffer.
  82     ///
  83     /// This includes the filled portion.
  84     #[inline]
  85     pub fn initialized(&self) -> &[u8] {
  86         //SAFETY: We only slice the initialized part of the buffer, which is always valid
  87         unsafe { MaybeUninit::slice_assume_init_ref(&self.buf[0..self.initialized]) }
  88     }
  89
  90     /// Returns a mutable reference to the initialized portion of the buffer.
  91     ///
  92     /// This includes the filled portion.
  93     #[inline]
  94     pub fn initialized_mut(&mut self) -> &mut [u8] {
  95         //SAFETY: We only slice the initialized part of the buffer, which is always valid
  96         unsafe { MaybeUninit::slice_assume_init_mut(&mut self.buf[0..self.initialized]) }
  97     }
  98
  99     /// Returns a mutable reference to the unfilled part of the buffer without ensuring that it has been fully
 100     /// initialized.
 101     ///
 102     /// # Safety
 103     ///
 104     /// The caller must not de-initialize portions of the buffer that have already been initialized.
 105     #[inline]
 106     pub unsafe fn unfilled_mut(&mut self) -> &mut [MaybeUninit<u8>] {
 107         &mut self.buf[self.filled..]
 108     }
 109
 110     /// Returns a mutable reference to the uninitialized part of the buffer.
 111     ///
 112     /// It is safe to uninitialize any of these bytes.
 113     #[inline]
 114     pub fn uninitialized_mut(&mut self) -> &mut [MaybeUninit<u8>] {
 115         &mut self.buf[self.initialized..]
 116     }
 117
 118     /// Returns a mutable reference to the unfilled part of the buffer, ensuring it is fully initialized.
 119     ///
 120     /// Since `ReadBuf` tracks the region of the buffer that has been initialized, this is effectively "free" after
 121     /// the first use.
 122     #[inline]
 123     pub fn initialize_unfilled(&mut self) -> &mut [u8] {
 124         // should optimize out the assertion
 125         self.initialize_unfilled_to(self.remaining())
 126     }
 127
 128     /// Returns a mutable reference to the first `n` bytes of the unfilled part of the buffer, ensuring it is
 129     /// fully initialized.
 130     ///
 131     /// # Panics
 132     ///
 133     /// Panics if `self.remaining()` is less than `n`.
 134     #[inline]
 135     pub fn initialize_unfilled_to(&mut self, n: usize) -> &mut [u8] {
 136         assert!(self.remaining() >= n);
 137
 138         let extra_init = self.initialized - self.filled;
 139         // If we don't have enough initialized, do zeroing
 140         if n > extra_init {
 141             let uninit = n - extra_init;
 142             let unfilled = &mut self.uninitialized_mut()[0..uninit];
 143
 144             for byte in unfilled.iter_mut() {
 145                 byte.write(0);
 146             }
 147
 148             // SAFETY: we just initialized uninit bytes, and the previous bytes were already init
 149             unsafe {
 150                 self.assume_init(n);
 151             }
 152         }
 153
 154         let filled = self.filled;
 155
 156         &mut self.initialized_mut()[filled..filled + n]
 157     }
 158
 159     /// Returns the number of bytes at the end of the slice that have not yet been filled.
 160     #[inline]
 161     pub fn remaining(&self) -> usize {
 162         self.capacity() - self.filled
 163     }
 164
 165     /// Clears the buffer, resetting the filled region to empty.
 166     ///
 167     /// The number of initialized bytes is not changed, and the contents of the buffer are not modified.
 168     #[inline]
 169     pub fn clear(&mut self) -> &mut Self {
 170         self.set_filled(0) // The assertion in `set_filled` is optimized out
 171     }
 172
 173     /// Increases the size of the filled region of the buffer.
 174     ///
 175     /// The number of initialized bytes is not changed.
 176     ///
 177     /// # Panics
 178     ///
 179     /// Panics if the filled region of the buffer would become larger than the initialized region.
 180     #[inline]
 181     pub fn add_filled(&mut self, n: usize) -> &mut Self {
 182         self.set_filled(self.filled + n)
 183     }
 184
 185     /// Sets the size of the filled region of the buffer.
 186     ///
 187     /// The number of initialized bytes is not changed.
 188     ///
 189     /// Note that this can be used to *shrink* the filled region of the buffer in addition to growing it (for
 190     /// example, by a `Read` implementation that compresses data in-place).
 191     ///
 192     /// # Panics
 193     ///
 194     /// Panics if the filled region of the buffer would become larger than the initialized region.
 195     #[inline]
 196     pub fn set_filled(&mut self, n: usize) -> &mut Self {
 197         assert!(n <= self.initialized);
 198
 199         self.filled = n;
 200         self
 201     }
 202
 203     /// Asserts that the first `n` unfilled bytes of the buffer are initialized.
 204     ///
 205     /// `ReadBuf` assumes that bytes are never de-initialized, so this method does nothing when called with fewer
 206     /// bytes than are already known to be initialized.
 207     ///
 208     /// # Safety
 209     ///
 210     /// The caller must ensure that the first `n` unfilled bytes of the buffer have already been initialized.
 211     #[inline]
 212     pub unsafe fn assume_init(&mut self, n: usize) -> &mut Self {
 213         self.initialized = cmp::max(self.initialized, self.filled + n);
 214         self
 215     }
 216
 217     /// Appends data to the buffer, advancing the written position and possibly also the initialized position.
 218     ///
 219     /// # Panics
 220     ///
 221     /// Panics if `self.remaining()` is less than `buf.len()`.
 222     #[inline]
 223     pub fn append(&mut self, buf: &[u8]) {
 224         assert!(self.remaining() >= buf.len());
 225
 226         // SAFETY: we do not de-initialize any of the elements of the slice
 227         unsafe {
 228             MaybeUninit::write_slice(&mut self.unfilled_mut()[..buf.len()], buf);
 229         }
 230
 231         // SAFETY: We just added the entire contents of buf to the filled section.
 232         unsafe {
 233             self.assume_init(buf.len());
 234         }
 235         self.add_filled(buf.len());
 236     }
 237
 238     /// Returns the amount of bytes that have been filled.
 239     #[inline]
 240     pub fn filled_len(&self) -> usize {
 241         self.filled
 242     }
 243
 244     /// Returns the amount of bytes that have been initialized.
 245     #[inline]
 246     pub fn initialized_len(&self) -> usize {
 247         self.initialized
 248     }
 249 }