Skip to content

Commit

Permalink
Auto merge of #105917 - a1phyr:read_chain_more_impls, r=workingjubilee
Browse files Browse the repository at this point in the history
Specialize some methods of `io::Chain`

This PR specializes the implementation of some methods of `io::Chain`, which could bring performance improvements when using it.
  • Loading branch information
bors committed Feb 19, 2024
2 parents d573564 + ebc5970 commit bea5beb
Show file tree
Hide file tree
Showing 2 changed files with 66 additions and 3 deletions.
58 changes: 55 additions & 3 deletions library/std/src/io/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2672,16 +2672,50 @@ impl<T: Read, U: Read> Read for Chain<T, U> {
}
self.second.read_vectored(bufs)
}

#[inline]
fn is_read_vectored(&self) -> bool {
self.first.is_read_vectored() || self.second.is_read_vectored()
}

fn read_to_end(&mut self, buf: &mut Vec<u8>) -> Result<usize> {
let mut read = 0;
if !self.done_first {
read += self.first.read_to_end(buf)?;
self.done_first = true;
}
read += self.second.read_to_end(buf)?;
Ok(read)
}

// We don't override `read_to_string` here because an UTF-8 sequence could
// be split between the two parts of the chain

fn read_buf(&mut self, mut buf: BorrowedCursor<'_>) -> Result<()> {
if buf.capacity() == 0 {
return Ok(());
}

if !self.done_first {
let old_len = buf.written();
self.first.read_buf(buf.reborrow())?;

if buf.written() != old_len {
return Ok(());
} else {
self.done_first = true;
}
}
self.second.read_buf(buf)
}
}

#[stable(feature = "chain_bufread", since = "1.9.0")]
impl<T: BufRead, U: BufRead> BufRead for Chain<T, U> {
fn fill_buf(&mut self) -> Result<&[u8]> {
if !self.done_first {
match self.first.fill_buf()? {
buf if buf.is_empty() => {
self.done_first = true;
}
buf if buf.is_empty() => self.done_first = true,
buf => return Ok(buf),
}
}
Expand All @@ -2691,6 +2725,24 @@ impl<T: BufRead, U: BufRead> BufRead for Chain<T, U> {
fn consume(&mut self, amt: usize) {
if !self.done_first { self.first.consume(amt) } else { self.second.consume(amt) }
}

fn read_until(&mut self, byte: u8, buf: &mut Vec<u8>) -> Result<usize> {
let mut read = 0;
if !self.done_first {
let n = self.first.read_until(byte, buf)?;
read += n;

match buf.last() {
Some(b) if *b == byte && n != 0 => return Ok(read),
_ => self.done_first = true,
}
}
read += self.second.read_until(byte, buf)?;
Ok(read)
}

// We don't override `read_line` here because an UTF-8 sequence could be
// split between the two parts of the chain
}

impl<T, U> SizeHint for Chain<T, U> {
Expand Down
11 changes: 11 additions & 0 deletions library/std/src/io/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -261,6 +261,17 @@ fn chain_bufread() {
cmp_bufread(chain1, chain2, &testdata[..]);
}

#[test]
fn chain_splitted_char() {
let chain = b"\xc3".chain(b"\xa9".as_slice());
assert_eq!(crate::io::read_to_string(chain).unwrap(), "é");

let mut chain = b"\xc3".chain(b"\xa9\n".as_slice());
let mut buf = String::new();
assert_eq!(chain.read_line(&mut buf).unwrap(), 3);
assert_eq!(buf, \n");
}

#[test]
fn bufreader_size_hint() {
let testdata = b"ABCDEFGHIJKL";
Expand Down

0 comments on commit bea5beb

Please sign in to comment.