accessor: efficient iterator method alternatives

Signed-off-by: Wolfgang Bumiller <w.bumiller@proxmox.com>

src/accessor.rs

@@ -15,7 +15,7 @@ use crate::decoder::{self, DecoderImpl};
 use crate::format::{self, GoodbyeItem};
 use crate::poll_fn::poll_fn;
 use crate::util;
-use crate::Entry;
+use crate::{Entry, EntryKind};
 
 pub mod aio;
 pub mod sync;
@@ -359,6 +359,20 @@ impl<T: Clone + ReadAt> FileEntryImpl<T> {
             .await
     }
 
+    pub async fn contents(&self) -> io::Result<FileContentsImpl<T>> {
+        let offset = self
+            .entry
+            .offset
+            .ok_or_else(|| io_format_err!("cannot open file, reader provided no offset"))?;
+        match self.entry.kind {
+            EntryKind::File { size } => Ok(FileContentsImpl::new(
+                self.input.clone(),
+                offset..(offset + size),
+            )),
+            _ => io_bail!("not a file"),
+        }
+    }
+
     #[inline]
     pub fn into_entry(self) -> Entry {
         self.entry
@@ -381,6 +395,7 @@ impl<'a, T: Clone + ReadAt> ReadDirImpl<'a, T> {
         Self { dir, at }
     }
 
+    /// Get the next entry.
     pub async fn next(&mut self) -> io::Result<Option<DirEntryImpl<'a, T>>> {
         if self.at == self.dir.table.len() {
             Ok(None)
@@ -390,6 +405,21 @@ impl<'a, T: Clone + ReadAt> ReadDirImpl<'a, T> {
             Ok(Some(cursor))
         }
     }
+
+    /// Efficient alternative to `Iterator::skip`.
+    #[inline]
+    pub fn skip(self, n: usize) -> Self {
+        Self {
+            at: (self.at + n).min(self.dir.table.len()),
+            dir: self.dir,
+        }
+    }
+
+    /// Efficient alternative to `Iterator::count`.
+    #[inline]
+    pub fn count(self) -> usize {
+        self.dir.table.len()
+    }
 }
 
 /// A cursor pointing to a file in a directory.
@@ -424,6 +454,39 @@ impl<'a, T: Clone + ReadAt> DirEntryImpl<'a, T> {
     }
 }
 
+/// A reader for file contents.
+pub struct FileContentsImpl<T> {
+    input: T,
+
+    /// Absolute offset inside the `input`.
+    range: Range<u64>,
+}
+
+impl<T: Clone + ReadAt> FileContentsImpl<T> {
+    pub fn new(input: T, range: Range<u64>) -> Self {
+        Self { input, range }
+    }
+
+    #[inline]
+    pub fn file_size(&self) -> u64 {
+        self.range.end - self.range.start
+    }
+
+    async fn read_at(&self, mut buf: &mut [u8], offset: u64) -> io::Result<usize> {
+        let size = self.file_size();
+        if offset >= size {
+            return Ok(0);
+        }
+        let remaining = size - offset;
+
+        if remaining < buf.len() as u64 {
+            buf = &mut buf[..(remaining as usize)];
+        }
+
+        (&self.input as &dyn ReadAt).read_at(buf, offset).await
+    }
+}
+
 #[doc(hidden)]
 pub struct SeqReadAtAdapter<T> {
     input: T,

src/accessor/sync.rs

@@ -189,6 +189,13 @@ impl<T: Clone + ReadAt> FileEntry<T> {
         )?))
     }
 
+    pub fn contents(&self) -> io::Result<FileContents<T>> {
+        Ok(FileContents {
+            inner: poll_result_once(self.inner.contents())?,
+            at: 0,
+        })
+    }
+
     #[inline]
     pub fn into_entry(self) -> Entry {
         self.inner.into_entry()
@@ -214,6 +221,22 @@ pub struct ReadDir<'a, T> {
     inner: accessor::ReadDirImpl<'a, T>,
 }
 
+impl<'a, T: Clone + ReadAt> ReadDir<'a, T> {
+    /// Efficient alternative to `Iterator::skip`.
+    #[inline]
+    pub fn skip(self, n: usize) -> Self {
+        Self {
+            inner: self.inner.skip(n),
+        }
+    }
+
+    /// Efficient alternative to `Iterator::count`.
+    #[inline]
+    pub fn count(self) -> usize {
+        self.inner.count()
+    }
+}
+
 impl<'a, T: Clone + ReadAt> Iterator for ReadDir<'a, T> {
     type Item = io::Result<DirEntry<'a, T>>;
 
@@ -244,3 +267,38 @@ impl<'a, T: Clone + ReadAt> DirEntry<'a, T> {
         poll_result_once(self.inner.get_entry()).map(|inner| FileEntry { inner })
     }
 }
+
+/// A reader for file contents.
+pub struct FileContents<T> {
+    inner: accessor::FileContentsImpl<T>,
+    at: u64,
+}
+
+impl<T: Clone + ReadAt> io::Read for FileContents<T> {
+    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
+        let got = poll_result_once(self.inner.read_at(buf, self.at))?;
+        self.at += got as u64;
+        Ok(got)
+    }
+}
+
+impl<T: Clone + ReadAt> FileExt for FileContents<T> {
+    fn read_at(&self, buf: &mut [u8], offset: u64) -> io::Result<usize> {
+        poll_result_once(self.inner.read_at(buf, offset))
+    }
+
+    fn write_at(&self, _buf: &[u8], _offset: u64) -> io::Result<usize> {
+        io_bail!("write_at on read-only file");
+    }
+}
+
+impl<T: Clone + ReadAt> ReadAt for FileContents<T> {
+    fn poll_read_at(
+        self: Pin<&Self>,
+        _cx: &mut Context,
+        buf: &mut [u8],
+        offset: u64,
+    ) -> Poll<io::Result<usize>> {
+        Poll::Ready(poll_result_once(self.get_ref().inner.read_at(buf, offset)))
+    }
+}