mirror of
https://github.com/rust-lang/rust.git
synced 2024-11-25 00:03:43 +00:00
additional io::copy specializations
- copying from `&[u8]` and `VecDeque<u8>` - copying to `Vec<u8>`
This commit is contained in:
parent
d1389b9b48
commit
6f8ba511fa
@ -1,4 +1,8 @@
|
||||
use super::{BorrowedBuf, BufReader, BufWriter, ErrorKind, Read, Result, Write, DEFAULT_BUF_SIZE};
|
||||
use crate::alloc::Allocator;
|
||||
use crate::cmp;
|
||||
use crate::collections::VecDeque;
|
||||
use crate::io::IoSlice;
|
||||
use crate::mem::MaybeUninit;
|
||||
|
||||
#[cfg(test)]
|
||||
@ -86,7 +90,7 @@ where
|
||||
|
||||
/// Specialization of the read-write loop that reuses the internal
|
||||
/// buffer of a BufReader. If there's no buffer then the writer side
|
||||
/// should be used intead.
|
||||
/// should be used instead.
|
||||
trait BufferedReaderSpec {
|
||||
fn buffer_size(&self) -> usize;
|
||||
|
||||
@ -104,7 +108,39 @@ where
|
||||
}
|
||||
|
||||
default fn copy_to(&mut self, _to: &mut (impl Write + ?Sized)) -> Result<u64> {
|
||||
unimplemented!("only called from specializations");
|
||||
unreachable!("only called from specializations")
|
||||
}
|
||||
}
|
||||
|
||||
impl BufferedReaderSpec for &[u8] {
|
||||
fn buffer_size(&self) -> usize {
|
||||
// prefer this specialization since the source "buffer" is all we'll ever need,
|
||||
// even if it's small
|
||||
usize::MAX
|
||||
}
|
||||
|
||||
fn copy_to(&mut self, to: &mut (impl Write + ?Sized)) -> Result<u64> {
|
||||
let len = self.len();
|
||||
to.write_all(self)?;
|
||||
*self = &self[len..];
|
||||
Ok(len as u64)
|
||||
}
|
||||
}
|
||||
|
||||
impl<A: Allocator> BufferedReaderSpec for VecDeque<u8, A> {
|
||||
fn buffer_size(&self) -> usize {
|
||||
// prefer this specialization since the source "buffer" is all we'll ever need,
|
||||
// even if it's small
|
||||
usize::MAX
|
||||
}
|
||||
|
||||
fn copy_to(&mut self, to: &mut (impl Write + ?Sized)) -> Result<u64> {
|
||||
let len = self.len();
|
||||
let (front, back) = self.as_slices();
|
||||
let bufs = &mut [IoSlice::new(front), IoSlice::new(back)];
|
||||
to.write_all_vectored(bufs)?;
|
||||
self.clear();
|
||||
Ok(len as u64)
|
||||
}
|
||||
}
|
||||
|
||||
@ -218,6 +254,47 @@ impl<I: Write + ?Sized> BufferedWriterSpec for BufWriter<I> {
|
||||
}
|
||||
}
|
||||
|
||||
impl<A: Allocator> BufferedWriterSpec for Vec<u8, A> {
|
||||
fn buffer_size(&self) -> usize {
|
||||
cmp::max(DEFAULT_BUF_SIZE, self.capacity() - self.len())
|
||||
}
|
||||
|
||||
fn copy_from<R: Read + ?Sized>(&mut self, reader: &mut R) -> Result<u64> {
|
||||
let mut bytes = 0;
|
||||
|
||||
// avoid allocating before we have determined that there's anything to read
|
||||
if self.capacity() == 0 {
|
||||
bytes = stack_buffer_copy(&mut reader.take(DEFAULT_BUF_SIZE as u64), self)?;
|
||||
if bytes == 0 {
|
||||
return Ok(0);
|
||||
}
|
||||
}
|
||||
|
||||
loop {
|
||||
self.reserve(DEFAULT_BUF_SIZE);
|
||||
let mut buf: BorrowedBuf<'_> = self.spare_capacity_mut().into();
|
||||
match reader.read_buf(buf.unfilled()) {
|
||||
Ok(()) => {}
|
||||
Err(e) if e.kind() == ErrorKind::Interrupted => continue,
|
||||
Err(e) => return Err(e),
|
||||
};
|
||||
|
||||
let read = buf.filled().len();
|
||||
if read == 0 {
|
||||
break;
|
||||
}
|
||||
|
||||
// SAFETY: BorrowedBuf guarantees all of its filled bytes are init
|
||||
// and the number of read bytes can't exceed the spare capacity since
|
||||
// that's what the buffer is borrowing from.
|
||||
unsafe { self.set_len(self.len() + read) };
|
||||
bytes += read as u64;
|
||||
}
|
||||
|
||||
Ok(bytes)
|
||||
}
|
||||
}
|
||||
|
||||
fn stack_buffer_copy<R: Read + ?Sized, W: Write + ?Sized>(
|
||||
reader: &mut R,
|
||||
writer: &mut W,
|
||||
|
@ -1,4 +1,6 @@
|
||||
use crate::cmp::{max, min};
|
||||
use crate::collections::VecDeque;
|
||||
use crate::io;
|
||||
use crate::io::*;
|
||||
|
||||
#[test]
|
||||
@ -19,7 +21,7 @@ struct ShortReader {
|
||||
|
||||
impl Read for ShortReader {
|
||||
fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
|
||||
let bytes = min(self.cap, self.read_size);
|
||||
let bytes = min(self.cap, self.read_size).min(buf.len());
|
||||
self.cap -= bytes;
|
||||
self.observed_buffer = max(self.observed_buffer, buf.len());
|
||||
Ok(bytes)
|
||||
@ -78,6 +80,40 @@ fn copy_specializes_bufreader() {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn copy_specializes_to_vec() {
|
||||
let cap = 123456;
|
||||
let mut source = ShortReader { cap, observed_buffer: 0, read_size: 1337 };
|
||||
let mut sink = Vec::new();
|
||||
assert_eq!(cap as u64, io::copy(&mut source, &mut sink).unwrap());
|
||||
assert!(
|
||||
source.observed_buffer > DEFAULT_BUF_SIZE,
|
||||
"expected a large buffer to be provided to the reader"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn copy_specializes_from_vecdeque() {
|
||||
let mut source = VecDeque::with_capacity(100 * 1024);
|
||||
for _ in 0..20 * 1024 {
|
||||
source.push_front(0);
|
||||
}
|
||||
for _ in 0..20 * 1024 {
|
||||
source.push_back(0);
|
||||
}
|
||||
let mut sink = WriteObserver { observed_buffer: 0 };
|
||||
assert_eq!(40 * 1024u64, io::copy(&mut source, &mut sink).unwrap());
|
||||
assert_eq!(20 * 1024, sink.observed_buffer);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn copy_specializes_from_slice() {
|
||||
let mut source = [1; 60 * 1024].as_slice();
|
||||
let mut sink = WriteObserver { observed_buffer: 0 };
|
||||
assert_eq!(60 * 1024u64, io::copy(&mut source, &mut sink).unwrap());
|
||||
assert_eq!(60 * 1024, sink.observed_buffer);
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
mod io_benches {
|
||||
use crate::fs::File;
|
||||
|
Loading…
Reference in New Issue
Block a user