library/std/src/io/buffered/bufreader/buffer.rs

   1 ///! An encapsulation of `BufReader`'s buffer management logic.
   2 ///
   3 /// This module factors out the basic functionality of `BufReader` in order to protect two core
   4 /// invariants:
   5 /// * `filled` bytes of `buf` are always initialized
   6 /// * `pos` is always <= `filled`
   7 /// Since this module encapsulates the buffer management logic, we can ensure that the range
   8 /// `pos..filled` is always a valid index into the initialized region of the buffer. This means
   9 /// that user code which wants to do reads from a `BufReader` via `buffer` + `consume` can do so
  10 /// without encountering any runtime bounds checks.
  11 use crate::cmp;
  12 use crate::io::{self, BorrowedBuf, Read};
  13 use crate::mem::MaybeUninit;
  14
  15 pub struct Buffer {
  16     // The buffer.
  17     buf: Box<[MaybeUninit<u8>]>,
  18     // The current seek offset into `buf`, must always be <= `filled`.
  19     pos: usize,
  20     // Each call to `fill_buf` sets `filled` to indicate how many bytes at the start of `buf` are
  21     // initialized with bytes from a read.
  22     filled: usize,
  23     // This is the max number of bytes returned across all `fill_buf` calls. We track this so that we
  24     // can accurately tell `read_buf` how many bytes of buf are initialized, to bypass as much of its
  25     // defensive initialization as possible. Note that while this often the same as `filled`, it
  26     // doesn't need to be. Calls to `fill_buf` are not required to actually fill the buffer, and
  27     // omitting this is a huge perf regression for `Read` impls that do not.
  28     initialized: usize,
  29 }
  30
  31 impl Buffer {
  32     #[inline]
  33     pub fn with_capacity(capacity: usize) -> Self {
  34         let buf = Box::new_uninit_slice(capacity);
  35         Self { buf, pos: 0, filled: 0, initialized: 0 }
  36     }
  37
  38     #[inline]
  39     pub fn buffer(&self) -> &[u8] {
  40         // SAFETY: self.pos and self.cap are valid, and self.cap => self.pos, and
  41         // that region is initialized because those are all invariants of this type.
  42         unsafe { MaybeUninit::slice_assume_init_ref(self.buf.get_unchecked(self.pos..self.filled)) }
  43     }
  44
  45     #[inline]
  46     pub fn capacity(&self) -> usize {
  47         self.buf.len()
  48     }
  49
  50     #[inline]
  51     pub fn filled(&self) -> usize {
  52         self.filled
  53     }
  54
  55     #[inline]
  56     pub fn pos(&self) -> usize {
  57         self.pos
  58     }
  59
  60     // This is only used by a test which asserts that the initialization-tracking is correct.
  61     #[cfg(test)]
  62     pub fn initialized(&self) -> usize {
  63         self.initialized
  64     }
  65
  66     #[inline]
  67     pub fn discard_buffer(&mut self) {
  68         self.pos = 0;
  69         self.filled = 0;
  70     }
  71
  72     #[inline]
  73     pub fn consume(&mut self, amt: usize) {
  74         self.pos = cmp::min(self.pos + amt, self.filled);
  75     }
  76
  77     /// If there are `amt` bytes available in the buffer, pass a slice containing those bytes to
  78     /// `visitor` and return true. If there are not enough bytes available, return false.
  79     #[inline]
  80     pub fn consume_with<V>(&mut self, amt: usize, mut visitor: V) -> bool
  81     where
  82         V: FnMut(&[u8]),
  83     {
  84         if let Some(claimed) = self.buffer().get(..amt) {
  85             visitor(claimed);
  86             // If the indexing into self.buffer() succeeds, amt must be a valid increment.
  87             self.pos += amt;
  88             true
  89         } else {
  90             false
  91         }
  92     }
  93
  94     #[inline]
  95     pub fn unconsume(&mut self, amt: usize) {
  96         self.pos = self.pos.saturating_sub(amt);
  97     }
  98
  99     #[inline]
 100     pub fn fill_buf(&mut self, mut reader: impl Read) -> io::Result<&[u8]> {
 101         // If we've reached the end of our internal buffer then we need to fetch
 102         // some more data from the reader.
 103         // Branch using `>=` instead of the more correct `==`
 104         // to tell the compiler that the pos..cap slice is always valid.
 105         if self.pos >= self.filled {
 106             debug_assert!(self.pos == self.filled);
 107
 108             let mut buf = BorrowedBuf::from(&mut *self.buf);
 109             // SAFETY: `self.filled` bytes will always have been initialized.
 110             unsafe {
 111                 buf.set_init(self.initialized);
 112             }
 113
 114             reader.read_buf(buf.unfilled())?;
 115
 116             self.pos = 0;
 117             self.filled = buf.len();
 118             self.initialized = buf.init_len();
 119         }
 120         Ok(self.buffer())
 121     }
 122 }