Create EntryReaderSpec trait

lib-core/src/entry/fs.rs

@@ -1,55 +1,56 @@
 use std::{fs, path::Path};
-use crate::fop::{FileOp, TypeBlacklist};
-use super::{EntryReader, EntrySaver, EntrySaverSpec, EntryType};
+use crate::fop::FileOp;
+use super::{EntryReaderSpec, EntrySaver, EntrySaverSpec, EntryType};
 
 /// An entry reader implementation for a file system. It reads a file tree from a provided directory.
 pub struct FSEntryReader {
-    src_dir: Box<Path>
+    src_dir: Box<Path>,
+    rrd: std::iter::Peekable<RecursiveReadDir>
 }
 impl FSEntryReader {
     /// Creates an entry reader with a source directory path.
-    pub const fn new(src_dir: Box<Path>) -> Self {
-        Self { src_dir }
+    pub fn new(src_dir: Box<Path>) -> Self {
+        let rrd = RecursiveReadDir::new(src_dir.clone()).peekable();
+        Self { src_dir, rrd }
     }
 }
-impl EntryReader for FSEntryReader {
-    fn read_entries(
-        self,
-        mut tx: impl FnMut(EntryType) -> crate::Result_<()>,
-        blacklist: &TypeBlacklist
-    ) -> crate::Result_<()> {
-        let mut vdir = vec![self.src_dir.clone()];
-        while let Some(px) = vdir.pop() {
-            let rd = fs::read_dir(px)?.collect::<Result<Vec<_>, _>>()?;
-            tx(EntryType::Count(rd.len()))?;
-            for de in rd {
-                let meta = de.metadata()?;
-                let et = if meta.is_dir() {
-                    let dp = de.path();
-                    let dname = if let Ok(d) = dp.strip_prefix(&self.src_dir) {
-                        d.to_string_lossy().to_string()
-                    } else {
-                        continue
-                    };
-                    vdir.push(dp.into_boxed_path());
-                    EntryType::dir(dname)
-                } else if meta.is_file() {
-                    let fp = de.path();
-                    let fname = if let Ok(d) = fp.strip_prefix(&self.src_dir) {
-                        d.to_string_lossy().to_string()
-                    } else {
-                        continue
-                    };
-                    let fop = FileOp::by_name(&fname, blacklist);
-                    let ff = fs::read(&fp)?;
-                    EntryType::file(fname, ff, fop)
-                } else {
-                    continue
-                };
-                tx(et)?;
-            }
-        }
-        Ok(())
+impl EntryReaderSpec for FSEntryReader {
+    fn len(&self) -> usize {
+        RecursiveReadDir::new(self.src_dir.clone()).filter(|x| x.is_ok()).count()
+    }
+    fn peek(&mut self) -> Option<(Option<bool>, Box<str>)> {
+        self.rrd.peek().map(|x| {
+            let Ok((is_dir, p)) = x else { return (None, "".into()) };
+            let lname = if let Ok(p) = p.strip_prefix(&self.src_dir) {
+                p.to_string_lossy().to_string()
+            } else {
+                return (None, "".into())
+            };
+            (*is_dir, lname.into_boxed_str())
+        })
+    }
+    fn skip(&mut self) {
+        self.rrd.next();
+    }
+    fn read(&mut self) -> crate::Result_<EntryType> {
+        let Some(r) = self.rrd.next() else {
+            anyhow::bail!("No more entries");
+        };
+        let (is_dir, p) = r?;
+        let lname = if let Ok(p) = p.strip_prefix(&self.src_dir) {
+            p.to_string_lossy().to_string()
+        } else {
+            anyhow::bail!("Invalid entry path: {}", p.display());
+        };
+        let et = match is_dir {
+            Some(true) => EntryType::dir(lname),
+            Some(false) => {
+                let ff = fs::read(&p)?;
+                EntryType::file(lname, ff, FileOp::Pass)
+            },
+            None => anyhow::bail!("Invalid entry type: {}", p.display()),
+        };
+        Ok(et)
     }
 }
 
@@ -77,3 +78,55 @@ impl EntrySaverSpec for FSEntrySaver {
         Ok(())
     }
 }
+
+struct RecursiveReadDir {
+    dirs: Vec<Box<Path>>,
+    cur: Option<fs::ReadDir>
+}
+impl RecursiveReadDir {
+    fn new(src_dir: Box<Path>) -> Self {
+        Self { dirs: vec![src_dir], cur: None }
+    }
+}
+impl Iterator for RecursiveReadDir {
+    type Item = std::io::Result<(Option<bool>, Box<Path>)>;
+    fn next(&mut self) -> Option<Self::Item> {
+        let rd = match self.cur {
+            None => {
+                let p = self.dirs.pop()?;
+                match fs::read_dir(p) {
+                    Ok(rd) => {
+                        self.cur = Some(rd);
+                        self.cur.as_mut().unwrap()
+                    },
+                    Err(e) => return Some(Err(e))
+                }
+            }
+            Some(ref mut rd) => rd
+        };
+        let e = match rd.next() {
+            None => {
+                self.cur = None;
+                return self.next()
+            }
+            Some(Ok(x)) => {
+                match x.file_type() {
+                    Ok(ft) => {
+                        let p = x.path().into_boxed_path();
+                        return Some(Ok((if ft.is_dir() {
+                            self.dirs.push(p.clone());
+                            Some(true)
+                        } else if ft.is_file() {
+                            Some(false)
+                        } else {
+                            None
+                        }, p)))
+                    }
+                    Err(e) => e
+                }
+            },
+            Some(Err(e)) => e
+        };
+        Some(Err(e))
+    }
+}

lib-core/src/entry/mod.rs

@@ -10,20 +10,56 @@ use crate::{cfg, errors::ErrorCollector, fop::{FileOp, TypeBlacklist}, ProgressS
 pub use fs::{FSEntryReader, FSEntrySaver};
 pub use zip::{ZipEntryReader, ZipEntrySaver};
 
-/// Trait for reading entries for further optimization. Typically used with `EntrySaver`.
-/// Any function that matches these arguments (excluding self) can be used as custom entry reader.
-pub trait EntryReader {
-    /// Reads entries and sends them via `tx`.
-    /// The `use_blacklist` parameter is used to ignore predefined file types written in `blacklist` module.
-    fn read_entries(
-        self,
-        tx: impl FnMut(EntryType) -> crate::Result_<()>,
+/// Reads entries from a file-based system.
+pub trait EntryReaderSpec {
+    /// Returns the number of entries.
+    fn len(&self) -> usize;
+    /// Peeks the next entry.
+    fn peek(&mut self) -> Option<(Option<bool>, Box<str>)>;
+    /// Skips the next entry.
+    fn skip(&mut self);
+    /// Reads the next entry.
+    fn read(&mut self) -> crate::Result_<EntryType>;
+    /// Returns `true` if there are no entries.
+    fn is_empty(&self) -> bool { self.len() == 0 }
+}
+
+/// A struct for reading entries for further optimization. Typically used with `EntrySaver`.
+pub struct EntryReader<R: EntryReaderSpec>(R);
+impl <R: EntryReaderSpec> EntryReader<R> {
+    /// Reads entries, checks if they are not blacklisted and sends them via `tx`.
+    pub fn read_entries(
+        mut self,
+        mut tx: impl FnMut(EntryType) -> crate::Result_<()>,
         blacklist: &TypeBlacklist
-    ) -> crate::Result_<()>;
+    ) -> crate::Result_<()> {
+        tx(EntryType::Count(self.0.len()))?;
+        while let Some((is_dir, name)) = self.0.peek() {
+            let fop = FileOp::by_name(&name, blacklist);
+            match is_dir {
+                Some(true) => {}
+                Some(false) => {
+                    if let FileOp::Ignore(_) = fop {
+                        self.0.skip();
+                        continue;
+                    }
+                }
+                None => {
+                    self.0.skip();
+                    continue;
+                }
+            }
+            let mut et = self.0.read()?;
+            if let EntryType::File(n, d, _) = et {
+                et = EntryType::File(n, d, fop);
+            }
+            tx(et)?;
+        }
+        Ok(())
+    }
 }
 
 /// A struct for saving entries that have been optimized. Typically used with `EntryReader`.
-/// Any function that matches these arguments (excluding self) can be used as custom entry saver.
 pub struct EntrySaver<S: EntrySaverSpec>(S);
 
 /// Saves entries in a file-based system.
@@ -34,7 +70,7 @@ pub trait EntrySaverSpec {
     fn save_file(&mut self, fname: &str, buf: &[u8], min_compress: u16) -> crate::Result_<()>;
 
 }
-impl<T: EntrySaverSpec> EntrySaver<T> {
+impl<S: EntrySaverSpec> EntrySaver<S> {
     /// Receives entries from `rx`, optimizes, sends progress (via `ps`), and saves them.
     /// Errors are collected with entry names.
     pub fn save_entries(

lib-core/src/entry/zip.rs

@@ -1,58 +1,99 @@
 use std::{io::{BufReader, BufWriter, Cursor, Read, Seek, Write}, sync::Arc};
 use zip::{write::{FileOptions, SimpleFileOptions}, CompressionMethod, ZipArchive, ZipWriter};
 
-use crate::{fop::{FileOp, TypeBlacklist}, Result_};
-use super::{EntryReader, EntrySaverSpec, EntrySaver, EntryType};
+use crate::{fop::FileOp, Result_};
+use super::{EntryReader, EntryReaderSpec, EntrySaver, EntrySaverSpec, EntryType};
 
 /// An entry reader implementation for ZIP archive. It reads its contents from a provided reader (with seeking).
 pub struct ZipEntryReader<R: Read + Seek> {
-    za: ZipArchive<R>
+    za: ZipArchive<R>,
+    cur: usize
 }
 impl <R: Read + Seek> ZipEntryReader<R> {
     /// Creates an entry reader with a specified reader.
-    pub fn new(r: R) -> Result_<Self> {
-        Ok(Self { za: ZipArchive::new(r)? })
+    pub fn new(r: R) -> Result_<EntryReader<Self>> {
+        Ok(EntryReader(Self { za: ZipArchive::new(r)?, cur: 0 }))
     }
 }
 impl <R: Read + Seek> ZipEntryReader<BufReader<R>> {
     /// Creates an entry reader wrapping a specified reader with a [`BufReader`].
-    pub fn new_buf(r: R) -> Result_<Self> {
-        Ok(Self { za: ZipArchive::new(BufReader::new(r))? })
+    #[inline]
+    pub fn new_buf(r: R) -> Result_<EntryReader<Self>> {
+        Self::new(BufReader::new(r))
     }
 }
 impl <T: AsRef<[u8]>> ZipEntryReader<Cursor<T>> {
     /// Creates an entry reader wrapping a specified reader with a [`Cursor`].
-    pub fn new_mem(t: T) -> Result_<Self> {
-        Ok(Self { za: ZipArchive::new(Cursor::new(t))? })
+    #[inline]
+    pub fn new_mem(t: T) -> Result_<EntryReader<Self>> {
+        Self::new(Cursor::new(t))
     }
 }
-impl <R: Read + Seek> EntryReader for ZipEntryReader<R> {
-    fn read_entries(
-        self,
-        mut tx: impl FnMut(EntryType) -> crate::Result_<()>,
-        blacklist: &TypeBlacklist
-    ) -> crate::Result_<()> {
-        let mut za = self.za;
+impl <R: Read + Seek> EntryReaderSpec for ZipEntryReader<R> {
+    fn len(&self) -> usize {
+        self.za.len()
+    }
+    fn peek(&mut self) -> Option<(Option<bool>, Box<str>)> {
+        let za = &self.za;
+        let jfc = za.len();
+        if self.cur >= jfc {
+            None
+        } else {
+            Some(za.name_for_index(self.cur).map_or_else(
+                || (None, "".into()),
+                |n| (Some(n.ends_with('/')), n.into())
+            ))
+        }
+    }
+    fn skip(&mut self) {
+        self.cur += 1;
+    }
+    fn read(&mut self) -> crate::Result_<EntryType> {
+        let za = &mut self.za;
         let jfc = za.len();
-        tx(EntryType::Count(jfc))?;
-        for i in 0..jfc {
-            let Some(name) = za.name_for_index(i) else { continue; };
-            let fname: Arc<str> = name.into();
-            tx(if fname.ends_with('/') {
-                EntryType::dir(fname)
+        if self.cur >= jfc {
+            anyhow::bail!("No more entries");
+        } else {
+            let i = self.cur;
+            self.cur += 1;
+            let name: Arc<str> = za.name_for_index(i).unwrap_or_default().into();
+            Ok(if name.ends_with('/') {
+                EntryType::dir(name)
             } else {
-                let fop = FileOp::by_name(&fname, blacklist);
                 let mut obuf = Vec::new();
-                if let FileOp::Ignore(_) = fop {} else {
-                    let mut jf = za.by_index(i)?;
-                    obuf.reserve_exact(jf.size() as usize);
-                    jf.read_to_end(&mut obuf)?;
-                }
-                EntryType::file(fname, obuf, fop)
-            })?;
+                let mut jf = za.by_index(i)?;
+                obuf.reserve_exact(jf.size() as usize);
+                jf.read_to_end(&mut obuf)?;
+                EntryType::file(name, obuf, FileOp::Pass)
+            })
         }
-        Ok(())
     }
+    // fn read_entries(
+    //     self,
+    //     mut tx: impl FnMut(EntryType) -> crate::Result_<()>,
+    //     blacklist: &TypeBlacklist
+    // ) -> crate::Result_<()> {
+    //     let mut za = self.za;
+    //     let jfc = za.len();
+    //     tx(EntryType::Count(jfc))?;
+    //     for i in 0..jfc {
+    //         let Some(name) = za.name_for_index(i) else { continue; };
+    //         let fname: Arc<str> = name.into();
+    //         tx(if fname.ends_with('/') {
+    //             EntryType::dir(fname)
+    //         } else {
+    //             let fop = FileOp::by_name(&fname, blacklist);
+    //             let mut obuf = Vec::new();
+    //             if let FileOp::Ignore(_) = fop {} else {
+    //                 let mut jf = za.by_index(i)?;
+    //                 obuf.reserve_exact(jf.size() as usize);
+    //                 jf.read_to_end(&mut obuf)?;
+    //             }
+    //             EntryType::file(fname, obuf, fop)
+    //         })?;
+    //     }
+    //     Ok(())
+    // }
 }
 
 #[cfg(feature = "zip-zopfli")]

src/bin/mc-repack/main.rs

@@ -4,7 +4,7 @@ use cli_args::RepackOpts;
 use crossbeam_channel::Sender;
 use indicatif::{ProgressBar, ProgressStyle, MultiProgress};
 
-use mc_repack_core::{cfg, entry::{self, EntryReader, EntrySaver, EntrySaverSpec, ZipEntryReader, ZipEntrySaver}, errors::{EntryRepackError, ErrorCollector}, fop::{FileType, TypeBlacklist}, ProgressState};
+use mc_repack_core::{cfg, entry::{self, EntryReader, EntryReaderSpec, EntrySaver, EntrySaverSpec, ZipEntryReader, ZipEntrySaver}, errors::{EntryRepackError, ErrorCollector}, fop::{FileType, TypeBlacklist}, ProgressState};
 
 mod cli_args;
 mod config;
@@ -229,8 +229,8 @@ impl std::fmt::Display for TaskError {
     }
 }
 
-pub fn optimize_with<R: EntryReader + Send + 'static, S: EntrySaverSpec>(
-    reader: R,
+pub fn optimize_with<R: EntryReaderSpec + Send + 'static, S: EntrySaverSpec>(
+    reader: EntryReader<R>,
     saver: EntrySaver<S>,
     cfgmap: &cfg::ConfigMap,
     ps: &Sender<ProgressState>,