Rollup merge of #108326 - tmiasko:read-buf, r=thomcc

Implement read_buf for a few more types

Implement read_buf for TcpStream, Stdin, StdinLock, ChildStdout,
ChildStderr (and internally for AnonPipe, Handle, Socket), so
that it skips buffer initialization.

The other provided methods like read_to_string and read_to_end are
implemented in terms of read_buf and so benefit from the optimization
as well.

This commit also implements read_vectored and is_read_vectored where
applicable.

library/std/src/fs/tests.rs

@@ -2,7 +2,8 @@ use crate::io::prelude::*;
 
 use crate::env;
 use crate::fs::{self, File, OpenOptions};
-use crate::io::{ErrorKind, SeekFrom};
+use crate::io::{BorrowedBuf, ErrorKind, SeekFrom};
+use crate::mem::MaybeUninit;
 use crate::path::Path;
 use crate::str;
 use crate::sync::Arc;
@@ -401,6 +402,23 @@ fn file_test_io_seek_read_write() {
     check!(fs::remove_file(&filename));
 }
 
+#[test]
+fn file_test_read_buf() {
+    let tmpdir = tmpdir();
+    let filename = &tmpdir.join("test");
+    check!(fs::write(filename, &[1, 2, 3, 4]));
+
+    let mut buf: [MaybeUninit<u8>; 128] = MaybeUninit::uninit_array();
+    let mut buf = BorrowedBuf::from(buf.as_mut_slice());
+    let mut file = check!(File::open(filename));
+    check!(file.read_buf(buf.unfilled()));
+    assert_eq!(buf.filled(), &[1, 2, 3, 4]);
+    // File::read_buf should omit buffer initialization.
+    assert_eq!(buf.init_len(), 4);
+
+    check!(fs::remove_file(filename));
+}
+
 #[test]
 fn file_test_stat_is_correct_on_is_file() {
     let tmpdir = tmpdir();

library/std/src/io/stdio.rs

@@ -8,7 +8,7 @@ use crate::io::prelude::*;
 use crate::cell::{Cell, RefCell};
 use crate::fmt;
 use crate::fs::File;
-use crate::io::{self, BufReader, IoSlice, IoSliceMut, LineWriter, Lines};
+use crate::io::{self, BorrowedCursor, BufReader, IoSlice, IoSliceMut, LineWriter, Lines};
 use crate::sync::atomic::{AtomicBool, Ordering};
 use crate::sync::{Arc, Mutex, MutexGuard, OnceLock, ReentrantMutex, ReentrantMutexGuard};
 use crate::sys::stdio;
@@ -97,6 +97,10 @@ impl Read for StdinRaw {
         handle_ebadf(self.0.read(buf), 0)
     }
 
+    fn read_buf(&mut self, buf: BorrowedCursor<'_>) -> io::Result<()> {
+        handle_ebadf(self.0.read_buf(buf), ())
+    }
+
     fn read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
         handle_ebadf(self.0.read_vectored(bufs), 0)
     }
@@ -418,6 +422,9 @@ impl Read for Stdin {
     fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
         self.lock().read(buf)
     }
+    fn read_buf(&mut self, buf: BorrowedCursor<'_>) -> io::Result<()> {
+        self.lock().read_buf(buf)
+    }
     fn read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
         self.lock().read_vectored(bufs)
     }
@@ -450,6 +457,10 @@ impl Read for StdinLock<'_> {
         self.inner.read(buf)
     }
 
+    fn read_buf(&mut self, buf: BorrowedCursor<'_>) -> io::Result<()> {
+        self.inner.read_buf(buf)
+    }
+
     fn read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
         self.inner.read_vectored(bufs)
     }

library/std/src/net/tcp.rs

@@ -6,7 +6,7 @@ mod tests;
 use crate::io::prelude::*;
 
 use crate::fmt;
-use crate::io::{self, IoSlice, IoSliceMut};
+use crate::io::{self, BorrowedCursor, IoSlice, IoSliceMut};
 use crate::iter::FusedIterator;
 use crate::net::{Shutdown, SocketAddr, ToSocketAddrs};
 use crate::sys_common::net as net_imp;
@@ -619,6 +619,10 @@ impl Read for TcpStream {
         self.0.read(buf)
     }
 
