wux 0.1.0

//! Rust implementation of the WUX compression format for Wii U disc images.
//!
//! Currently, only decompression is supported.

#![warn(missing_docs)]

use binrw::prelude::*;
use std::fs::File;
use std::io::{self, prelude::*, SeekFrom};
use thiserror::Error;

/// An error when processing a WUX file.
#[derive(Error, Debug)]
pub enum Error {
    /// An error when reading or writing a file.
    #[error("{0}")]
    IoError(#[from] io::Error),
    /// An error when attempting to parse an input WUX file.
    #[error("{0}")]
    ParseError(#[from] binrw::Error),
}

type Result<T> = std::result::Result<T, Error>;

#[derive(Debug)]
#[binread]
#[brw(little, magic = b"WUX0\x2e\xd0\x99\x10")]
struct WuxHeader {
    #[brw(pad_after = 4)]
    pub sector_size: u32,
    pub uncompressed_size: u64,
    #[brw(pad_after = 4)]
    pub _flags: u32,
}

#[binread]
#[brw(little, import(header: &WuxHeader))]
struct WuxLookup {
    #[br(count = (header.uncompressed_size + header.sector_size as u64 - 1) / (header.sector_size as u64))]
    #[brw(align_after = header.sector_size)]
    pub lookup_index: Vec<u32>,
}

/// The current progress of an operation.
#[derive(Debug, Copy, Clone)]
pub struct Progress {
    /// The number of bytes processed so far.
    pub bytes_processed: u64,
    /// The total bytes in the file.
    pub total_bytes: u64,
}

fn decompress_impl<R, W, F>(
    reader: &mut R,
    writer: &mut W,
    mut progress: F,
    mut copy_range: impl FnMut(&mut R, &mut W, u64, u64) -> io::Result<()>,
) -> Result<()>
where
    R: Read + Seek,
    W: Write,
    F: FnMut(Progress),
{
    let header = WuxHeader::read(reader)?;

    progress(Progress {
        bytes_processed: 0,
        total_bytes: header.uncompressed_size,
    });

    let lookup = WuxLookup::read_args(reader, (&header,))?;

    let data_start = reader.stream_position()?;
    let mut bytes_processed = 0;

    for idx in lookup.lookup_index {
        let sector_start = data_start + (idx as u64 * header.sector_size as u64);
        let sector_size = if bytes_processed + header.sector_size as u64 > header.uncompressed_size
        {
            header.uncompressed_size - bytes_processed
        } else {
            header.sector_size as u64
        };

        copy_range(reader, writer, sector_start, sector_size)?;

        bytes_processed += sector_size;

        progress(Progress {
            bytes_processed,
            total_bytes: header.uncompressed_size,
        });
    }

    Ok(())
}

/// Decompress a WUX file from `reader`, writing the resulting WUD to `writer`.
pub fn decompress<R, W>(reader: &mut R, writer: &mut W) -> Result<()>
where
    R: Read + Seek,
    W: Write,
{
    decompress_with_progress(reader, writer, drop)
}

/// Decompress a WUX file from `reader`, writing the resulting WUD to `writer` and calling
/// `progress` after each sector is written.
pub fn decompress_with_progress<R, W, F>(reader: &mut R, writer: &mut W, progress: F) -> Result<()>
where
    R: Read + Seek,
    W: Write,
    F: FnMut(Progress),
{
    decompress_impl(reader, writer, progress, |reader, writer, offset, size| {
        reader.seek(SeekFrom::Start(offset))?;
        io::copy(&mut reader.take(size), writer)?;
        Ok(())
    })
}

/// Decompress a WUX file on disk, writing the resulting WUD to another file on disk. This enables
/// certain performance improvements (currently only implemented for Linux).
pub fn decompress_file(reader: &mut File, writer: &mut File) -> Result<()> {
    decompress_file_with_progress(reader, writer, drop)
}

/// Decompress a WUX file on disk, writing the resulting WUD to another file on disk and calling
/// `progress` after each sector is written. This enables certain performance improvements
/// (currently only implemented for Linux).
pub fn decompress_file_with_progress<F>(
    reader: &mut File,
    writer: &mut File,
    progress: F,
) -> Result<()>
where
    F: FnMut(Progress),
{
    cfg_if::cfg_if! {
        if #[cfg(target_os = "linux")] {
            decompress_impl(reader, writer, progress, |reader, writer, offset, size| {
                debug_assert!(size < usize::MAX as u64);
                let mut offset = offset as i64;
                let mut size = size as usize;
                while size > 0 {
                    size -= nix::sys::sendfile::sendfile64(&writer, &reader, Some(&mut offset), size)?;
                }
                Ok(())
            })
        } else {
            decompress_with_progress(reader, writer, progress)
        }
    }
}