Add {IoSlice, IoSliceMut}::advance

9 files changed, 351 insertions, 1 deletions

diff --git a/src/libstd/io/mod.rs b/src/libstd/io/mod.rs
index f89a7144437..f2b6ce6feb2 100644
--- a/src/libstd/io/mod.rs
+++ b/src/libstd/io/mod.rs
@@ -942,6 +942,62 @@ impl<'a> IoSliceMut<'a> {
     pub fn new(buf: &'a mut [u8]) -> IoSliceMut<'a> {
         IoSliceMut(sys::io::IoSliceMut::new(buf))
     }
+
+    /// Advance the internal cursor of the slice.
+    ///
+    /// # Notes
+    ///
+    /// Elements in the slice may be modified if the cursor is not advanced to
+    /// the end of the slice. For example if we have a slice of buffers with 2
+    /// `IoSliceMut`s, both of length 8, and we advance the cursor by 10 bytes
+    /// the first `IoSliceMut` will be untouched however the second will be
+    /// modified to remove the first 2 bytes (10 - 8).
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(io_slice_advance)]
+    ///
+    /// use std::io::IoSliceMut;
+    /// use std::mem;
+    /// use std::ops::Deref;
+    ///
+    /// let mut buf1 = [1; 8];
+    /// let mut buf2 = [2; 16];
+    /// let mut buf3 = [3; 8];
+    /// let mut bufs = &mut [
+    ///     IoSliceMut::new(&mut buf1),
+    ///     IoSliceMut::new(&mut buf2),
+    ///     IoSliceMut::new(&mut buf3),
+    /// ][..];
+    ///
+    /// // Mark 10 bytes as read.
+    /// bufs = IoSliceMut::advance(mem::replace(&mut bufs, &mut []), 10);
+    /// assert_eq!(bufs[0].deref(), [2; 14].as_ref());
+    /// assert_eq!(bufs[1].deref(), [3; 8].as_ref());
+    /// ```
+    #[unstable(feature = "io_slice_advance", issue = "62726")]
+    #[inline]
+    pub fn advance<'b>(bufs: &'b mut [IoSliceMut<'a>], n: usize) -> &'b mut [IoSliceMut<'a>] {
+        // Number of buffers to remove.
+        let mut remove = 0;
+        // Total length of all the to be removed buffers.
+        let mut accumulated_len = 0;
+        for buf in bufs.iter() {
+            if accumulated_len + buf.len() > n {
+                break;
+            } else {
+                accumulated_len += buf.len();
+                remove += 1;
+            }
+        }
+
+        let bufs = &mut bufs[remove..];
+        if !bufs.is_empty() {
+            bufs[0].0.advance(n - accumulated_len)
+        }
+        bufs
+    }
 }
 
 #[stable(feature = "iovec", since = "1.36.0")]
@@ -989,6 +1045,61 @@ impl<'a> IoSlice<'a> {
     pub fn new(buf: &'a [u8]) -> IoSlice<'a> {
         IoSlice(sys::io::IoSlice::new(buf))
     }
+
+    /// Advance the internal cursor of the slice.
+    ///
+    /// # Notes
+    ///
+    /// Elements in the slice may be modified if the cursor is not advanced to
+    /// the end of the slice. For example if we have a slice of buffers with 2
+    /// `IoSlice`s, both of length 8, and we advance the cursor by 10 bytes the
+    /// first `IoSlice` will be untouched however the second will be modified to
+    /// remove the first 2 bytes (10 - 8).
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(io_slice_advance)]
+    ///
+    /// use std::io::IoSlice;
+    /// use std::mem;
+    /// use std::ops::Deref;
+    ///
+    /// let mut buf1 = [1; 8];
+    /// let mut buf2 = [2; 16];
+    /// let mut buf3 = [3; 8];
+    /// let mut bufs = &mut [
+    ///     IoSlice::new(&mut buf1),
+    ///     IoSlice::new(&mut buf2),
+    ///     IoSlice::new(&mut buf3),
+    /// ][..];
+    ///
+    /// // Mark 10 bytes as written.
+    /// bufs = IoSlice::advance(mem::replace(&mut bufs, &mut []), 10);
+    /// assert_eq!(bufs[0].deref(), [2; 14].as_ref());
+    /// assert_eq!(bufs[1].deref(), [3; 8].as_ref());
+    #[unstable(feature = "io_slice_advance", issue = "62726")]
+    #[inline]
+    pub fn advance<'b>(bufs: &'b mut [IoSlice<'a>], n: usize) -> &'b mut [IoSlice<'a>] {
+        // Number of buffers to remove.
+        let mut remove = 0;
+        // Total length of all the to be removed buffers.
+        let mut accumulated_len = 0;
+        for buf in bufs.iter() {
+            if accumulated_len + buf.len() > n {
+                break;
+            } else {
+                accumulated_len += buf.len();
+                remove += 1;
+            }
+        }
+
+        let bufs = &mut bufs[remove..];
+        if !bufs.is_empty() {
+            bufs[0].0.advance(n - accumulated_len)
+        }
+        bufs
+    }
 }
 
 #[stable(feature = "iovec", since = "1.36.0")]
@@ -2268,8 +2379,10 @@ impl<B: BufRead> Iterator for Lines<B> {
 #[cfg(test)]
 mod tests {
     use crate::io::prelude::*;
-    use crate::io;
     use super::{Cursor, SeekFrom, repeat};
+    use crate::io::{self, IoSlice, IoSliceMut};
+    use crate::mem;
+    use crate::ops::Deref;
 
     #[test]
     #[cfg_attr(target_os = "emscripten", ignore)]
@@ -2537,4 +2650,89 @@ mod tests {
 
         Ok(())
     }
+
+    #[test]
+    fn io_slice_mut_advance() {
+        let mut buf1 = [1; 8];
+        let mut buf2 = [2; 16];
+        let mut buf3 = [3; 8];
+        let mut bufs = &mut [
+            IoSliceMut::new(&mut buf1),
+            IoSliceMut::new(&mut buf2),
+            IoSliceMut::new(&mut buf3),
+        ][..];
+
+        // Only in a single buffer..
+        bufs = IoSliceMut::advance(mem::replace(&mut bufs, &mut []), 1);
+        assert_eq!(bufs[0].deref(), [1; 7].as_ref());
+        assert_eq!(bufs[1].deref(), [2; 16].as_ref());
+        assert_eq!(bufs[2].deref(), [3; 8].as_ref());
+
+        // Removing a buffer, leaving others as is.
+        bufs = IoSliceMut::advance(mem::replace(&mut bufs, &mut []), 7);
+        assert_eq!(bufs[0].deref(), [2; 16].as_ref());
+        assert_eq!(bufs[1].deref(), [3; 8].as_ref());
+
+        // Removing a buffer and removing from the next buffer.
+        bufs = IoSliceMut::advance(mem::replace(&mut bufs, &mut []), 18);
+        assert_eq!(bufs[0].deref(), [3; 6].as_ref());
+    }
+
+    #[test]
+    fn io_slice_mut_advance_empty_slice() {
+        let mut empty_bufs = &mut [][..];
+        // Shouldn't panic.
+        IoSliceMut::advance(&mut empty_bufs, 1);
+    }
+
+    #[test]
+    fn io_slice_mut_advance_beyond_total_length() {
+        let mut buf1 = [1; 8];
+        let mut bufs = &mut [IoSliceMut::new(&mut buf1)][..];
+
+        // Going beyond the total length should be ok.
+        bufs = IoSliceMut::advance(mem::replace(&mut bufs, &mut []), 9);
+        assert!(bufs.is_empty());
+    }
+
+    #[test]
+    fn io_slice_advance() {
+        let mut buf1 = [1; 8];
+        let mut buf2 = [2; 16];
+        let mut buf3 = [3; 8];
+        let mut bufs =
+            &mut [IoSlice::new(&mut buf1), IoSlice::new(&mut buf2), IoSlice::new(&mut buf3)][..];
+
+        // Only in a single buffer..
+        bufs = IoSlice::advance(mem::replace(&mut bufs, &mut []), 1);
+        assert_eq!(bufs[0].deref(), [1; 7].as_ref());
+        assert_eq!(bufs[1].deref(), [2; 16].as_ref());
+        assert_eq!(bufs[2].deref(), [3; 8].as_ref());
+
+        // Removing a buffer, leaving others as is.
+        bufs = IoSlice::advance(mem::replace(&mut bufs, &mut []), 7);
+        assert_eq!(bufs[0].deref(), [2; 16].as_ref());
+        assert_eq!(bufs[1].deref(), [3; 8].as_ref());
+
+        // Removing a buffer and removing from the next buffer.
+        bufs = IoSlice::advance(mem::replace(&mut bufs, &mut []), 18);
+        assert_eq!(bufs[0].deref(), [3; 6].as_ref());
+    }
+
+    #[test]
+    fn io_slice_advance_empty_slice() {
+        let mut empty_bufs = &mut [][..];
+        // Shouldn't panic.
+        IoSlice::advance(&mut empty_bufs, 1);
+    }
+
+    #[test]
+    fn io_slice_advance_beyond_total_length() {
+        let mut buf1 = [1; 8];
+        let mut bufs = &mut [IoSlice::new(&mut buf1)][..];
+
+        // Going beyond the total length should be ok.
+        bufs = IoSlice::advance(mem::replace(&mut bufs, &mut []), 9);
+        assert!(bufs.is_empty());
+    }
 }
diff --git a/src/libstd/sys/cloudabi/io.rs b/src/libstd/sys/cloudabi/io.rs
index 4b423a5cbc1..976e122463d 100644
--- a/src/libstd/sys/cloudabi/io.rs
+++ b/src/libstd/sys/cloudabi/io.rs
@@ -1,3 +1,5 @@
+use crate::mem;
+
 pub struct IoSlice<'a>(&'a [u8]);
 
 impl<'a> IoSlice<'a> {
@@ -7,6 +9,11 @@ impl<'a> IoSlice<'a> {
     }
 
     #[inline]
+    pub fn advance(&mut self, n: usize) {
+        self.0 = &self.0[n..]
+    }
+
+    #[inline]
     pub fn as_slice(&self) -> &[u8] {
         self.0
     }
@@ -21,6 +28,13 @@ impl<'a> IoSliceMut<'a> {
     }
 
     #[inline]
+    pub fn advance(&mut self, n: usize) {
+        let slice = mem::replace(&mut self.0, &mut []);
+        let (_, remaining) = slice.split_at_mut(n);
+        self.0 = remaining;
+    }
+
+    #[inline]
     pub fn as_slice(&self) -> &[u8] {
         self.0
     }
diff --git a/src/libstd/sys/redox/io.rs b/src/libstd/sys/redox/io.rs
index 4b423a5cbc1..976e122463d 100644
--- a/src/libstd/sys/redox/io.rs
+++ b/src/libstd/sys/redox/io.rs
@@ -1,3 +1,5 @@
+use crate::mem;
+
 pub struct IoSlice<'a>(&'a [u8]);
 
 impl<'a> IoSlice<'a> {
@@ -7,6 +9,11 @@ impl<'a> IoSlice<'a> {
     }
 
     #[inline]
+    pub fn advance(&mut self, n: usize) {
+        self.0 = &self.0[n..]
+    }
+
+    #[inline]
     pub fn as_slice(&self) -> &[u8] {
         self.0
     }
@@ -21,6 +28,13 @@ impl<'a> IoSliceMut<'a> {
     }
 
     #[inline]
+    pub fn advance(&mut self, n: usize) {
+        let slice = mem::replace(&mut self.0, &mut []);
+        let (_, remaining) = slice.split_at_mut(n);
+        self.0 = remaining;
+    }
+
+    #[inline]
     pub fn as_slice(&self) -> &[u8] {
         self.0
     }
diff --git a/src/libstd/sys/sgx/io.rs b/src/libstd/sys/sgx/io.rs
index 4b423a5cbc1..976e122463d 100644
--- a/src/libstd/sys/sgx/io.rs
+++ b/src/libstd/sys/sgx/io.rs
@@ -1,3 +1,5 @@
+use crate::mem;
+
 pub struct IoSlice<'a>(&'a [u8]);
 
 impl<'a> IoSlice<'a> {
@@ -7,6 +9,11 @@ impl<'a> IoSlice<'a> {
     }
 
     #[inline]
+    pub fn advance(&mut self, n: usize) {
+        self.0 = &self.0[n..]
+    }
+
+    #[inline]
     pub fn as_slice(&self) -> &[u8] {
         self.0
     }
@@ -21,6 +28,13 @@ impl<'a> IoSliceMut<'a> {
     }
 
     #[inline]
+    pub fn advance(&mut self, n: usize) {
+        let slice = mem::replace(&mut self.0, &mut []);
+        let (_, remaining) = slice.split_at_mut(n);
+        self.0 = remaining;
+    }
+
+    #[inline]
     pub fn as_slice(&self) -> &[u8] {
         self.0
     }
diff --git a/src/libstd/sys/unix/io.rs b/src/libstd/sys/unix/io.rs
index bc854e772e1..a3a72919176 100644
--- a/src/libstd/sys/unix/io.rs
+++ b/src/libstd/sys/unix/io.rs
@@ -22,6 +22,18 @@ impl<'a> IoSlice<'a> {
     }
 
     #[inline]
+    pub fn advance(&mut self, n: usize) {
+        if self.vec.iov_len < n {
+            panic!("advancing IoSlice beyond its length");
+        }
+
+        unsafe {
+            self.vec.iov_len -= n;
+            self.vec.iov_base = self.vec.iov_base.add(n);
+        }
+    }
+
+    #[inline]
     pub fn as_slice(&self) -> &[u8] {
         unsafe {
             slice::from_raw_parts(self.vec.iov_base as *mut u8, self.vec.iov_len)
@@ -48,6 +60,18 @@ impl<'a> IoSliceMut<'a> {
     }
 
     #[inline]
+    pub fn advance(&mut self, n: usize) {
+        if self.vec.iov_len < n {
+            panic!("advancing IoSliceMut beyond its length");
+        }
+
+        unsafe {
+            self.vec.iov_len -= n;
+            self.vec.iov_base = self.vec.iov_base.add(n);
+        }
+    }
+
+    #[inline]
     pub fn as_slice(&self) -> &[u8] {
         unsafe {
             slice::from_raw_parts(self.vec.iov_base as *mut u8, self.vec.iov_len)
diff --git a/src/libstd/sys/vxworks/io.rs b/src/libstd/sys/vxworks/io.rs
index 72954ff20ef..8cd11cbf5df 100644
--- a/src/libstd/sys/vxworks/io.rs
+++ b/src/libstd/sys/vxworks/io.rs
@@ -22,6 +22,18 @@ impl<'a> IoSlice<'a> {
     }
 
     #[inline]
+    pub fn advance(&mut self, n: usize) {
+        if self.vec.iov_len < n {
+            panic!("advancing IoSlice beyond its length");
+        }
+
+        unsafe {
+            self.vec.iov_len -= n;
+            self.vec.iov_base = self.vec.iov_base.add(n);
+        }
+    }
+
+    #[inline]
     pub fn as_slice(&self) -> &[u8] {
         unsafe {
             slice::from_raw_parts(self.vec.iov_base as *mut u8, self.vec.iov_len)
@@ -47,6 +59,18 @@ impl<'a> IoSliceMut<'a> {
     }
 
     #[inline]
+    pub fn advance(&mut self, n: usize) {
+        if self.vec.iov_len < n {
+            panic!("advancing IoSliceMut beyond its length");
+        }
+
+        unsafe {
+            self.vec.iov_len -= n;
+            self.vec.iov_base = self.vec.iov_base.add(n);
+        }
+    }
+
+    #[inline]
     pub fn as_slice(&self) -> &[u8] {
         unsafe {
             slice::from_raw_parts(self.vec.iov_base as *mut u8, self.vec.iov_len)
diff --git a/src/libstd/sys/wasi/io.rs b/src/libstd/sys/wasi/io.rs
index a5bddad708b..ffecca5d1b6 100644
--- a/src/libstd/sys/wasi/io.rs
+++ b/src/libstd/sys/wasi/io.rs
@@ -22,6 +22,18 @@ impl<'a> IoSlice<'a> {
     }
 
     #[inline]
+    pub fn advance(&mut self, n: usize) {
+        if self.vec.buf_len < n {
+            panic!("advancing IoSlice beyond its length");
+        }
+
+        unsafe {
+            self.vec.buf_len -= n;
+            self.vec.buf = self.vec.buf.add(n);
+        }
+    }
+
+    #[inline]
     pub fn as_slice(&self) -> &[u8] {
         unsafe {
             slice::from_raw_parts(self.vec.buf as *const u8, self.vec.buf_len)
@@ -48,6 +60,18 @@ impl<'a> IoSliceMut<'a> {
     }
 
     #[inline]
+    pub fn advance(&mut self, n: usize) {
+        if self.vec.buf_len < n {
+            panic!("advancing IoSlice beyond its length");
+        }
+
+        unsafe {
+            self.vec.buf_len -= n;
+            self.vec.buf = self.vec.buf.add(n);
+        }
+    }
+
+    #[inline]
     pub fn as_slice(&self) -> &[u8] {
         unsafe {
             slice::from_raw_parts(self.vec.buf as *const u8, self.vec.buf_len)
diff --git a/src/libstd/sys/wasm/io.rs b/src/libstd/sys/wasm/io.rs
index 4b423a5cbc1..976e122463d 100644
--- a/src/libstd/sys/wasm/io.rs
+++ b/src/libstd/sys/wasm/io.rs
@@ -1,3 +1,5 @@
+use crate::mem;
+
 pub struct IoSlice<'a>(&'a [u8]);
 
 impl<'a> IoSlice<'a> {
@@ -7,6 +9,11 @@ impl<'a> IoSlice<'a> {
     }
 
     #[inline]
+    pub fn advance(&mut self, n: usize) {
+        self.0 = &self.0[n..]
+    }
+
+    #[inline]
     pub fn as_slice(&self) -> &[u8] {
         self.0
     }
@@ -21,6 +28,13 @@ impl<'a> IoSliceMut<'a> {
     }
 
     #[inline]
+    pub fn advance(&mut self, n: usize) {
+        let slice = mem::replace(&mut self.0, &mut []);
+        let (_, remaining) = slice.split_at_mut(n);
+        self.0 = remaining;
+    }
+
+    #[inline]
     pub fn as_slice(&self) -> &[u8] {
         self.0
     }
diff --git a/src/libstd/sys/windows/io.rs b/src/libstd/sys/windows/io.rs
index f0da2323f4f..e44dcbe164d 100644
--- a/src/libstd/sys/windows/io.rs
+++ b/src/libstd/sys/windows/io.rs
@@ -22,6 +22,18 @@ impl<'a> IoSlice<'a> {
     }
 
     #[inline]
+    pub fn advance(&mut self, n: usize) {
+        if (self.vec.len as usize) < n {
+            panic!("advancing IoSlice beyond its length");
+        }
+
+        unsafe {
+            self.vec.len -= n as c::ULONG;
+            self.vec.buf = self.vec.buf.add(n);
+        }
+    }
+
+    #[inline]
     pub fn as_slice(&self) -> &[u8] {
         unsafe {
             slice::from_raw_parts(self.vec.buf as *mut u8, self.vec.len as usize)
@@ -49,6 +61,18 @@ impl<'a> IoSliceMut<'a> {
     }
 
     #[inline]
+    pub fn advance(&mut self, n: usize) {
+        if (self.vec.len as usize) < n {
+            panic!("advancing IoSliceMut beyond its length");
+        }
+
+        unsafe {
+            self.vec.len -= n as c::ULONG;
+            self.vec.buf = self.vec.buf.add(n);
+        }
+    }
+
+    #[inline]
     pub fn as_slice(&self) -> &[u8] {
         unsafe {
             slice::from_raw_parts(self.vec.buf as *mut u8, self.vec.len as usize)