Auto merge of #116785 - nnethercote:spec-Bytes-read, r=the8472

bors · bors · commit e092502f6363 · 2023-10-19T23:31:58.000Z
Specialize `Bytes&lt;R&gt;::next` when `R` is a `BufReader`.

This reduces the runtime for a simple program using `Bytes::next` to iterate through a file from 220ms to 70ms on my Linux box.

r? `@the8472`
diff --git a/library/std/src/io/buffered/bufreader.rs b/library/std/src/io/buffered/bufreader.rs
@@ -2,7 +2,8 @@ mod buffer;
 
 use crate::fmt;
 use crate::io::{
-    self, BorrowedCursor, BufRead, IoSliceMut, Read, Seek, SeekFrom, SizeHint, DEFAULT_BUF_SIZE,
+    self, uninlined_slow_read_byte, BorrowedCursor, BufRead, IoSliceMut, Read, Seek, SeekFrom,
+    SizeHint, SpecReadByte, DEFAULT_BUF_SIZE,
 };
 use buffer::Buffer;
 
@@ -259,6 +260,22 @@ impl<R: ?Sized + Seek> BufReader<R> {
     }
 }
 
+impl<R> SpecReadByte for BufReader<R>
+where
+    Self: Read,
+{
+    #[inline]
+    fn spec_read_byte(&mut self) -> Option<io::Result<u8>> {
+        let mut byte = 0;
+        if self.buf.consume_with(1, |claimed| byte = claimed[0]) {
+            return Some(Ok(byte));
+        }
+
+        // Fallback case, only reached once per buffer refill.
+        uninlined_slow_read_byte(self)
+    }
+}
+
 #[stable(feature = "rust1", since = "1.0.0")]
 impl<R: ?Sized + Read> Read for BufReader<R> {
     fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
@@ -269,10 +286,8 @@ impl<R: ?Sized + Read> Read for BufReader<R> {
             self.discard_buffer();
             return self.inner.read(buf);
         }
-        let nread = {
-            let mut rem = self.fill_buf()?;
-            rem.read(buf)?
-        };
+        let mut rem = self.fill_buf()?;
+        let nread = rem.read(buf)?;
         self.consume(nread);
         Ok(nread)
     }
diff --git a/library/std/src/io/mod.rs b/library/std/src/io/mod.rs
@@ -2777,17 +2777,10 @@ pub struct Bytes<R> {
 impl<R: Read> Iterator for Bytes<R> {
     type Item = Result<u8>;
 
-    #[inline]
+    // Not `#[inline]`. This function gets inlined even without it, but having
+    // the inline annotation can result in worse code generation. See #116785.
     fn next(&mut self) -> Option<Result<u8>> {
-        let mut byte = 0;
-        loop {
-            return match self.inner.read(slice::from_mut(&mut byte)) {
-                Ok(0) => None,
-                Ok(..) => Some(Ok(byte)),
-                Err(ref e) if e.is_interrupted() => continue,
-                Err(e) => Some(Err(e)),
-            };
-        }
+        SpecReadByte::spec_read_byte(&mut self.inner)
     }
 
     #[inline]
@@ -2796,6 +2789,43 @@ impl<R: Read> Iterator for Bytes<R> {
     }
 }
 
+/// For the specialization of `Bytes::next`.
+trait SpecReadByte {
+    fn spec_read_byte(&mut self) -> Option<Result<u8>>;
+}
+
+impl<R> SpecReadByte for R
+where
+    Self: Read,
+{
+    #[inline]
+    default fn spec_read_byte(&mut self) -> Option<Result<u8>> {
+        inlined_slow_read_byte(self)
+    }
+}
+
+/// Read a single byte in a slow, generic way. This is used by the default
+/// `spec_read_byte`.
+#[inline]
+fn inlined_slow_read_byte<R: Read>(reader: &mut R) -> Option<Result<u8>> {
+    let mut byte = 0;
+    loop {
+        return match reader.read(slice::from_mut(&mut byte)) {
+            Ok(0) => None,
+            Ok(..) => Some(Ok(byte)),
+            Err(ref e) if e.is_interrupted() => continue,
+            Err(e) => Some(Err(e)),
+        };
+    }
+}
+
+// Used by `BufReader::spec_read_byte`, for which the `inline(ever)` is
+// important.
+#[inline(never)]
+fn uninlined_slow_read_byte<R: Read>(reader: &mut R) -> Option<Result<u8>> {
+    inlined_slow_read_byte(reader)
+}
+
 trait SizeHint {
     fn lower_bound(&self) -> usize;
 
