Avoid defensive re-initialization of the BufReader buffer

3 files changed, 48 insertions, 3 deletions

diff --git a/library/std/src/io/buffered/bufreader.rs b/library/std/src/io/buffered/bufreader.rs
index 88ad92d8a98..4f339a18a48 100644
--- a/library/std/src/io/buffered/bufreader.rs
+++ b/library/std/src/io/buffered/bufreader.rs
@@ -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.
diff --git a/library/std/src/io/buffered/bufreader/buffer.rs b/library/std/src/io/buffered/bufreader/buffer.rs
index 867c22c6041..e9e29d60ca2 100644
--- a/library/std/src/io/buffered/bufreader/buffer.rs
+++ b/library/std/src/io/buffered/bufreader/buffer.rs
@@ -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]
@@ -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;
@@ -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())
     }
diff --git a/library/std/src/io/buffered/tests.rs b/library/std/src/io/buffered/tests.rs
index bd6d95242ad..f4e688eb926 100644
--- a/library/std/src/io/buffered/tests.rs
+++ b/library/std/src/io/buffered/tests.rs
@@ -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());
+}