diff --git a/library/std/src/io/mod.rs b/library/std/src/io/mod.rs index 38c8120d66fc6..4d261a6389ac0 100644 --- a/library/std/src/io/mod.rs +++ b/library/std/src/io/mod.rs @@ -425,6 +425,14 @@ pub(crate) fn default_read_to_end( .and_then(|s| s.checked_add(1024)?.checked_next_multiple_of(DEFAULT_BUF_SIZE)) .unwrap_or(DEFAULT_BUF_SIZE); + // Additively counts how many bytes we wrote into buf in each + // iteration of the loop; it's used to see if we should initialize + // more bytes or re-use the initialized buffer space. + let mut written_bytes = 0; + // Tracks whether we have an initialized buffer from previous loop + // or not + let mut is_init = false; + const PROBE_SIZE: usize = 32; fn small_probe_read(r: &mut R, buf: &mut Vec) -> Result { @@ -463,21 +471,37 @@ pub(crate) fn default_read_to_end( if read == 0 { return Ok(buf.len() - start_len); + } else { + written_bytes = 0; } } if buf.len() == buf.capacity() { // buf is full, need more space buf.try_reserve(PROBE_SIZE)?; + written_bytes = 0; } let mut spare = buf.spare_capacity_mut(); - let buf_len = cmp::min(spare.len(), max_read_size); + // If our read from previous loop (or earlier) indicated that we + // read less than max_read_size bytes into buf, then we should + // aim to re-utilize the initialized portion of spare buffer. + let buf_len = if written_bytes < max_read_size { + cmp::min(spare.len(), max_read_size - written_bytes) + } else { + written_bytes = 0; + cmp::min(spare.len(), max_read_size) + }; spare = &mut spare[..buf_len]; let mut read_buf: BorrowedBuf<'_, u8> = spare.into(); - // Note that we don't track already initialized bytes here, but this is fine - // because we explicitly limit the read size + // We track if written_bytes is not equal to 0 because it could've + // been reset by buf_len if our written_bytes == max_read_size + if is_init && written_bytes != 0 { + // SAFETY: These bytes were initialized but not filled in the previous loop + unsafe { read_buf.set_init() }; + } + let mut cursor = read_buf.unfilled(); let result = loop { match r.read_buf(cursor.reborrow()) { @@ -489,7 +513,8 @@ pub(crate) fn default_read_to_end( }; let bytes_read = cursor.written(); - let is_init = read_buf.is_init(); + written_bytes += bytes_read; + is_init = read_buf.is_init(); // SAFETY: BorrowedBuf's invariants mean this much memory is initialized. unsafe {