Skip to content

Commit

Permalink
Avoid defensive re-initialization of the BufReader buffer
Browse files Browse the repository at this point in the history
  • Loading branch information
saethlin committed Oct 7, 2022
1 parent 0ca3565 commit 95ae993
Show file tree
Hide file tree
Showing 3 changed files with 48 additions and 3 deletions.
8 changes: 8 additions & 0 deletions library/std/src/io/buffered/bufreader.rs
Original file line number Diff line number Diff line change
Expand Up @@ -224,6 +224,14 @@ impl<R> BufReader<R> {
}
}

// This is only used by a test which asserts that the initialization-tracking is correct.
#[cfg(test)]
impl<R> BufReader<R> {
pub fn initialized(&self) -> usize {
self.buf.initialized()
}
}

impl<R: Seek> BufReader<R> {
/// Seeks relative to the current position. If the new position lies within the buffer,
/// the buffer will not be flushed, allowing for more efficient seeks.
Expand Down
19 changes: 16 additions & 3 deletions library/std/src/io/buffered/bufreader/buffer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -20,13 +20,19 @@ pub struct Buffer {
// 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,
// This is the max number of bytes returned across all `fill_buf` calls. We track this so that we
// can accurately tell `read_buf` how many bytes of buf are initialized, to bypass as much of its
// defensive initialization as possible. Note that while this often the same as `filled`, it
// doesn't need to be. Calls to `fill_buf` are not required to actually fill the buffer, and
// omitting this is a huge perf regression for `Read` impls that do not.
initialized: usize,
}

impl Buffer {
#[inline]
pub fn with_capacity(capacity: usize) -> Self {
let buf = Box::new_uninit_slice(capacity);
Self { buf, pos: 0, filled: 0 }
Self { buf, pos: 0, filled: 0, initialized: 0 }
}

#[inline]
Expand All @@ -51,6 +57,12 @@ impl Buffer {
self.pos
}

// This is only used by a test which asserts that the initialization-tracking is correct.
#[cfg(test)]
pub fn initialized(&self) -> usize {
self.initialized
}

#[inline]
pub fn discard_buffer(&mut self) {
self.pos = 0;
Expand Down Expand Up @@ -96,13 +108,14 @@ impl Buffer {
let mut buf = BorrowedBuf::from(&mut *self.buf);
// SAFETY: `self.filled` bytes will always have been initialized.
unsafe {
buf.set_init(self.filled);
buf.set_init(self.initialized);
}

reader.read_buf(buf.unfilled())?;

self.filled = buf.len();
self.pos = 0;
self.filled = buf.len();
self.initialized = buf.init_len();
}
Ok(self.buffer())
}
Expand Down
24 changes: 24 additions & 0 deletions library/std/src/io/buffered/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1039,3 +1039,27 @@ fn single_formatted_write() {
writeln!(&mut writer, "{}, {}!", "hello", "world").unwrap();
assert_eq!(writer.get_ref().events, [RecordedEvent::Write("hello, world!\n".to_string())]);
}

#[test]
fn bufreader_full_initialize() {
struct OneByteReader;
impl Read for OneByteReader {
fn read(&mut self, buf: &mut [u8]) -> crate::io::Result<usize> {
if buf.len() > 0 {
buf[0] = 0;
Ok(1)
} else {
Ok(0)
}
}
}
let mut reader = BufReader::new(OneByteReader);
// Nothing is initialized yet.
assert_eq!(reader.initialized(), 0);

let buf = reader.fill_buf().unwrap();
// We read one byte...
assert_eq!(buf.len(), 1);
// But we initialized the whole buffer!
assert_eq!(reader.initialized(), reader.capacity());
}

0 comments on commit 95ae993

Please sign in to comment.