Add std::io::Seek instance for std::io::Take

Author: melrief

library/std/src/io/mod.rs

@@ -1173,7 +1173,7 @@ pub trait Read {
     where
         Self: Sized,
     {
-        Take { inner: self, limit }
+        Take { inner: self, len: limit, limit }
     }
 }
 
@@ -2822,6 +2822,7 @@ impl<T, U> SizeHint for Chain<T, U> {
 #[derive(Debug)]
 pub struct Take<T> {
     inner: T,
+    len: u64,
     limit: u64,
 }
 
@@ -2856,6 +2857,12 @@ impl<T> Take<T> {
         self.limit
     }
 
+    /// Returns the number of bytes read so far.
+    #[stable(feature = "seek_io_take", since = "CURRENT_RUSTC_VERSION")]
+    pub fn position(&self) -> u64 {
+        self.len - self.limit
+    }
+
     /// Sets the number of bytes that can be read before this instance will
     /// return EOF. This is the same as constructing a new `Take` instance, so
     /// the amount of bytes read and the previous limit value don't matter when
@@ -2881,6 +2888,7 @@ impl<T> Take<T> {
     /// ```
     #[stable(feature = "take_set_limit", since = "1.27.0")]
     pub fn set_limit(&mut self, limit: u64) {
+        self.len = limit;
         self.limit = limit;
     }
 
@@ -3064,6 +3072,53 @@ impl<T> SizeHint for Take<T> {
     }
 }
 
+#[stable(feature = "seek_io_take", since = "CURRENT_RUSTC_VERSION")]
+impl<T: Seek> Seek for Take<T> {
+    fn seek(&mut self, pos: SeekFrom) -> Result<u64> {
+        let offset_from_start = match pos {
+            SeekFrom::Start(offset) => offset,
+            SeekFrom::End(offset) => {
+                if offset > 0 {
+                    return Ok(self.position());
+                }
+                if offset.unsigned_abs() > self.len {
+                    return Err(Error::from(ErrorKind::InvalidInput));
+                }
+                self.len - offset.unsigned_abs()
+            }
+            SeekFrom::Current(offset) => {
+                if offset >= 0 {
+                    self.position() + offset.unsigned_abs()
+                } else {
+                    self.position() - offset.unsigned_abs()
+                }
+            }
+        };
+        let offset_from_start = offset_from_start.min(self.len);
+        if offset_from_start > self.position() {
+            let mut offset_from_current = offset_from_start - self.position();
+            while offset_from_current > i64::MAX as u64 {
+                self.inner.seek(SeekFrom::Current(i64::MAX))?;
+                self.limit -= i64::MAX as u64;
+                offset_from_current -= i64::MAX as u64;
+            }
+            self.inner.seek(SeekFrom::Current(offset_from_current as i64))?;
+            self.limit -= offset_from_current;
+            Ok(self.position())
+        } else {
+            let mut offset_from_current = self.position() - offset_from_start;
+            while offset_from_current > i64::MIN.unsigned_abs() {
+                self.inner.seek(SeekFrom::Current(i64::MIN))?;
+                self.limit += i64::MIN.unsigned_abs();
+                offset_from_current -= i64::MIN.unsigned_abs();
+            }
+            self.inner.seek(SeekFrom::Current(-(offset_from_current as i64)))?;
+            self.limit += offset_from_current;
+            Ok(self.position())
+        }
+    }
+}
+
 /// An iterator over `u8` values of a reader.
 ///
 /// This struct is generally created by calling [`bytes`] on a reader.

library/std/src/io/tests.rs

@@ -416,6 +416,116 @@ fn seek_position() -> io::Result<()> {
     Ok(())
 }
 
+#[test]
+fn take_seek() -> io::Result<()> {
+    let mut buf = Cursor::new(b"0123456789");
+    buf.set_position(2);
+    let mut take = buf.by_ref().take(4);
+    let mut buf1 = [0u8; 1];
+    let mut buf2 = [0u8; 2];
+    assert_eq!(take.position(), 0);
+
+    assert_eq!(take.seek(SeekFrom::Start(0))?, 0);
+    take.read_exact(&mut buf2)?;
+    assert_eq!(buf2, [b'2', b'3']);
+    assert_eq!(take.seek(SeekFrom::Start(1))?, 1);
+    take.read_exact(&mut buf2)?;
+    assert_eq!(buf2, [b'3', b'4']);
+    assert_eq!(take.seek(SeekFrom::Start(2))?, 2);
+    take.read_exact(&mut buf2)?;
+    assert_eq!(buf2, [b'4', b'5']);
+    assert_eq!(take.seek(SeekFrom::Start(3))?, 3);
+    take.read_exact(&mut buf1)?;
+    assert_eq!(buf1, [b'5']);
+    assert_eq!(take.seek(SeekFrom::Start(4))?, 4);
+    assert_eq!(take.seek(SeekFrom::Start(5))?, 4);
+
+    assert_eq!(take.seek(SeekFrom::End(1))?, 4);
+    assert_eq!(take.seek(SeekFrom::End(0))?, 4);
+    assert_eq!(take.seek(SeekFrom::End(-1))?, 3);
+    take.read_exact(&mut buf1)?;
+    assert_eq!(buf1, [b'5']);
+    assert_eq!(take.seek(SeekFrom::End(-2))?, 2);
+    take.read_exact(&mut buf2)?;
+    assert_eq!(buf2, [b'4', b'5']);
+    assert_eq!(take.seek(SeekFrom::End(-3))?, 1);
+    take.read_exact(&mut buf2)?;
+    assert_eq!(buf2, [b'3', b'4']);
+    assert_eq!(take.seek(SeekFrom::End(-4))?, 0);
+    take.read_exact(&mut buf2)?;
+    assert_eq!(buf2, [b'2', b'3']);
+
+    assert_eq!(take.seek(SeekFrom::Current(0))?, 2);
+    take.read_exact(&mut buf2)?;
+    assert_eq!(buf2, [b'4', b'5']);
+
+    assert_eq!(take.seek(SeekFrom::Current(-3))?, 1);
+    take.read_exact(&mut buf2)?;
+    assert_eq!(buf2, [b'3', b'4']);
+
+    assert_eq!(take.seek(SeekFrom::Current(-1))?, 2);
+    take.read_exact(&mut buf2)?;
+    assert_eq!(buf2, [b'4', b'5']);
+
+    assert_eq!(take.seek(SeekFrom::Current(-4))?, 0);
+    take.read_exact(&mut buf2)?;
+    assert_eq!(buf2, [b'2', b'3']);
+
+    assert_eq!(take.seek(SeekFrom::Current(2))?, 4);
+    assert_eq!(take.seek(SeekFrom::Current(10))?, 4);
+
+    Ok(())
+}
+
+struct ExampleHugeRangeOfZeroes {
+    position: u64,
+}
+
+impl Read for ExampleHugeRangeOfZeroes {
+    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
+        let max = buf.len().min(usize::MAX);
+        for i in 0..max {
+            if self.position == u64::MAX {
+                return Ok(i);
+            }
+            self.position += 1;
+            buf[i] = 0;
+        }
+        Ok(max)
+    }
+}
+
+impl Seek for ExampleHugeRangeOfZeroes {
+    fn seek(&mut self, pos: io::SeekFrom) -> io::Result<u64> {
+        match pos {
+            io::SeekFrom::Start(i) => self.position = i,
+            io::SeekFrom::End(i) if i >= 0 => self.position = u64::MAX,
+            io::SeekFrom::End(i) => self.position = self.position - i.unsigned_abs(),
+            io::SeekFrom::Current(i) => {
+                self.position = if i >= 0 {
+                    self.position.saturating_add(i.unsigned_abs())
+                } else {
+                    self.position.saturating_sub(i.unsigned_abs())
+                };
+            }
+        }
+        Ok(self.position)
+    }
+}
+
+#[test]
+fn take_seek_big_offsets() -> io::Result<()> {
+    let inner = ExampleHugeRangeOfZeroes { position: 1 };
+    let mut take = inner.take(u64::MAX - 2);
+    assert_eq!(take.seek(io::SeekFrom::Start(u64::MAX))?, u64::MAX - 2);
+    assert_eq!(take.inner.position, u64::MAX - 1);
+    assert_eq!(take.seek(io::SeekFrom::Start(0))?, 0);
+    assert_eq!(take.inner.position, 1);
+    assert_eq!(take.seek(io::SeekFrom::End(-1))?, u64::MAX - 3);
+    assert_eq!(take.inner.position, u64::MAX - 2);
+    Ok(())
+}
+
 // A simple example reader which uses the default implementation of
 // read_to_end.
 struct ExampleSliceReader<'a> {