haggis-rs/src/stream.rs

use crate::MAGIC;

use {
    crate::{Error, Node},
    std::{
        io::{ErrorKind, Read},
        iter::Iterator,
    },
};
#[cfg(feature = "parallel")]
use {
    rayon::{iter::ParallelBridge, prelude::ParallelIterator},
    std::{ops::DerefMut, sync::mpsc::Sender},
};

/// An iterator over a series of archive `Node`'s. This struct is generic over any
/// type which implements `Read`, such as a file or a network stream.
#[derive(Debug)]
pub struct Stream<R: Read + Send> {
    reader: R,
}

impl<R: Read + Send> Iterator for Stream<R> {
    type Item = Result<Node, Error>;

    fn next(&mut self) -> Option<Self::Item> {
        match Node::read(&mut self.reader) {
            Err(Error::Io(e)) if e.kind() == ErrorKind::UnexpectedEof => None,
            Ok(f) => match f.filetype {
                crate::FileType::Eof => None,
                _ => Some(Ok(f)),
            },
            x => Some(x),
        }
    }
}

#[cfg(feature = "parallel")]
#[derive(Debug)]
pub enum Message {
    FileExtracted { name: String, size: u64 },
    LinkCreated { name: String, target: String },
    DirectoryCreated { name: String },
    DeviceCreated { name: String },
    Err(Error),
    Eof,
}

impl<R: Read + Send> Stream<R> {
    pub fn new(mut reader: R) -> Result<Self, Error> {
        let mut buf = [0; 7];
        reader.read_exact(&mut buf)?;
        if buf == MAGIC {
            Ok(Self { reader })
        } else {
            Err(Error::InvalidMagic)
        }
    }

    pub fn extract(&mut self, prefix: Option<&str>) -> Result<(), Error> {
        for node in self {
            node?.extract(prefix)?;
        }
        Ok(())
    }

    #[cfg(feature = "parallel")]
    pub fn par_extract(
        &mut self,
        prefix: Option<&str>,
        sender: &Sender<Message>,
    ) -> Result<(), Error> {
        let s = sender.clone();
        self.into_iter().par_bridge().try_for_each_with(s, |s, n| {
            n?.extract(prefix)?;
            Ok::<(), Error>(())
        })?;
        Ok(())
    }
}