Auto merge of #58357 - sfackler:vectored-io, r=alexcrichton

Add vectored read and write support

This functionality has lived for a while in the tokio ecosystem, where
it can improve performance by minimizing copies.

r? @alexcrichton

4 files changed, 144 insertions, 14 deletions

diff --git a/src/libstd/sys/windows/c.rs b/src/libstd/sys/windows/c.rs
index 28fd4df386e..a78b599204b 100644
--- a/src/libstd/sys/windows/c.rs
+++ b/src/libstd/sys/windows/c.rs
@@ -57,6 +57,9 @@ pub type LPWSAPROTOCOL_INFO = *mut WSAPROTOCOL_INFO;
 pub type LPSTR = *mut CHAR;
 pub type LPWSTR = *mut WCHAR;
 pub type LPFILETIME = *mut FILETIME;
+pub type LPWSABUF = *mut WSABUF;
+pub type LPWSAOVERLAPPED = *mut c_void;
+pub type LPWSAOVERLAPPED_COMPLETION_ROUTINE = *mut c_void;
 
 pub type PCONDITION_VARIABLE = *mut CONDITION_VARIABLE;
 pub type PLARGE_INTEGER = *mut c_longlong;
@@ -325,6 +328,12 @@ pub struct WSADATA {
 }
 
 #[repr(C)]
