hl 0.4.0

Highlight patterns.
Documentation
use crossterm::style::{Color, ResetColor, SetForegroundColor};
use regex::Regex;
use std::collections::HashMap;
use std::io::{BufRead, Write};

#[derive(Debug)]
pub struct Opt {
    pub color: Color,
    pub regex: Regex,
    pub index: usize,
}

#[derive(Debug)]
enum Operation {
    Start,
    End,
}

#[derive(Debug)]
struct Style {
    operation: Operation,
    color: Color,
    order: usize,
}

impl Style {
    fn start(color: Color, order: usize) -> Self {
        Self {
            operation: Operation::Start,
            color,
            order,
        }
    }

    fn end(color: Color, order: usize) -> Self {
        Self {
            operation: Operation::End,
            color,
            order,
        }
    }
}

#[macro_export]
macro_rules! opts {
    ( $( $color:path => ($name:expr, $short:expr) ),* ) => {
        {
            let mut app = clap::App::new("hl")
                .version(clap::crate_version!())
                .author(clap::crate_authors!())
                .about("Highlight patterns.");
            $(
                let mut arg = clap::Arg::with_name($name)
                    .long($name)
                    .value_name("PATTERN")
                    .help(concat!("Highlight PATTERN in ", $name))
                    .takes_value(true);
                let short: Option<&str> = $short;
                if let Some(s) = short {
                    arg = arg.short(s);
                }
                app = app.arg(arg);
            )*
            let matches = app.get_matches();

            let mut opts = Vec::new();
            $(
                matches.value_of($name).map(|pattern| {
                    let regex = regex::Regex::new(&pattern).unwrap_or_else(|_| {
                        eprintln!("Invalid regex: {:?}", pattern);
                        std::process::exit(1);
                    });
                    let index = matches.index_of($name).unwrap();
                    opts.push(hl::Opt {
                        color: $color,
                        regex,
                        index,
                    });
                });
            )*
            opts
        }
    };
}

pub fn hl<W, R>(opts: &[Opt], reader: &mut R, writer: &mut W)
where
    W: Write,
    R: BufRead,
{
    let mut reset = false;
    let mut input = String::new();
    loop {
        let len = reader.read_line(&mut input).unwrap_or_else(|e| {
            eprintln!("{}", e);
            std::process::exit(e.raw_os_error().unwrap_or(1));
        });
        if len == 0 {
            break;
        }

        let mut indices = HashMap::<usize, Vec<Style>>::new();
        for opt in opts {
            for mat in opt.regex.find_iter(&input) {
                indices
                    .entry(mat.start())
                    .or_insert_with(Vec::new)
                    .push(Style::start(opt.color, opt.index));
                indices
                    .entry(mat.end())
                    .or_insert_with(Vec::new)
                    .push(Style::end(opt.color, opt.index));
            }
        }
        for (_, v) in indices.iter_mut() {
            v.sort_by(|a, b| a.order.cmp(&b.order));
        }

        let mut stack = Stack::new();
        input.chars().enumerate().for_each(|(i, c)| {
            if let Some(styles) = indices.get(&i) {
                for style in styles {
                    match style.operation {
                        Operation::Start => {
                            reset = false;
                            stack.push(style.color);
                            let _ = write!(writer, "{}", SetForegroundColor(style.color));
                        }
                        Operation::End => {
                            stack.pop(style.color);
                            reset = true;
                            let _ = write!(writer, "{}", ResetColor);
                            for color in stack.items().iter() {
                                reset = false;
                                let _ = write!(writer, "{}", SetForegroundColor(*color));
                            }
                        }
                    }
                }
            }
            let _ = write!(writer, "{}", c);
        });
        input.clear();
    }
    if !reset {
        let _ = write!(writer, "{}", ResetColor);
    }
}

struct Stack {
    items: Vec<Color>,
}

impl Stack {
    fn new() -> Self {
        Self { items: Vec::new() }
    }

    fn push(&mut self, color: Color) {
        self.items.push(color);
    }

    fn pop(&mut self, color: Color) {
        if let Some(pos) = self.items.iter().rposition(|item| item == &color) {
            self.items.remove(pos);
        }
    }

    fn items(&self) -> &Vec<Color> {
        &self.items
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::BufReader;

    #[test]
    fn multi_line() {
        let opts = vec![
            Opt {
                color: Color::Red,
                regex: Regex::new("foo bar baz").unwrap(),
                index: 0,
            },
            Opt {
                color: Color::Blue,
                regex: Regex::new("ba").unwrap(),
                index: 1,
            },
            Opt {
                color: Color::Green,
                regex: Regex::new("bar").unwrap(),
                index: 2,
            },
        ];

        let mut output = Vec::new();
        let input = include_bytes!("../fixtures/example.txt").to_vec();
        let mut input = BufReader::new(&input[..]);
        hl(&opts, &mut input, &mut output);
        let output = String::from_utf8(output).unwrap();

        let red = format!("{}", SetForegroundColor(Color::Red));
        let green = format!("{}", SetForegroundColor(Color::Green));
        let blue = format!("{}", SetForegroundColor(Color::Blue));
        let reset = format!("{}", ResetColor);

        let expected = format!(
            "{red}foo {blue}{green}ba{reset}{red}{green}r{reset}{red} {blue}ba{reset}{red}z{reset} qux \
{blue}{green}ba{reset}{green}r{reset} {blue}{green}ba{reset}{green}r{reset}
foo {blue}{green}ba{reset}{green}r{reset}
qux {blue}ba{reset}z
{blue}{green}ba{reset}{green}r{reset}
{red}foo {blue}{green}ba{reset}{red}{green}r{reset}{red} {blue}ba{reset}{red}z{reset} {blue}{green}ba{reset}{green}r{reset}
",
            red = red,
            green = green,
            blue = blue,
            reset = reset
        );

        println!("Output:\n{}{}", &output, reset);
        println!("Expected:\n{}{}", &expected, reset);
        assert_eq!(output, expected);
    }

    #[test]
    fn no_final_linebreak() {
        let opts = vec![Opt {
            color: Color::Red,
            regex: Regex::new("foo").unwrap(),
            index: 0,
        }];

        let mut output = Vec::new();
        let input = "foo".as_bytes();
        let mut input = BufReader::new(input);
        hl(&opts, &mut input, &mut output);
        let output = String::from_utf8(output).unwrap();

        let red = format!("{}", SetForegroundColor(Color::Red));
        let reset = format!("{}", ResetColor);

        let expected = format!("{red}foo{reset}", red = red, reset = reset);

        println!("Output:\n{}{}", &output, reset);
        println!("Expected:\n{}{}", &expected, reset);
        assert_eq!(output, expected);
    }
}