+    fn read_buf(&mut self, buf: BorrowedCursor<'_>) -> io::Result<()> {
+        self.0.read_buf(buf)
+    }
+
     fn read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
         self.0.read_vectored(bufs)
     }
@@ -653,6 +657,10 @@ impl Read for &TcpStream {
         self.0.read(buf)
     }
 
+    fn read_buf(&mut self, buf: BorrowedCursor<'_>) -> io::Result<()> {
+        self.0.read_buf(buf)
+    }
+
     fn read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
         self.0.read_vectored(bufs)
     }

library/std/src/net/tcp/tests.rs

@@ -1,6 +1,7 @@
 use crate::fmt;
 use crate::io::prelude::*;
-use crate::io::{ErrorKind, IoSlice, IoSliceMut};
+use crate::io::{BorrowedBuf, ErrorKind, IoSlice, IoSliceMut};
+use crate::mem::MaybeUninit;
 use crate::net::test::{next_test_ip4, next_test_ip6};
 use crate::net::*;
 use crate::sync::mpsc::channel;
@@ -279,6 +280,31 @@ fn partial_read() {
     })
 }
 
+#[test]
+fn read_buf() {
+    each_ip(&mut |addr| {
+        let srv = t!(TcpListener::bind(&addr));
+        let t = thread::spawn(move || {
+            let mut s = t!(TcpStream::connect(&addr));
+            s.write_all(&[1, 2, 3, 4]).unwrap();
+        });
+
+        let mut s = t!(srv.accept()).0;
+        let mut buf: [MaybeUninit<u8>; 128] = MaybeUninit::uninit_array();
+        let mut buf = BorrowedBuf::from(buf.as_mut_slice());
+        t!(s.read_buf(buf.unfilled()));
+        assert_eq!(buf.filled(), &[1, 2, 3, 4]);
+
+        // FIXME: sgx uses default_read_buf that initializes the buffer.
+        if cfg!(not(target_env = "sgx")) {
+            // TcpStream::read_buf should omit buffer initialization.
+            assert_eq!(buf.init_len(), 4);
+        }
+
+        t.join().ok().expect("thread panicked");
+    })
+}
+
 #[test]
 fn read_vectored() {
     each_ip(&mut |addr| {

library/std/src/process.rs

@@ -110,7 +110,7 @@ use crate::convert::Infallible;
 use crate::ffi::OsStr;
 use crate::fmt;
 use crate::fs;
-use crate::io::{self, IoSlice, IoSliceMut};
+use crate::io::{self, BorrowedCursor, IoSlice, IoSliceMut};
 use crate::num::NonZeroI32;
 use crate::path::Path;
 use crate::str;
@@ -354,6 +354,10 @@ impl Read for ChildStdout {
         self.inner.read(buf)
     }
 
+    fn read_buf(&mut self, buf: BorrowedCursor<'_>) -> io::Result<()> {
+        self.inner.read_buf(buf)
+    }
+
     fn read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
         self.inner.read_vectored(bufs)
     }
@@ -419,6 +423,10 @@ impl Read for ChildStderr {
         self.inner.read(buf)
     }
 
+    fn read_buf(&mut self, buf: BorrowedCursor<'_>) -> io::Result<()> {
+        self.inner.read_buf(buf)
+    }
+
     fn read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
         self.inner.read_vectored(bufs)
     }

library/std/src/process/tests.rs

@@ -1,7 +1,8 @@
 use crate::io::prelude::*;
 
 use super::{Command, Output, Stdio};
-use crate::io::ErrorKind;
+use crate::io::{BorrowedBuf, ErrorKind};
+use crate::mem::MaybeUninit;
 use crate::str;
 
 fn known_command() -> Command {
@@ -119,6 +120,37 @@ fn stdin_works() {
     assert_eq!(out, "foobar\n");
 }
 
+#[test]
+#[cfg_attr(any(target_os = "vxworks"), ignore)]
+fn child_stdout_read_buf() {
+    let mut cmd = if cfg!(target_os = "windows") {
+        let mut cmd = Command::new("cmd");
+        cmd.arg("/C").arg("echo abc");
+        cmd
+    } else {
+        let mut cmd = shell_cmd();
+        cmd.arg("-c").arg("echo abc");
+        cmd
+    };
+    cmd.stdin(Stdio::null());
+    cmd.stdout(Stdio::piped());
+    let child = cmd.spawn().unwrap();
+
+    let mut stdout = child.stdout.unwrap();
+    let mut buf: [MaybeUninit<u8>; 128] = MaybeUninit::uninit_array();
+    let mut buf = BorrowedBuf::from(buf.as_mut_slice());
+    stdout.read_buf(buf.unfilled()).unwrap();
+
+    // ChildStdout::read_buf should omit buffer initialization.
+    if cfg!(target_os = "windows") {
+        assert_eq!(buf.filled(), b"abc\r\n");
+        assert_eq!(buf.init_len(), 5);
+    } else {
+        assert_eq!(buf.filled(), b"abc\n");
+        assert_eq!(buf.init_len(), 4);
+    };
+}
+
 #[test]
 #[cfg_attr(any(target_os = "vxworks"), ignore)]
 fn test_process_status() {

library/std/src/sys/sgx/fd.rs

@@ -1,7 +1,7 @@
 use fortanix_sgx_abi::Fd;
 
 use super::abi::usercalls;
-use crate::io::{self, IoSlice, IoSliceMut};
+use crate::io::{self, BorrowedCursor, IoSlice, IoSliceMut};
 use crate::mem;
 use crate::sys::{AsInner, FromInner, IntoInner};
 
@@ -30,6 +30,10 @@ impl FileDesc {
         usercalls::read(self.fd, &mut [IoSliceMut::new(buf)])
     }
 
+    pub fn read_buf(&self, buf: BorrowedCursor<'_>) -> io::Result<()> {
+        crate::io::default_read_buf(|b| self.read(b), buf)
+    }
+
     pub fn read_vectored(&self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
         usercalls::read(self.fd, bufs)
     }

library/std/src/sys/sgx/net.rs

@@ -1,6 +1,6 @@
 use crate::error;
 use crate::fmt;
-use crate::io::{self, IoSlice, IoSliceMut};
+use crate::io::{self, BorrowedCursor, IoSlice, IoSliceMut};
 use crate::net::{Ipv4Addr, Ipv6Addr, Shutdown, SocketAddr, ToSocketAddrs};
 use crate::sync::Arc;
 use crate::sys::fd::FileDesc;
@@ -144,6 +144,10 @@ impl TcpStream {
         self.inner.inner.read(buf)
     }
 
+    pub fn read_buf(&self, buf: BorrowedCursor<'_>) -> io::Result<()> {
+        self.inner.inner.read_buf(buf)
+    }
+
     pub fn read_vectored(&self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
         self.inner.inner.read_vectored(bufs)
     }

library/std/src/sys/unix/fd.rs

@@ -469,6 +469,15 @@ impl<'a> Read for &'a FileDesc {
     fn read_buf(&mut self, cursor: BorrowedCursor<'_>) -> io::Result<()> {
         (**self).read_buf(cursor)
     }
+
+    fn read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
+        (**self).read_vectored(bufs)
+    }
+
+    #[inline]
+    fn is_read_vectored(&self) -> bool {
+        (**self).is_read_vectored()
+    }
 }
 
 impl AsInner<OwnedFd> for FileDesc {

library/std/src/sys/unix/net.rs

@@ -1,6 +1,6 @@
 use crate::cmp;
 use crate::ffi::CStr;
-use crate::io::{self, IoSlice, IoSliceMut};
+use crate::io::{self, BorrowedBuf, BorrowedCursor, IoSlice, IoSliceMut};
 use crate::mem;
 use crate::net::{Shutdown, SocketAddr};
 use crate::os::unix::io::{AsFd, AsRawFd, BorrowedFd, FromRawFd, IntoRawFd, RawFd};
@@ -242,19 +242,35 @@ impl Socket {
         self.0.duplicate().map(Socket)
     }
 
-    fn recv_with_flags(&self, buf: &mut [u8], flags: c_int) -> io::Result<usize> {
+    fn recv_with_flags(&self, mut buf: BorrowedCursor<'_>, flags: c_int) -> io::Result<()> {
         let ret = cvt(unsafe {
-            libc::recv(self.as_raw_fd(), buf.as_mut_ptr() as *mut c_void, buf.len(), flags)
+            libc::recv(
+                self.as_raw_fd(),
+                buf.as_mut().as_mut_ptr() as *mut c_void,
+                buf.capacity(),
+                flags,
+            )
         })?;
-        Ok(ret as usize)
+        unsafe {
+            buf.advance(ret as usize);
+        }
+        Ok(())
     }
 
     pub fn read(&self, buf: &mut [u8]) -> io::Result<usize> {
-        self.recv_with_flags(buf, 0)
+        let mut buf = BorrowedBuf::from(buf);
+        self.recv_with_flags(buf.unfilled(), 0)?;
+        Ok(buf.len())
     }
 
     pub fn peek(&self, buf: &mut [u8]) -> io::Result<usize> {
-        self.recv_with_flags(buf, MSG_PEEK)
+        let mut buf = BorrowedBuf::from(buf);
+        self.recv_with_flags(buf.unfilled(), MSG_PEEK)?;
+        Ok(buf.len())
+    }
+
+    pub fn read_buf(&self, buf: BorrowedCursor<'_>) -> io::Result<()> {
+        self.recv_with_flags(buf, 0)
     }
 
     pub fn read_vectored(&self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {

library/std/src/sys/unix/pipe.rs

@@ -1,4 +1,4 @@
-use crate::io::{self, IoSlice, IoSliceMut};
+use crate::io::{self, BorrowedCursor, IoSlice, IoSliceMut};
 use crate::mem;
 use crate::os::unix::io::{AsFd, AsRawFd, BorrowedFd, FromRawFd, IntoRawFd, RawFd};
 use crate::sys::fd::FileDesc;
@@ -49,6 +49,10 @@ impl AnonPipe {
         self.0.read(buf)
     }
 
+    pub fn read_buf(&self, buf: BorrowedCursor<'_>) -> io::Result<()> {
+        self.0.read_buf(buf)
+    }
+
     pub fn read_vectored(&self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
         self.0.read_vectored(bufs)
     }

library/std/src/sys/unix/stdio.rs

@@ -1,4 +1,4 @@
-use crate::io::{self, IoSlice, IoSliceMut};
+use crate::io::{self, BorrowedCursor, IoSlice, IoSliceMut};
 use crate::mem::ManuallyDrop;
 use crate::os::unix::io::FromRawFd;
 use crate::sys::fd::FileDesc;
@@ -18,6 +18,10 @@ impl io::Read for Stdin {
         unsafe { ManuallyDrop::new(FileDesc::from_raw_fd(libc::STDIN_FILENO)).read(buf) }
     }
 
+    fn read_buf(&mut self, buf: BorrowedCursor<'_>) -> io::Result<()> {
+        unsafe { ManuallyDrop::new(FileDesc::from_raw_fd(libc::STDIN_FILENO)).read_buf(buf) }
+    }
+
     fn read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
         unsafe { ManuallyDrop::new(FileDesc::from_raw_fd(libc::STDIN_FILENO)).read_vectored(bufs) }
     }

library/std/src/sys/unsupported/net.rs

@@ -1,5 +1,5 @@
 use crate::fmt;
-use crate::io::{self, IoSlice, IoSliceMut};
+use crate::io::{self, BorrowedCursor, IoSlice, IoSliceMut};
 use crate::net::{Ipv4Addr, Ipv6Addr, Shutdown, SocketAddr};
 use crate::sys::unsupported;
 use crate::time::Duration;
@@ -39,6 +39,10 @@ impl TcpStream {
         self.0
     }
 
+    pub fn read_buf(&self, _buf: BorrowedCursor<'_>) -> io::Result<()> {
+        self.0
+    }
+
     pub fn read_vectored(&self, _: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
         self.0
     }

library/std/src/sys/unsupported/pipe.rs

@@ -1,4 +1,4 @@
-use crate::io::{self, IoSlice, IoSliceMut};
+use crate::io::{self, BorrowedCursor, IoSlice, IoSliceMut};
 
 pub struct AnonPipe(!);
 
@@ -7,6 +7,10 @@ impl AnonPipe {
         self.0
     }
 
+    pub fn read_buf(&self, _buf: BorrowedCursor<'_>) -> io::Result<()> {
+        self.0
+    }
+
     pub fn read_vectored(&self, _bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
         self.0
     }