+pub struct WSABUF {
+    pub len: ULONG,
+    pub buf: *mut CHAR,
+}
+
+#[repr(C)]
 pub struct WSAPROTOCOL_INFO {
     pub dwServiceFlags1: DWORD,
     pub dwServiceFlags2: DWORD,
@@ -988,6 +997,22 @@ extern "system" {
                                dwProcessId: DWORD,
                                lpProtocolInfo: LPWSAPROTOCOL_INFO)
                                -> c_int;
+    pub fn WSASend(s: SOCKET,
+                   lpBuffers: LPWSABUF,
+                   dwBufferCount: DWORD,
+                   lpNumberOfBytesSent: LPDWORD,
+                   dwFlags: DWORD,
+                   lpOverlapped: LPWSAOVERLAPPED,
+                   lpCompletionRoutine: LPWSAOVERLAPPED_COMPLETION_ROUTINE)
+                   -> c_int;
+    pub fn WSARecv(s: SOCKET,
+                   lpBuffers: LPWSABUF,
+                   dwBufferCount: DWORD,
+                   lpNumberOfBytesRecvd: LPDWORD,
+                   lpFlags: LPDWORD,
+                   lpOverlapped: LPWSAOVERLAPPED,
+                   lpCompletionRoutine: LPWSAOVERLAPPED_COMPLETION_ROUTINE)
+                   -> c_int;
     pub fn GetCurrentProcessId() -> DWORD;
     pub fn WSASocketW(af: c_int,
                       kind: c_int,
diff --git a/src/libstd/sys/windows/io.rs b/src/libstd/sys/windows/io.rs
new file mode 100644
index 00000000000..662e3047923
--- /dev/null
+++ b/src/libstd/sys/windows/io.rs
@@ -0,0 +1,63 @@
+use marker::PhantomData;
+use slice;
+use sys::c;
+
+#[repr(transparent)]
+pub struct IoVec<'a> {
+    vec: c::WSABUF,
+    _p: PhantomData<&'a [u8]>,
+}
+
+impl<'a> IoVec<'a> {
+    #[inline]
+    pub fn new(buf: &'a [u8]) -> IoVec<'a> {
+        assert!(buf.len() <= c::ULONG::max_value() as usize);
+        IoVec {
+            vec: c::WSABUF {
+                len: buf.len() as c::ULONG,
+                buf: buf.as_ptr() as *mut u8 as *mut c::CHAR,
+            },
+            _p: PhantomData,
+        }
+    }
+
+    #[inline]
+    pub fn as_slice(&self) -> &[u8] {
+        unsafe {
+            slice::from_raw_parts(self.vec.buf as *mut u8, self.vec.len as usize)
+        }
+    }
+}
+
+pub struct IoVecMut<'a> {
+    vec: c::WSABUF,
+    _p: PhantomData<&'a mut [u8]>,
+}
+
+impl<'a> IoVecMut<'a> {
+    #[inline]
+    pub fn new(buf: &'a mut [u8]) -> IoVecMut<'a> {
+        assert!(buf.len() <= c::ULONG::max_value() as usize);
+        IoVecMut {
+            vec: c::WSABUF {
+                len: buf.len() as c::ULONG,
+                buf: buf.as_mut_ptr() as *mut c::CHAR,
+            },
+            _p: PhantomData,
+        }
+    }
+
+    #[inline]
+    pub fn as_slice(&self) -> &[u8] {
+        unsafe {
+            slice::from_raw_parts(self.vec.buf as *mut u8, self.vec.len as usize)
+        }
+    }
+
+    #[inline]
+    pub fn as_mut_slice(&mut self) -> &mut [u8] {
+        unsafe {
+            slice::from_raw_parts_mut(self.vec.buf as *mut u8, self.vec.len as usize)
+        }
+    }
+}
diff --git a/src/libstd/sys/windows/mod.rs b/src/libstd/sys/windows/mod.rs
index e97e436efbf..56c76a169fe 100644
--- a/src/libstd/sys/windows/mod.rs
+++ b/src/libstd/sys/windows/mod.rs
@@ -2,7 +2,7 @@
 
 use ptr;
 use ffi::{OsStr, OsString};
-use io::{self, ErrorKind};
+use io::ErrorKind;
 use os::windows::ffi::{OsStrExt, OsStringExt};
 use path::PathBuf;
 use time::Duration;
@@ -26,6 +26,7 @@ pub mod ext;
 pub mod fast_thread_local;
 pub mod fs;
 pub mod handle;
+pub mod io;
 pub mod memchr;
 pub mod mutex;
 pub mod net;
@@ -75,12 +76,12 @@ pub fn decode_error_kind(errno: i32) -> ErrorKind {
     }
 }
 
-pub fn to_u16s<S: AsRef<OsStr>>(s: S) -> io::Result<Vec<u16>> {
-    fn inner(s: &OsStr) -> io::Result<Vec<u16>> {
+pub fn to_u16s<S: AsRef<OsStr>>(s: S) -> ::io::Result<Vec<u16>> {
+    fn inner(s: &OsStr) -> ::io::Result<Vec<u16>> {
         let mut maybe_result: Vec<u16> = s.encode_wide().collect();
         if maybe_result.iter().any(|&u| u == 0) {
-            return Err(io::Error::new(io::ErrorKind::InvalidInput,
-                                      "strings passed to WinAPI cannot contain NULs"));
+            return Err(::io::Error::new(::io::ErrorKind::InvalidInput,
+                                        "strings passed to WinAPI cannot contain NULs"));
         }
         maybe_result.push(0);
         Ok(maybe_result)
@@ -102,7 +103,7 @@ pub fn to_u16s<S: AsRef<OsStr>>(s: S) -> io::Result<Vec<u16>> {
 // Once the syscall has completed (errors bail out early) the second closure is
 // yielded the data which has been read from the syscall. The return value
 // from this closure is then the return value of the function.
-fn fill_utf16_buf<F1, F2, T>(mut f1: F1, f2: F2) -> io::Result<T>
+fn fill_utf16_buf<F1, F2, T>(mut f1: F1, f2: F2) -> ::io::Result<T>
     where F1: FnMut(*mut u16, c::DWORD) -> c::DWORD,
           F2: FnOnce(&[u16]) -> T
 {
@@ -134,7 +135,7 @@ fn fill_utf16_buf<F1, F2, T>(mut f1: F1, f2: F2) -> io::Result<T>
             c::SetLastError(0);
             let k = match f1(buf.as_mut_ptr(), n as c::DWORD) {
                 0 if c::GetLastError() == 0 => 0,
-                0 => return Err(io::Error::last_os_error()),
+                0 => return Err(::io::Error::last_os_error()),
                 n => n,
             } as usize;
             if k == n && c::GetLastError() == c::ERROR_INSUFFICIENT_BUFFER {
@@ -157,7 +158,7 @@ fn wide_char_to_multi_byte(code_page: u32,
                            flags: u32,
                            s: &[u16],
                            no_default_char: bool)
-                           -> io::Result<Vec<i8>> {
+                           -> ::io::Result<Vec<i8>> {
     unsafe {
         let mut size = c::WideCharToMultiByte(code_page,
                                               flags,
@@ -168,7 +169,7 @@ fn wide_char_to_multi_byte(code_page: u32,
                                               ptr::null(),
                                               ptr::null_mut());
         if size == 0 {
-            return Err(io::Error::last_os_error());
+            return Err(::io::Error::last_os_error());
         }
 
         let mut buf = Vec::with_capacity(size as usize);
@@ -185,10 +186,10 @@ fn wide_char_to_multi_byte(code_page: u32,
                                       if no_default_char { &mut used_default_char }
                                       else { ptr::null_mut() });
         if size == 0 {
-            return Err(io::Error::last_os_error());
+            return Err(::io::Error::last_os_error());
         }
         if no_default_char && used_default_char == c::TRUE {
-            return Err(io::Error::new(io::ErrorKind::InvalidData,
+            return Err(::io::Error::new(::io::ErrorKind::InvalidData,
                                       "string cannot be converted to requested code page"));
         }
 
@@ -220,9 +221,9 @@ macro_rules! impl_is_zero {
 
 impl_is_zero! { i8 i16 i32 i64 isize u8 u16 u32 u64 usize }
 
-pub fn cvt<I: IsZero>(i: I) -> io::Result<I> {
+pub fn cvt<I: IsZero>(i: I) -> ::io::Result<I> {
     if i.is_zero() {
-        Err(io::Error::last_os_error())
+        Err(::io::Error::last_os_error())
     } else {
         Ok(i)
     }
diff --git a/src/libstd/sys/windows/net.rs b/src/libstd/sys/windows/net.rs
index acda81dcde5..76be26a9d1a 100644
--- a/src/libstd/sys/windows/net.rs
+++ b/src/libstd/sys/windows/net.rs
@@ -1,7 +1,7 @@
 #![unstable(issue = "0", feature = "windows_net")]
 
 use cmp;
-use io::{self, Read};
+use io::{self, Read, IoVec, IoVecMut};
 use libc::{c_int, c_void, c_ulong, c_long};
 use mem;
 use net::{SocketAddr, Shutdown};
@@ -207,6 +207,30 @@ impl Socket {
         self.recv_with_flags(buf, 0)
     }
 
+    pub fn read_vectored(&self, bufs: &mut [IoVecMut<'_>]) -> io::Result<usize> {
+        // On unix when a socket is shut down all further reads return 0, so we
+        // do the same on windows to map a shut down socket to returning EOF.
+        let len = cmp::min(bufs.len(), c::DWORD::max_value() as usize) as c::DWORD;
+        let mut nread = 0;
+        let mut flags = 0;
+        unsafe {
+            let ret = c::WSARecv(
+                self.0,
+                bufs.as_mut_ptr() as *mut c::WSABUF,
+                len,
+                &mut nread,
+                &mut flags,
+                ptr::null_mut(),
+                ptr::null_mut(),
+            );
+            match ret {
+                0 => Ok(nread as usize),
+                _ if c::WSAGetLastError() == c::WSAESHUTDOWN => Ok(0),
+                _ => Err(last_error()),
+            }
+        }
+    }
+
     pub fn peek(&self, buf: &mut [u8]) -> io::Result<usize> {
         self.recv_with_flags(buf, c::MSG_PEEK)
     }
@@ -243,6 +267,23 @@ impl Socket {
         self.recv_from_with_flags(buf, c::MSG_PEEK)
     }
 
+    pub fn write_vectored(&self, bufs: &[IoVec<'_>]) -> io::Result<usize> {
+        let len = cmp::min(bufs.len(), c::DWORD::max_value() as usize) as c::DWORD;
+        let mut nwritten = 0;
+        unsafe {
+            cvt(c::WSASend(
+                self.0,
+                bufs.as_ptr() as *const c::WSABUF as *mut c::WSABUF,
+                len,
+                &mut nwritten,
+                0,
+                ptr::null_mut(),
+                ptr::null_mut(),
+            ))?;
+        }
+        Ok(nwritten as usize)
+    }
+
     pub fn set_timeout(&self, dur: Option<Duration>,
                        kind: c_int) -> io::Result<()> {
         let timeout = match dur {