Skip to content

Commit e326aa1

Browse files
committed
Auto merge of #24176 - kballard:bufreader-seek-impl, r=aturon
2 parents e57410c + 1605205 commit e326aa1

File tree

1 file changed

+131
-2
lines changed

1 file changed

+131
-2
lines changed

src/libstd/io/buffered.rs

+131-2
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ use io::prelude::*;
1818
use cmp;
1919
use error;
2020
use fmt;
21-
use io::{self, DEFAULT_BUF_SIZE, Error, ErrorKind};
21+
use io::{self, DEFAULT_BUF_SIZE, Error, ErrorKind, SeekFrom};
2222
use ptr;
2323
use iter;
2424

@@ -120,6 +120,52 @@ impl<R> fmt::Debug for BufReader<R> where R: fmt::Debug {
120120
}
121121
}
122122

123+
#[unstable(feature = "buf_seek", reason = "recently added")]
124+
impl<R: Seek> Seek for BufReader<R> {
125+
/// Seek to an offset, in bytes, in the underlying reader.
126+
///
127+
/// The position used for seeking with `SeekFrom::Current(_)` is the
128+
/// position the underlying reader would be at if the `BufReader` had no
129+
/// internal buffer.
130+
///
131+
/// Seeking always discards the internal buffer, even if the seek position
132+
/// would otherwise fall within it. This guarantees that calling
133+
/// `.unwrap()` immediately after a seek yields the underlying reader at
134+
/// the same position.
135+
///
136+
/// See `std::io::Seek` for more details.
137+
///
138+
/// Note: In the edge case where you're seeking with `SeekFrom::Current(n)`
139+
/// where `n` minus the internal buffer length underflows an `i64`, two
140+
/// seeks will be performed instead of one. If the second seek returns
141+
/// `Err`, the underlying reader will be left at the same position it would
142+
/// have if you seeked to `SeekFrom::Current(0)`.
143+
fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
144+
let result: u64;
145+
if let SeekFrom::Current(n) = pos {
146+
let remainder = (self.cap - self.pos) as i64;
147+
// it should be safe to assume that remainder fits within an i64 as the alternative
148+
// means we managed to allocate 8 ebibytes and that's absurd.
149+
// But it's not out of the realm of possibility for some weird underlying reader to
150+
// support seeking by i64::min_value() so we need to handle underflow when subtracting
151+
// remainder.
152+
if let Some(offset) = n.checked_sub(remainder) {
153+
result = try!(self.inner.seek(SeekFrom::Current(offset)));
154+
} else {
155+
// seek backwards by our remainder, and then by the offset
156+
try!(self.inner.seek(SeekFrom::Current(-remainder)));
157+
self.pos = self.cap; // empty the buffer
158+
result = try!(self.inner.seek(SeekFrom::Current(n)));
159+
}
160+
} else {
161+
// Seeking with Start/End doesn't care about our buffer length.
162+
result = try!(self.inner.seek(pos));
163+
}
164+
self.pos = self.cap; // empty the buffer
165+
Ok(result)
166+
}
167+
}
168+
123169
/// Wraps a Writer and buffers output to it
124170
///
125171
/// It can be excessively inefficient to work directly with a `Write`. For
@@ -238,6 +284,16 @@ impl<W: Write> fmt::Debug for BufWriter<W> where W: fmt::Debug {
238284
}
239285
}
240286

