Implement io::Seek on io::BufReader and io::BufWriter

src/libstd/io/buffered.rs

@@ -18,7 +18,7 @@ use io::prelude::*;
 use cmp;
 use error;
 use fmt;
-use io::{self, DEFAULT_BUF_SIZE, Error, ErrorKind};
+use io::{self, DEFAULT_BUF_SIZE, Error, ErrorKind, SeekFrom};
 use ptr;
 use iter;
 
@@ -120,6 +120,52 @@ impl<R> fmt::Debug for BufReader<R> where R: fmt::Debug {
     }
 }
 
+#[unstable(feature = "buf_seek", reason = "recently added")]
+impl<R: Seek> Seek for BufReader<R> {
+    /// Seek to an offset, in bytes, in the underlying reader.
+    ///
+    /// The position used for seeking with `SeekFrom::Current(_)` is the
+    /// position the underlying reader would be at if the `BufReader` had no
+    /// internal buffer.
+    ///
+    /// Seeking always discards the internal buffer, even if the seek position
+    /// would otherwise fall within it. This guarantees that calling
+    /// `.unwrap()` immediately after a seek yields the underlying reader at
+    /// the same position.
+    ///
+    /// See `std::io::Seek` for more details.
+    ///
+    /// Note: In the edge case where you're seeking with `SeekFrom::Current(n)`
+    /// where `n` minus the internal buffer length underflows an `i64`, two
+    /// seeks will be performed instead of one. If the second seek returns
+    /// `Err`, the underlying reader will be left at the same position it would
+    /// have if you seeked to `SeekFrom::Current(0)`.
+    fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
+        let result: u64;
+        if let SeekFrom::Current(n) = pos {
+            let remainder = (self.cap - self.pos) as i64;
+            // it should be safe to assume that remainder fits within an i64 as the alternative
+            // means we managed to allocate 8 ebibytes and that's absurd.
+            // But it's not out of the realm of possibility for some weird underlying reader to
+            // support seeking by i64::min_value() so we need to handle underflow when subtracting
+            // remainder.
+            if let Some(offset) = n.checked_sub(remainder) {
+                result = try!(self.inner.seek(SeekFrom::Current(offset)));
+            } else {
+                // seek backwards by our remainder, and then by the offset
+                try!(self.inner.seek(SeekFrom::Current(-remainder)));
+                self.pos = self.cap; // empty the buffer
+                result = try!(self.inner.seek(SeekFrom::Current(n)));
+            }
+        } else {
+            // Seeking with Start/End doesn't care about our buffer length.
+            result = try!(self.inner.seek(pos));
+        }
+        self.pos = self.cap; // empty the buffer
+        Ok(result)
+    }
+}
+
 /// Wraps a Writer and buffers output to it
 ///
 /// 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 {
     }
 }
 
+#[unstable(feature = "buf_seek", reason = "recently added")]
+impl<W: Write+Seek> Seek for BufWriter<W> {
+    /// Seek to the offset, in bytes, in the underlying writer.
+    ///
+    /// Seeking always writes out the internal buffer before seeking.
+    fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
+        self.flush_buf().and_then(|_| self.get_mut().seek(pos))
+    }
+}
+
 #[unsafe_destructor]
 impl<W: Write> Drop for BufWriter<W> {
     fn drop(&mut self) {
@@ -478,7 +534,7 @@ impl<S: Write> fmt::Debug for BufStream<S> where S: fmt::Debug {
 mod tests {
     use prelude::v1::*;
     use io::prelude::*;
-    use io::{self, BufReader, BufWriter, BufStream, Cursor, LineWriter};
+    use io::{self, BufReader, BufWriter, BufStream, Cursor, LineWriter, SeekFrom};
     use test;
 
     /// A dummy reader intended at testing short-reads propagation.
@@ -533,6 +589,67 @@ mod tests {
         assert_eq!(reader.read(&mut buf).unwrap(), 0);
     }
 
+    #[test]
+    fn test_buffered_reader_seek() {
+        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.seek(SeekFrom::Start(3)).ok(), Some(3));
+        assert_eq!(reader.fill_buf().ok(), Some(&[0, 1][..]));
+        assert_eq!(reader.seek(SeekFrom::Current(0)).ok(), Some(3));
+        assert_eq!(reader.fill_buf().ok(), Some(&[0, 1][..]));
+        assert_eq!(reader.seek(SeekFrom::Current(1)).ok(), Some(4));
+        assert_eq!(reader.fill_buf().ok(), Some(&[1, 2][..]));
+        reader.consume(1);
+        assert_eq!(reader.seek(SeekFrom::Current(-2)).ok(), Some(3));
+    }
+
+    #[test]
+    fn test_buffered_reader_seek_underflow() {
+        // gimmick reader that yields its position modulo 256 for each byte
+        struct PositionReader {
+            pos: u64
+        }
+        impl Read for PositionReader {
+            fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
+                let len = buf.len();
+                for x in buf {
+                    *x = self.pos as u8;
+                    self.pos = self.pos.wrapping_add(1);
+                }
+                Ok(len)
+            }
+        }
+        impl Seek for PositionReader {
+            fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
+                match pos {
+                    SeekFrom::Start(n) => {
+                        self.pos = n;
+                    }
+                    SeekFrom::Current(n) => {
+                        self.pos = self.pos.wrapping_add(n as u64);
+                    }
+                    SeekFrom::End(n) => {
+                        self.pos = u64::max_value().wrapping_add(n as u64);
+                    }
+                }
+                Ok(self.pos)
+            }
+        }
+
+        let mut reader = BufReader::with_capacity(5, PositionReader { pos: 0 });
+        assert_eq!(reader.fill_buf().ok(), Some(&[0, 1, 2, 3, 4][..]));
+        assert_eq!(reader.seek(SeekFrom::End(-5)).ok(), Some(u64::max_value()-5));
+        assert_eq!(reader.fill_buf().ok().map(|s| s.len()), Some(5));
+        // the following seek will require two underlying seeks
+        let expected = 9223372036854775802;
+        assert_eq!(reader.seek(SeekFrom::Current(i64::min_value())).ok(), Some(expected));
+        assert_eq!(reader.fill_buf().ok().map(|s| s.len()), Some(5));
+        // seeking to 0 should empty the buffer.
+        assert_eq!(reader.seek(SeekFrom::Current(0)).ok(), Some(expected));
+        assert_eq!(reader.get_ref().pos, expected);
+    }
+
     #[test]
     fn test_buffered_writer() {
         let inner = Vec::new();
@@ -576,6 +693,18 @@ mod tests {
         assert_eq!(w, [0, 1]);
     }
 
+    #[test]
+    fn test_buffered_writer_seek() {
+        let mut w = BufWriter::with_capacity(3, io::Cursor::new(Vec::new()));
+        w.write_all(&[0, 1, 2, 3, 4, 5]).unwrap();
+        w.write_all(&[6, 7]).unwrap();
+        assert_eq!(w.seek(SeekFrom::Current(0)).ok(), Some(8));
+        assert_eq!(&w.get_ref().get_ref()[..], &[0, 1, 2, 3, 4, 5, 6, 7][..]);
+        assert_eq!(w.seek(SeekFrom::Start(2)).ok(), Some(2));
+        w.write_all(&[8, 9]).unwrap();
+        assert_eq!(&w.into_inner().unwrap().into_inner()[..], &[0, 1, 8, 9, 4, 5, 6, 7]);
+    }
+
     // This is just here to make sure that we don't infinite loop in the
     // newtype struct autoderef weirdness
     #[test]