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.
11use crate::cmp;
12use crate::io::{self, BorrowedBuf, Read};
13use crate::mem::MaybeUninit;
14
15pub 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
31impl 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}
123