287+
#[unstable(feature = "buf_seek", reason = "recently added")]
288+
impl<W: Write+Seek> Seek for BufWriter<W> {
289+
/// Seek to the offset, in bytes, in the underlying writer.
290+
///
291+
/// Seeking always writes out the internal buffer before seeking.
292+
fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
293+
self.flush_buf().and_then(|_| self.get_mut().seek(pos))
294+
}
295+
}
296+
241297
#[unsafe_destructor]
242298
impl<W: Write> Drop for BufWriter<W> {
243299
fn drop(&mut self) {
@@ -478,7 +534,7 @@ impl<S: Write> fmt::Debug for BufStream<S> where S: fmt::Debug {
478534
mod tests {
479535
use prelude::v1::*;
480536
use io::prelude::*;
481-
use io::{self, BufReader, BufWriter, BufStream, Cursor, LineWriter};
537+
use io::{self, BufReader, BufWriter, BufStream, Cursor, LineWriter, SeekFrom};
482538
use test;
483539

484540
/// A dummy reader intended at testing short-reads propagation.
@@ -533,6 +589,67 @@ mod tests {
533589
assert_eq!(reader.read(&mut buf).unwrap(), 0);
534590
}
535591

592+
#[test]
593+
fn test_buffered_reader_seek() {
594+
let inner: &[u8] = &[5, 6, 7, 0, 1, 2, 3, 4];
595+
let mut reader = BufReader::with_capacity(2, io::Cursor::new(inner));
596+
597+
assert_eq!(reader.seek(SeekFrom::Start(3)).ok(), Some(3));
598+
assert_eq!(reader.fill_buf().ok(), Some(&[0, 1][..]));
599+
assert_eq!(reader.seek(SeekFrom::Current(0)).ok(), Some(3));
600+
assert_eq!(reader.fill_buf().ok(), Some(&[0, 1][..]));
601+
assert_eq!(reader.seek(SeekFrom::Current(1)).ok(), Some(4));
602+
assert_eq!(reader.fill_buf().ok(), Some(&[1, 2][..]));
603+
reader.consume(1);
604+
assert_eq!(reader.seek(SeekFrom::Current(-2)).ok(), Some(3));
605+
}
606+
607+
#[test]
608+
fn test_buffered_reader_seek_underflow() {
609+
// gimmick reader that yields its position modulo 256 for each byte
610+
struct PositionReader {
611+
pos: u64
612+
}
613+
impl Read for PositionReader {
614+
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
615+
let len = buf.len();
616+
for x in buf {
617+
*x = self.pos as u8;
618+
self.pos = self.pos.wrapping_add(1);
619+
}
620+
Ok(len)
621+
}
622+
}
623+
impl Seek for PositionReader {
624+
fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
625+
match pos {
626+
SeekFrom::Start(n) => {
627+
self.pos = n;
628+
}
629+
SeekFrom::Current(n) => {
630+
self.pos = self.pos.wrapping_add(n as u64);
631+
}
632+
SeekFrom::End(n) => {
633+
self.pos = u64::max_value().wrapping_add(n as u64);
634+
}
635+
}
636+
Ok(self.pos)
637+
}
638+
}
639+
640+
let mut reader = BufReader::with_capacity(5, PositionReader { pos: 0 });
641+
assert_eq!(reader.fill_buf().ok(), Some(&[0, 1, 2, 3, 4][..]));
642+
assert_eq!(reader.seek(SeekFrom::End(-5)).ok(), Some(u64::max_value()-5));
643+
assert_eq!(reader.fill_buf().ok().map(|s| s.len()), Some(5));
644+
// the following seek will require two underlying seeks
645+
let expected = 9223372036854775802;
646+
assert_eq!(reader.seek(SeekFrom::Current(i64::min_value())).ok(), Some(expected));
647+
assert_eq!(reader.fill_buf().ok().map(|s| s.len()), Some(5));
648+
// seeking to 0 should empty the buffer.
649+
assert_eq!(reader.seek(SeekFrom::Current(0)).ok(), Some(expected));
650+
assert_eq!(reader.get_ref().pos, expected);
651+
}
652+
536653
#[test]
537654
fn test_buffered_writer() {
538655
let inner = Vec::new();
@@ -576,6 +693,18 @@ mod tests {
576693
assert_eq!(w, [0, 1]);
577694
}
578695

696+
#[test]
697+
fn test_buffered_writer_seek() {
698+
let mut w = BufWriter::with_capacity(3, io::Cursor::new(Vec::new()));
699+
w.write_all(&[0, 1, 2, 3, 4, 5]).unwrap();
700+
w.write_all(&[6, 7]).unwrap();
701+
assert_eq!(w.seek(SeekFrom::Current(0)).ok(), Some(8));
702+
assert_eq!(&w.get_ref().get_ref()[..], &[0, 1, 2, 3, 4, 5, 6, 7][..]);
703+
assert_eq!(w.seek(SeekFrom::Start(2)).ok(), Some(2));
704+
w.write_all(&[8, 9]).unwrap();
705+
assert_eq!(&w.into_inner().unwrap().into_inner()[..], &[0, 1, 8, 9, 4, 5, 6, 7]);
706+
}
707+
579708
// This is just here to make sure that we don't infinite loop in the
580709
// newtype struct autoderef weirdness
581710
#[test]

0 commit comments

Comments
 (0)