shitbox/corebox/commands/wc/mod.rs

use super::Cmd;
use bitflags::BitFlags;
use clap::{Arg, ArgAction, ArgMatches, Command};
use std::{
    cmp,
    error::Error,
    fmt::{self, Write},
    fs,
    io::{stdin, Read},
    ops::{AddAssign, BitAnd, BitAndAssign, BitOr, BitOrAssign},
    str::FromStr,
};

#[derive(Debug, Default)]
pub struct Wc;

impl Cmd for Wc {
    fn cli(&self) -> clap::Command {
        Command::new("wc")
            .version(env!("CARGO_PKG_VERSION"))
            .author("Nathan Fisher")
            .about("Print newline, word, and byte counts for each file")
            .args([
                Arg::new("INPUT")
                    .help("The input file to use")
                    .num_args(1..),
                Arg::new("BYTES")
                    .help("Print the byte counts")
                    .short('c')
                    .long("bytes")
                    .action(ArgAction::SetTrue),
                Arg::new("CHARS")
                    .help("Print the character counts")
                    .short('m')
                    .long("chars")
                    .action(ArgAction::SetTrue),
                Arg::new("LINES")
                    .help("Print the line counts")
                    .short('l')
                    .long("lines")
                    .action(ArgAction::SetTrue),
                Arg::new("MAX")
                    .help("Print the maximum display width")
                    .short('L')
                    .long("max-line-length")
                    .action(ArgAction::SetTrue),
                Arg::new("WORDS")
                    .help("Print the word counts")
                    .short('w')
                    .long("words")
                    .action(ArgAction::SetTrue),
            ])
    }

    fn run(&self, matches: &ArgMatches) -> Result<(), Box<dyn Error>> {
        let mut flags = 0;
        for arg in &["LINES", "WORDS", "CHARS", "BYTES", "MAX"] {
            if matches.get_flag(arg) {
                flags |= arg.parse::<Flags>()?;
            }
        }
        if flags == 0 {
            flags = 0b1011;
        }
        let files: Vec<_> = match matches.get_many::<String>("INPUT") {
            Some(c) => c.map(|x| x.to_owned()).collect(),
            None => vec!["-".to_string()],
        };
        let mut totals = Values {
            name: "Total",
            ..Values::default()
        };
        files
            .iter()
            .try_for_each(|f| get_values(f, &mut totals, flags))?;
        if files.len() > 1 {
            totals.print_values(flags)?;
        }
        Ok(())
    }

    fn path(&self) -> Option<shitbox::Path> {
        Some(shitbox::Path::UsrBin)
    }
}

fn get_values<'a>(
    file: &'a str,
    totals: &mut Values<'a>,
    flags: u32,
) -> Result<(), Box<dyn Error>> {
    let contents = if file == "-" {
        let mut buf = String::new();
        stdin().read_to_string(&mut buf)?;
        buf
    } else {
        fs::read_to_string(file)?
    };
    let mut f = Values {
        name: file,
        ..Values::default()
    };
    if flags.contains(Flags::Lines) {
        f.lines = contents.lines().count();
    }
    if flags.contains(Flags::Words) {
        f.words = contents.split_whitespace().count();
    }
    if flags.contains(Flags::Chars) {
        f.chars = contents.chars().count();
    }
    if flags.contains(Flags::Bytes) {
        f.bytes = contents.len();
    }
    if flags.contains(Flags::Max) {
        f.max = 0;
        contents.lines().for_each(|line| {
            let max = line.chars().count();
            f.max = cmp::max(max, f.max);
        });
    }
    *totals += f;
    f.print_values(flags)
}

#[derive(Clone, Copy, Default)]
struct Values<'a> {
    name: &'a str,
    lines: usize,
    words: usize,
    chars: usize,
    bytes: usize,
    max: usize,
}

impl AddAssign for Values<'_> {
    fn add_assign(&mut self, rhs: Self) {
        self.lines += rhs.lines;
        self.words += rhs.words;
        self.chars += rhs.chars;
        self.bytes += rhs.bytes;
        self.max = cmp::max(self.max, rhs.max);
    }
}

impl Values<'_> {
    fn print_values(&self, flags: u32) -> Result<(), Box<dyn Error>> {
        let mut line = String::new();
        if flags.contains(Flags::Lines) {
            write!(line, "\t{}", self.lines)?;
        }
        if flags.contains(Flags::Words) {
            write!(line, "\t{}", self.words)?;
        }
        if flags.contains(Flags::Chars) {
            write!(line, "\t{}", self.chars)?;
        }
        if flags.contains(Flags::Bytes) {
            write!(line, "\t{}", self.bytes)?;
        }
        if flags.contains(Flags::Max) {
            write!(line, "\t{}", self.max)?;
        }
        if self.name != "-" {
            write!(line, "\t{}", self.name).unwrap();
        }
        println!("{}", line);
        Ok(())
    }
}

#[derive(Debug)]
pub struct ParseFlagsError;

impl fmt::Display for ParseFlagsError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{self:?}")
    }
}

impl Error for ParseFlagsError {}

#[derive(Clone, Copy)]
enum Flags {
    Lines = 0b1,
    Words = 0b10,
    Chars = 0b100,
    Bytes = 0b1000,
    Max = 0b10000,
}

impl FromStr for Flags {
    type Err = ParseFlagsError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "LINES" => Ok(Self::Lines),
            "WORDS" => Ok(Self::Words),
            "CHARS" => Ok(Self::Chars),
            "BYTES" => Ok(Self::Bytes),
            "MAX" => Ok(Self::Max),
            _ => Err(ParseFlagsError),
        }
    }
}

impl BitAnd<u32> for Flags {
    type Output = u32;

    fn bitand(self, rhs: u32) -> Self::Output {
        self as u32 & rhs
    }
}

impl BitAnd<Flags> for u32 {
    type Output = u32;

    fn bitand(self, rhs: Flags) -> Self::Output {
        self & rhs as u32
    }
}

impl BitAndAssign<Flags> for u32 {
    fn bitand_assign(&mut self, rhs: Flags) {
        *self = *self & rhs;
    }
}

impl BitOr for Flags {
    type Output = u32;

    fn bitor(self, rhs: Self) -> Self::Output {
        self as u32 | rhs as u32
    }
}

impl BitOr<Flags> for u32 {
    type Output = u32;

    fn bitor(self, rhs: Flags) -> Self::Output {
        self | rhs as u32
    }
}

impl BitOrAssign<Flags> for u32 {
    fn bitor_assign(&mut self, rhs: Flags) {
        *self = *self | rhs;
    }
}