Skip to content

Commit 27ede55

Browse files
committed
Auto merge of #46830 - Diggsey:cursor-vec-mut, r=alexcrichton
Implement `Write` for `Cursor<&mut Vec<T>>` Fixes #30132 r? @dtolnay (I'm just going through `feature-accepted` issues I swear 😛)
2 parents 92c32d2 + 77b3090 commit 27ede55

File tree

1 file changed

+58
-32
lines changed

1 file changed

+58
-32
lines changed

src/libstd/io/cursor.rs

+58-32
Original file line numberDiff line numberDiff line change
@@ -252,44 +252,62 @@ impl<T> BufRead for Cursor<T> where T: AsRef<[u8]> {
252252
fn consume(&mut self, amt: usize) { self.pos += amt as u64; }
253253
}
254254

255+
// Non-resizing write implementation
256+
fn slice_write(pos_mut: &mut u64, slice: &mut [u8], buf: &[u8]) -> io::Result<usize> {
257+
let pos = cmp::min(*pos_mut, slice.len() as u64);
258+
let amt = (&mut slice[(pos as usize)..]).write(buf)?;
259+
*pos_mut += amt as u64;
260+
Ok(amt)
261+
}
262+
263+
// Resizing write implementation
264+
fn vec_write(pos_mut: &mut u64, vec: &mut Vec<u8>, buf: &[u8]) -> io::Result<usize> {
265+
let pos: usize = (*pos_mut).try_into().map_err(|_| {
266+
Error::new(ErrorKind::InvalidInput,
267+
"cursor position exceeds maximum possible vector length")
268+
})?;
269+
// Make sure the internal buffer is as least as big as where we
270+
// currently are
271+
let len = vec.len();
272+
if len < pos {
273+
// use `resize` so that the zero filling is as efficient as possible
274+
vec.resize(pos, 0);
275+
}
276+
// Figure out what bytes will be used to overwrite what's currently
277+
// there (left), and what will be appended on the end (right)
278+
{
279+
let space = vec.len() - pos;
280+
let (left, right) = buf.split_at(cmp::min(space, buf.len()));
281+
vec[pos..pos + left.len()].copy_from_slice(left);
282+
vec.extend_from_slice(right);
283+
}
284+
285+
// Bump us forward
286+
*pos_mut = (pos + buf.len()) as u64;
287+
Ok(buf.len())
288+
}
289+
255290
#[stable(feature = "rust1", since = "1.0.0")]
256291
impl<'a> Write for Cursor<&'a mut [u8]> {
257292
#[inline]
258-
fn write(&mut self, data: &[u8]) -> io::Result<usize> {
259-
let pos = cmp::min(self.pos, self.inner.len() as u64);
260-
let amt = (&mut self.inner[(pos as usize)..]).write(data)?;
261-
self.pos += amt as u64;
262-
Ok(amt)
293+
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
294+
slice_write(&mut self.pos, self.inner, buf)
295+
}
296+
fn flush(&mut self) -> io::Result<()> { Ok(()) }
297+
}
298+
299+
#[unstable(feature = "cursor_mut_vec", issue = "30132")]
300+
impl<'a> Write for Cursor<&'a mut Vec<u8>> {
301+
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
302+
vec_write(&mut self.pos, self.inner, buf)
263303
}
264304
fn flush(&mut self) -> io::Result<()> { Ok(()) }
265305
}
266306

267307
#[stable(feature = "rust1", since = "1.0.0")]
268308
impl Write for Cursor<Vec<u8>> {
269309
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
270-
let pos: usize = self.position().try_into().map_err(|_| {
271-
Error::new(ErrorKind::InvalidInput,
272-
"cursor position exceeds maximum possible vector length")
273-
})?;
274-
// Make sure the internal buffer is as least as big as where we
275-
// currently are
276-
let len = self.inner.len();
277-
if len < pos {
278-
// use `resize` so that the zero filling is as efficient as possible
279-
self.inner.resize(pos, 0);
280-
}
281-
// Figure out what bytes will be used to overwrite what's currently
282-
// there (left), and what will be appended on the end (right)
283-
{
284-
let space = self.inner.len() - pos;
285-
let (left, right) = buf.split_at(cmp::min(space, buf.len()));
286-
self.inner[pos..pos + left.len()].copy_from_slice(left);
287-
self.inner.extend_from_slice(right);
288-
}
289-
290-
// Bump us forward
291-
self.set_position((pos + buf.len()) as u64);
292-
Ok(buf.len())
310+
vec_write(&mut self.pos, &mut self.inner, buf)
293311
}
294312
fn flush(&mut self) -> io::Result<()> { Ok(()) }
295313
}
@@ -298,10 +316,7 @@ impl Write for Cursor<Vec<u8>> {
298316
impl Write for Cursor<Box<[u8]>> {
299317
#[inline]
300318
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
301-
let pos = cmp::min(self.pos, self.inner.len() as u64);
302-
let amt = (&mut self.inner[(pos as usize)..]).write(buf)?;
303-
self.pos += amt as u64;
304-
Ok(amt)
319+
slice_write(&mut self.pos, &mut self.inner, buf)
305320
}
306321
fn flush(&mut self) -> io::Result<()> { Ok(()) }
307322
}
@@ -331,6 +346,17 @@ mod tests {
331346
assert_eq!(&writer.get_ref()[..], b);
332347
}
333348

349+
#[test]
350+
fn test_mem_mut_writer() {
351+
let mut vec = Vec::new();
352+
let mut writer = Cursor::new(&mut vec);
353+
assert_eq!(writer.write(&[0]).unwrap(), 1);
354+
assert_eq!(writer.write(&[1, 2, 3]).unwrap(), 3);
355+
assert_eq!(writer.write(&[4, 5, 6, 7]).unwrap(), 4);
356+
let b: &[_] = &[0, 1, 2, 3, 4, 5, 6, 7];
357+
assert_eq!(&writer.get_ref()[..], b);
358+
}
359+
334360
#[test]
335361
fn test_box_slice_writer() {
336362
let mut writer = Cursor::new(vec![0u8; 9].into_boxed_slice());

0 commit comments

Comments
 (0)