Skip to content

Commit

Permalink
Implement Seek::stream_position() for BufReader
Browse files Browse the repository at this point in the history
Optimization over BufReader::seek() for getting the current position
without flushing the internal buffer.

Related to #31100. Based on code in #70577.
  • Loading branch information
t-rapp committed Sep 7, 2020
1 parent 445f34b commit 246d327
Show file tree
Hide file tree
Showing 2 changed files with 87 additions and 0 deletions.
45 changes: 45 additions & 0 deletions library/std/src/io/buffered.rs
Expand Up @@ -386,6 +386,51 @@ impl<R: Seek> Seek for BufReader<R> {
self.discard_buffer();
Ok(result)
}

/// Returns the current seek position from the start of the stream.
///
/// The value returned is equivalent to `self.seek(SeekFrom::Current(0))`
/// but does not flush the internal buffer. Due to this optimization the
/// function does not guarantee that calling `.into_inner()` immediately
/// afterwards will yield the underlying reader at the same position. Use
/// [`BufReader::seek`] instead if you require that guarantee.
///
/// # Panics
///
/// This function will panic if the position of the inner reader is smaller
/// than the amount of buffered data. That can happen if the inner reader
/// has an incorrect implementation of [`Seek::stream_position`], or if the
/// position has gone out of sync due to calling [`Seek::seek`] directly on
/// the underlying reader.
///
/// # Example
///
/// ```no_run
/// #![feature(seek_convenience)]
/// use std::{
/// io::{self, BufRead, BufReader, Seek},
/// fs::File,
/// };
///
/// fn main() -> io::Result<()> {
/// let mut f = BufReader::new(File::open("foo.txt")?);
///
/// let before = f.stream_position()?;
/// f.read_line(&mut String::new())?;
/// let after = f.stream_position()?;
///
/// println!("The first line was {} bytes long", after - before);
/// Ok(())
/// }
/// ```
fn stream_position(&mut self) -> io::Result<u64> {
let remainder = (self.cap - self.pos) as u64;
self.inner.stream_position().map(|pos| {
pos.checked_sub(remainder).expect(
"overflow when subtracting remaining buffer size from inner stream position",
)
})
}
}

/// Wraps a writer and buffers its output.
Expand Down
42 changes: 42 additions & 0 deletions library/std/src/io/buffered/tests.rs
@@ -1,5 +1,6 @@
use crate::io::prelude::*;
use crate::io::{self, BufReader, BufWriter, ErrorKind, IoSlice, LineWriter, SeekFrom};
use crate::panic;
use crate::sync::atomic::{AtomicUsize, Ordering};
use crate::thread;

Expand Down Expand Up @@ -86,6 +87,47 @@ fn test_buffered_reader_seek_relative() {
assert_eq!(reader.fill_buf().ok(), Some(&[2, 3][..]));
}

#[test]
fn test_buffered_reader_stream_position() {
let inner: &[u8] = &[5, 6, 7, 0, 1, 2, 3, 4];
let mut reader = BufReader::with_capacity(2, io::Cursor::new(inner));

assert_eq!(reader.stream_position().ok(), Some(0));
assert_eq!(reader.seek(SeekFrom::Start(3)).ok(), Some(3));
assert_eq!(reader.stream_position().ok(), Some(3));
// relative seeking within the buffer and reading position should keep the buffer
assert_eq!(reader.fill_buf().ok(), Some(&[0, 1][..]));
assert!(reader.seek_relative(0).is_ok());
assert_eq!(reader.stream_position().ok(), Some(3));
assert_eq!(reader.buffer(), &[0, 1][..]);
assert!(reader.seek_relative(1).is_ok());
assert_eq!(reader.stream_position().ok(), Some(4));
assert_eq!(reader.buffer(), &[1][..]);
assert!(reader.seek_relative(-1).is_ok());
assert_eq!(reader.stream_position().ok(), Some(3));
assert_eq!(reader.buffer(), &[0, 1][..]);
// relative seeking outside the buffer will discard it
assert!(reader.seek_relative(2).is_ok());
assert_eq!(reader.stream_position().ok(), Some(5));
assert_eq!(reader.buffer(), &[][..]);
}

#[test]
fn test_buffered_reader_stream_position_panic() {
let inner: &[u8] = &[5, 6, 7, 0, 1, 2, 3, 4];
let mut reader = BufReader::with_capacity(4, io::Cursor::new(inner));

// cause internal buffer to be filled but read only partially
let mut buffer = [0, 0];
assert!(reader.read_exact(&mut buffer).is_ok());
// rewinding the internal reader will cause buffer to loose sync
let inner = reader.get_mut();
assert!(inner.seek(SeekFrom::Start(0)).is_ok());
// overflow when subtracting the remaining buffer size from current position
let result = panic::catch_unwind(panic::AssertUnwindSafe(|| reader.stream_position().ok()));
assert!(result.is_err());
}

#[test]
fn test_buffered_reader_invalidated_after_read() {
let inner: &[u8] = &[5, 6, 7, 0, 1, 2, 3, 4];
Expand Down

0 comments on commit 246d327

Please sign in to comment.