takes/
lib.rs

1//! A Seekable `Take` implementation.
2
3#![cfg_attr(not(debug_assertions), deny(warnings, missing_docs, clippy::dbg_macro))]
4#![cfg_attr(feature = "read_initializer", feature(read_initializer))]
5
6use std::cmp;
7use std::io::{Error, ErrorKind, Read, Result, Seek, SeekFrom};
8
9/// Extension trait for `Read + Seek` to support `takes`
10pub trait Ext: Read + Seek + Sized {
11    /// Returns a seekable Take.
12    ///
13    /// # Errors
14    /// Returns an error if the current offset could not be seeked.
15    fn takes(self, limit: u64) -> Result<Takes<Self>>;
16}
17
18impl<R: Read + Seek> Ext for R {
19    fn takes(mut self, limit: u64) -> Result<Takes<Self>> {
20        let start = self.seek(SeekFrom::Current(0))?;
21
22        Ok(Takes {
23            inner: self,
24            start,
25            limit,
26            current: 0,
27        })
28    }
29}
30
31/// A Seekable `Take` implementation.
32///
33/// Note that Seek offsets may not start at zero.
34pub struct Takes<R> {
35    inner: R,
36    start: u64,
37    limit: u64,
38    current: u64, // number of bytes of current pointer from start
39}
40
41// Code based on std::io::Take implementation
42impl<R: Read> Read for Takes<R> {
43    fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
44        let rem = self.limit - self.current;
45        // Don't call into inner reader at all at EOF because it may still block
46        if rem == 0 {
47            return Ok(0);
48        }
49
50        let max = cmp::min(buf.len() as u64, rem) as usize;
51        let n = self.inner.read(&mut buf[..max])?;
52        self.current += n as u64;
53        Ok(n)
54    }
55
56    #[cfg(feature = "read_initializer")]
57    unsafe fn initializer(&self) -> Initializer {
58        self.inner.initializer()
59    }
60}
61
62/// The absolute offsets used in the Seek implementation are *identical* to those in the underlying
63/// Read.
64/// In other words, `SeekFrom::Start(0)` may seek beyond range and cause error.
65impl<R: Seek> Seek for Takes<R> {
66    fn seek(&mut self, seek: SeekFrom) -> Result<u64> {
67        Ok(match seek {
68            SeekFrom::Start(offset) => {
69                if offset < self.start || offset > self.current {
70                    return Err(Error::new(
71                        ErrorKind::UnexpectedEof,
72                        "cannot seek beyond Takes range",
73                    ));
74                }
75                self.inner.seek(SeekFrom::Start(offset))?
76            }
77            SeekFrom::Current(delta) => {
78                let dest = (self.current as i64) + delta;
79                if dest < 0 || (dest as u64) > self.limit {
80                    return Err(Error::new(
81                        ErrorKind::UnexpectedEof,
82                        "cannot seek beyond Takes range",
83                    ));
84                }
85                self.inner.seek(SeekFrom::Current(delta))?
86            }
87            SeekFrom::End(_) => unimplemented!("SeekFrom::End implementation would be ambiguous"),
88        })
89    }
90}
takes/lib.rs

takes/
lib.rs
