huffman/src/main.rs

use std::io::{self, Read, Write};
use std::path::Path;
use std::path::PathBuf;

use anyhow::{Context, anyhow};
use clap::Parser;
use huffman::{cli, hufftree, storage};

fn main() -> Result<(), anyhow::Error> {
    let args = cli::Args::parse();

    let inputf = args.input_file;
    let outputf = args.output_file;
    let mode = args.mode;

    let is_stdin = inputf == Path::new("-");

    // Read all input into memory upfront so we know its size and can inspect content.
    let input_bytes: Vec<u8> = if is_stdin {
        let mut buf = Vec::new();
        io::stdin()
            .read_to_end(&mut buf)
            .context("Could not read from stdin.")?;
        buf
    } else {
        if !inputf.exists() {
            return Err(anyhow!("Input file did not exist."));
        }
        std::fs::read(&inputf).context("Could not read input file.")?
    };
    let in_size = input_bytes.len() as u64;

    let mode = match mode {
        Some(m) => m,
        None => {
            if is_stdin {
                // No filename to inspect — infer from content validity.
                if std::str::from_utf8(&input_bytes).is_ok() {
                    cli::Mode::C
                } else {
                    cli::Mode::X
                }
            } else {
                determine_mode(&inputf, outputf.as_ref())
            }
        }
    };

    // None means write to stdout.
    let output_path: Option<PathBuf> = if is_stdin && outputf.is_none() {
        None
    } else {
        Some(match outputf {
            Some(p) => p,
            None => match mode {
                cli::Mode::X => {
                    if let Some(ext) = inputf.extension()
                        && ext.eq("z")
                    {
                        inputf.with_extension("")
                    } else {
                        inputf.with_extension("unhuffed")
                    }
                }
                cli::Mode::C => match inputf.extension() {
                    Some(ext) => {
                        let ext = ext
                            .to_str()
                            .ok_or(anyhow!("Input file path was not valid unicode."))?;
                        inputf.with_extension(ext.to_string() + ".z")
                    }
                    None => inputf.with_extension("z"),
                },
            },
        })
    };

    // When the output is stdout, status messages go to stderr to avoid corrupting binary output.
    macro_rules! status {
        ($($arg:tt)*) => {
            if output_path.is_none() {
                eprintln!($($arg)*);
            } else {
                println!($($arg)*);
            }
        };
    }

    let mut writer: Box<dyn Write> = match output_path {
        Some(ref p) => {
            Box::new(std::fs::File::create(p).context("Could not create output file.")?)
        }
        None => Box::new(io::stdout()),
    };

    status!("Read: {} bytes.", in_size);

    match mode {
        cli::Mode::X => {
            status!("Decoding text...");
            let decoded_text = huffman::storage::read_tree_and_text(&mut &input_bytes[..])?;
            status!("Decoded!");

            writer
                .write_all(decoded_text.as_bytes())
                .context("Could not write decoded text to output.")?;

            let out_size = decoded_text.len() as u64;
            status!("Stored: {} bytes.", out_size);
            let (compressed, original) = (in_size, out_size);
            status!("Compression Ratio: {:.2}.", compressed as f64 / original as f64);
        }

        cli::Mode::C => {
            let input_text =
                String::from_utf8(input_bytes).context("Input is not valid UTF-8.")?;
            status!("Encoding text...");
            let char_f = huffman::hufftree::base::get_char_frequencies(&input_text);

            let base_tree = huffman::hufftree::base::Hufftree::new(char_f);
            let canonical_tree = hufftree::canonical::CanonicalHufftree::from_tree(base_tree);

            // Buffer encoded output so we can report its size before writing.
            let mut out_buf: Vec<u8> = Vec::new();
            storage::store_tree_and_text(canonical_tree, &mut out_buf, &input_text)
                .expect("Could not store the tree and text.");
            let out_size = out_buf.len() as u64;

            writer
                .write_all(&out_buf)
                .context("Could not write encoded data to output.")?;

            status!("Encoded!");
            status!("Stored: {} bytes.", out_size);
            let (compressed, original) = (out_size, in_size);
            status!("Compression Ratio: {:.2}.", compressed as f64 / original as f64);
        }
    }

    Ok(())
}

fn determine_mode(inputf: &Path, _outputf: Option<&PathBuf>) -> cli::Mode {
    // If '.z' at end of inputf -> Decompress.
    if let Some(extension) = inputf.extension()
        && extension.eq("z")
    {
        cli::Mode::X
    } else {
        // Otherwise compress
        cli::Mode::C
    }
}