std/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
12use crate::cmp;
13use crate::io::{self, BorrowedBuf, ErrorKind, Read};
14use crate::mem::MaybeUninit;
15
16pub struct Buffer {
17    // The buffer.
18    buf: Box<[MaybeUninit<u8>]>,
19    // The current seek offset into `buf`, must always be <= `filled`.
20    pos: usize,
21    // Each call to `fill_buf` sets `filled` to indicate how many bytes at the start of `buf` are
22    // initialized with bytes from a read.
23    filled: usize,
24}
25
26impl Buffer {
27    #[inline]
28    pub fn with_capacity(capacity: usize) -> Self {
29        let buf = Box::new_uninit_slice(capacity);
30        Self { buf, pos: 0, filled: 0 }
31    }
32
33    #[inline]
34    pub fn try_with_capacity(capacity: usize) -> io::Result<Self> {
35        match Box::try_new_uninit_slice(capacity) {
36            Ok(buf) => Ok(Self { buf, pos: 0, filled: 0 }),
37            Err(_) => {
38                Err(io::const_error!(ErrorKind::OutOfMemory, "failed to allocate read buffer"))
39            }
40        }
41    }
42
43    #[inline]
44    pub fn buffer(&self) -> &[u8] {
45        // SAFETY: self.pos and self.cap are valid, and self.cap => self.pos, and
46        // that region is initialized because those are all invariants of this type.
47        unsafe { self.buf.get_unchecked(self.pos..self.filled).assume_init_ref() }
48    }
49
50    #[inline]
51    pub fn capacity(&self) -> usize {
52        self.buf.len()
53    }
54
55    #[inline]
56    pub fn filled(&self) -> usize {
57        self.filled
58    }
59
60    #[inline]
61    pub fn pos(&self) -> usize {
62        self.pos
63    }
64
65    #[inline]
66    pub fn discard_buffer(&mut self) {
67        self.pos = 0;
68        self.filled = 0;
69    }
70
71    #[inline]
72    pub fn consume(&mut self, amt: usize) {
73        self.pos = cmp::min(self.pos + amt, self.filled);
74    }
75
76    /// If there are `amt` bytes available in the buffer, pass a slice containing those bytes to
77    /// `visitor` and return true. If there are not enough bytes available, return false.
78    #[inline]
79    pub fn consume_with<V>(&mut self, amt: usize, mut visitor: V) -> bool
80    where
81        V: FnMut(&[u8]),
82    {
83        if let Some(claimed) = self.buffer().get(..amt) {
84            visitor(claimed);
85            // If the indexing into self.buffer() succeeds, amt must be a valid increment.
86            self.pos += amt;
87            true
88        } else {
89            false
90        }
91    }
92
93    #[inline]
94    pub fn unconsume(&mut self, amt: usize) {
95        self.pos = self.pos.saturating_sub(amt);
96    }
97
98    /// Read more bytes into the buffer without discarding any of its contents
99    pub fn read_more(&mut self, mut reader: impl Read) -> io::Result<usize> {
100        let mut buf = BorrowedBuf::from(&mut self.buf[self.filled..]);
101        reader.read_buf(buf.unfilled())?;
102        self.filled += buf.len();
103        Ok(buf.len())
104    }
105
106    /// Remove bytes that have already been read from the buffer.
107    pub fn backshift(&mut self) {
108        self.buf.copy_within(self.pos..self.filled, 0);
109        self.filled -= self.pos;
110        self.pos = 0;
111    }
112
113    #[inline]
114    pub fn fill_buf(&mut self, mut reader: impl Read) -> io::Result<&[u8]> {
115        // If we've reached the end of our internal buffer then we need to fetch
116        // some more data from the reader.
117        // Branch using `>=` instead of the more correct `==`
118        // to tell the compiler that the pos..cap slice is always valid.
119        if self.pos >= self.filled {
120            debug_assert!(self.pos == self.filled);
121
122            let mut buf = BorrowedBuf::from(&mut *self.buf);
123            let result = reader.read_buf(buf.unfilled());
124
125            self.pos = 0;
126            self.filled = buf.len();
127
128            result?;
129        }
130        Ok(self.buffer())
131    }
132}