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 BufRead for Cursor 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 { + 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, buf: &[u8]) -> io::Result { + 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 { - 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 { + 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> { + fn write(&mut self, buf: &[u8]) -> io::Result { + 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> { fn write(&mut self, buf: &[u8]) -> io::Result { - 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> { impl Write for Cursor> { #[inline] fn write(&mut self, buf: &[u8]) -> io::Result { - 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(()) } } @@ -331,6 +346,17 @@ mod tests { assert_eq!(&writer.get_ref()[..], b); } + #[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());