-
Notifications
You must be signed in to change notification settings - Fork 13k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Auto merge of #98748 - saethlin:optimize-bufreader, r=Mark-Simulacrum
Remove some redundant checks from BufReader The implementation of BufReader contains a lot of redundant checks. While any one of these checks is not particularly expensive to execute, especially when taken together they dramatically inhibit LLVM's ability to make subsequent optimizations by confusing data flow increasing the code size of anything that uses BufReader. In particular, these changes have a ~2x increase on the benchmark that this adds a `black_box` to. I'm adding that `black_box` here just in case LLVM gets clever enough to remove the reads entirely. Right now it can't, but these optimizations are really setting it up to do so. We get this optimization by factoring all the actual buffer management and bounds-checking logic into a new module inside `bufreader` with a new `Buffer` type. This makes it much easier to ensure that we have correctly encapsulated the management of the region of the buffer that we have read bytes into, and it lets us provide a new faster way to do small reads. `Buffer::consume_with` lets a caller do a read from the buffer with a single bounds check, instead of the double-check that's required to use `buffer` + `consume`. Unfortunately I'm not aware of a lot of open-source usage of `BufReader` in perf-critical environments. Some time ago I tweaked this code because I saw `BufReader` in a profile at work, and I contributed some benchmarks to the `bincode` crate which exercise `BufReader::buffer`. These changes appear to help those benchmarks at little, but all these sorts of benchmarks are kind of fragile so I'm wary of quoting anything specific.
- Loading branch information
Showing
3 changed files
with
134 additions
and
53 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,105 @@ | ||
| ///! An encapsulation of `BufReader`'s buffer management logic. | ||
| /// | ||
| /// This module factors out the basic functionality of `BufReader` in order to protect two core | ||
| /// invariants: | ||
| /// * `filled` bytes of `buf` are always initialized | ||
| /// * `pos` is always <= `filled` | ||
| /// Since this module encapsulates the buffer management logic, we can ensure that the range | ||
| /// `pos..filled` is always a valid index into the initialized region of the buffer. This means | ||
| /// that user code which wants to do reads from a `BufReader` via `buffer` + `consume` can do so | ||
| /// without encountering any runtime bounds checks. | ||
| use crate::cmp; | ||
| use crate::io::{self, Read, ReadBuf}; | ||
| use crate::mem::MaybeUninit; | ||
|
|
||
| pub struct Buffer { | ||
| // The buffer. | ||
| buf: Box<[MaybeUninit<u8>]>, | ||
| // The current seek offset into `buf`, must always be <= `filled`. | ||
| pos: usize, | ||
| // Each call to `fill_buf` sets `filled` to indicate how many bytes at the start of `buf` are | ||
| // initialized with bytes from a read. | ||
| filled: usize, | ||
| } | ||
|
|
||
| impl Buffer { | ||
| #[inline] | ||
| pub fn with_capacity(capacity: usize) -> Self { | ||
| let buf = Box::new_uninit_slice(capacity); | ||
| Self { buf, pos: 0, filled: 0 } | ||
| } | ||
|
|
||
| #[inline] | ||
| pub fn buffer(&self) -> &[u8] { | ||
| // SAFETY: self.pos and self.cap are valid, and self.cap => self.pos, and | ||
| // that region is initialized because those are all invariants of this type. | ||
| unsafe { MaybeUninit::slice_assume_init_ref(self.buf.get_unchecked(self.pos..self.filled)) } | ||
| } | ||
|
|
||
| #[inline] | ||
| pub fn capacity(&self) -> usize { | ||
| self.buf.len() | ||
| } | ||
|
|
||
| #[inline] | ||
| pub fn filled(&self) -> usize { | ||
| self.filled | ||
| } | ||
|
|
||
| #[inline] | ||
| pub fn pos(&self) -> usize { | ||
| self.pos | ||
| } | ||
|
|
||
| #[inline] | ||
| pub fn discard_buffer(&mut self) { | ||
| self.pos = 0; | ||
| self.filled = 0; | ||
| } | ||
|
|
||
| #[inline] | ||
| pub fn consume(&mut self, amt: usize) { | ||
| self.pos = cmp::min(self.pos + amt, self.filled); | ||
| } | ||
|
|
||
| /// If there are `amt` bytes available in the buffer, pass a slice containing those bytes to | ||
| /// `visitor` and return true. If there are not enough bytes available, return false. | ||
| #[inline] | ||
| pub fn consume_with<V>(&mut self, amt: usize, mut visitor: V) -> bool | ||
| where | ||
| V: FnMut(&[u8]), | ||
| { | ||
| if let Some(claimed) = self.buffer().get(..amt) { | ||
| visitor(claimed); | ||
| // If the indexing into self.buffer() succeeds, amt must be a valid increment. | ||
| self.pos += amt; | ||
| true | ||
| } else { | ||
| false | ||
| } | ||
| } | ||
|
|
||
| #[inline] | ||
| pub fn unconsume(&mut self, amt: usize) { | ||
| self.pos = self.pos.saturating_sub(amt); | ||
| } | ||
|
|
||
| #[inline] | ||
| pub fn fill_buf(&mut self, mut reader: impl Read) -> io::Result<&[u8]> { | ||
| // If we've reached the end of our internal buffer then we need to fetch | ||
| // some more data from the reader. | ||
| // Branch using `>=` instead of the more correct `==` | ||
| // to tell the compiler that the pos..cap slice is always valid. | ||
| if self.pos >= self.filled { | ||
| debug_assert!(self.pos == self.filled); | ||
|
|
||
| let mut readbuf = ReadBuf::uninit(&mut self.buf); | ||
|
|
||
| reader.read_buf(&mut readbuf)?; | ||
|
|
||
| self.filled = readbuf.filled_len(); | ||
| self.pos = 0; | ||
| } | ||
| Ok(self.buffer()) | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters