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

1 files changed, 120 insertions, 0 deletions

diff --git a/library/std/src/io/tests.rs b/library/std/src/io/tests.rs
index fd962b0415c..b22988d4a8a 100644
--- a/library/std/src/io/tests.rs
+++ b/library/std/src/io/tests.rs
@@ -416,6 +416,126 @@ 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.read(&mut buf1)?, 0);
+
+    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.read(&mut buf1)?, 0);
+
+    Ok(())
+}
+
+#[test]
+fn take_seek_error() {
+    let buf = Cursor::new(b"0123456789");
+    let mut take = buf.take(2);
+    assert!(take.seek(SeekFrom::Start(3)).is_err());
+    assert!(take.seek(SeekFrom::End(1)).is_err());
+    assert!(take.seek(SeekFrom::End(-3)).is_err());
+    assert!(take.seek(SeekFrom::Current(-1)).is_err());
+    assert!(take.seek(SeekFrom::Current(3)).is_err());
+}
+
+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 - 2))?, 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> {