Implement `Write` for `Cursor<&mut Vec<T>>`

1 files changed, 58 insertions, 32 deletions

diff --git a/src/libstd/io/cursor.rs b/src/libstd/io/cursor.rs
index b5ea5531b65..c8447707d5b 100644
--- a/src/libstd/io/cursor.rs
+++ b/src/libstd/io/cursor.rs
@@ -252,14 +252,54 @@ impl<T> BufRead for Cursor<T> where T: AsRef<[u8]> {
     fn consume(&mut self, amt: usize) { self.pos += amt as u64; }
 }
 
+// Non-resizing write implementation
+fn slice_write(pos_mut: &mut u64, slice: &mut [u8], buf: &[u8]) -> io::Result<usize> {
+    let pos = cmp::min(*pos_mut, slice.len() as u64);
+    let amt = (&mut slice[(pos as usize)..]).write(buf)?;
+    *pos_mut += amt as u64;
+    Ok(amt)
+}
+
+// Resizing write implementation
+fn vec_write(pos_mut: &mut u64, vec: &mut Vec<u8>, buf: &[u8]) -> io::Result<usize> {
+    let pos: usize = (*pos_mut).try_into().map_err(|_| {
+        Error::new(ErrorKind::InvalidInput,
+                    "cursor position exceeds maximum possible vector length")
+    })?;
+    // Make sure the internal buffer is as least as big as where we
+    // currently are
+    let len = vec.len();
+    if len < pos {
+        // use `resize` so that the zero filling is as efficient as possible
+        vec.resize(pos, 0);
+    }
+    // Figure out what bytes will be used to overwrite what's currently
+    // there (left), and what will be appended on the end (right)
+    {
+        let space = vec.len() - pos;
+        let (left, right) = buf.split_at(cmp::min(space, buf.len()));
+        vec[pos..pos + left.len()].copy_from_slice(left);
+        vec.extend_from_slice(right);
+    }
+
+    // Bump us forward
+    *pos_mut = (pos + buf.len()) as u64;
+    Ok(buf.len())
+}
+
 #[stable(feature = "rust1", since = "1.0.0")]
 impl<'a> Write for Cursor<&'a mut [u8]> {
     #[inline]
-    fn write(&mut self, data: &[u8]) -> io::Result<usize> {
-        let pos = cmp::min(self.pos, self.inner.len() as u64);
-        let amt = (&mut self.inner[(pos as usize)..]).write(data)?;
-        self.pos += amt as u64;
-        Ok(amt)
+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        slice_write(&mut self.pos, self.inner, buf)
+    }
+    fn flush(&mut self) -> io::Result<()> { Ok(()) }
+}
+
+#[unstable(feature = "cursor_mut_vec", issue = "30132")]
+impl<'a> Write for Cursor<&'a mut Vec<u8>> {
+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        vec_write(&mut self.pos, self.inner, buf)
     }
     fn flush(&mut self) -> io::Result<()> { Ok(()) }
 }
@@ -267,29 +307,7 @@ impl<'a> Write for Cursor<&'a mut [u8]> {
 #[stable(feature = "rust1", since = "1.0.0")]
 impl Write for Cursor<Vec<u8>> {
     fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
-        let pos: usize = self.position().try_into().map_err(|_| {
-            Error::new(ErrorKind::InvalidInput,
-                       "cursor position exceeds maximum possible vector length")
-        })?;
-        // Make sure the internal buffer is as least as big as where we
-        // currently are
-        let len = self.inner.len();
-        if len < pos {
-            // use `resize` so that the zero filling is as efficient as possible
-            self.inner.resize(pos, 0);
-        }
-        // Figure out what bytes will be used to overwrite what's currently
-        // there (left), and what will be appended on the end (right)
-        {
-            let space = self.inner.len() - pos;
-            let (left, right) = buf.split_at(cmp::min(space, buf.len()));
-            self.inner[pos..pos + left.len()].copy_from_slice(left);
-            self.inner.extend_from_slice(right);
-        }
-
-        // Bump us forward
-        self.set_position((pos + buf.len()) as u64);
-        Ok(buf.len())
+        vec_write(&mut self.pos, &mut self.inner, buf)
     }
     fn flush(&mut self) -> io::Result<()> { Ok(()) }
 }
@@ -298,10 +316,7 @@ impl Write for Cursor<Vec<u8>> {
 impl Write for Cursor<Box<[u8]>> {
     #[inline]
     fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
-        let pos = cmp::min(self.pos, self.inner.len() as u64);
-        let amt = (&mut self.inner[(pos as usize)..]).write(buf)?;
-        self.pos += amt as u64;
-        Ok(amt)
+        slice_write(&mut self.pos, &mut self.inner, buf)
     }
     fn flush(&mut self) -> io::Result<()> { Ok(()) }
 }
@@ -332,6 +347,17 @@ mod tests {
     }
 
     #[test]
+    fn test_mem_mut_writer() {
+        let mut vec = Vec::new();
+        let mut writer = Cursor::new(&mut vec);
+        assert_eq!(writer.write(&[0]).unwrap(), 1);
+        assert_eq!(writer.write(&[1, 2, 3]).unwrap(), 3);
+        assert_eq!(writer.write(&[4, 5, 6, 7]).unwrap(), 4);
+        let b: &[_] = &[0, 1, 2, 3, 4, 5, 6, 7];
+        assert_eq!(&writer.get_ref()[..], b);
+    }
+
+    #[test]
     fn test_box_slice_writer() {
         let mut writer = Cursor::new(vec![0u8; 9].into_boxed_slice());
         assert_eq!(writer.position(), 0);