diff --git a/src/main.rs b/src/main.rs index 2e7ca25..9d8125d 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,15 +1,45 @@ #![allow(clippy::cast_precision_loss)] use std::{ collections::HashMap, + env, io::{self, Read}, + process, }; +const USAGE: &str = "\ +usage: (ngram-entropy) | (ngram-entropy N) | (ngram-entropy MIN MAX)\ +"; + fn main() { + let arg1 = env::args().nth(1); + let arg2 = env::args().nth(2); + + let (min, max) = match (arg1, arg2) { + (Some(n), None) => { + if let Ok(n) = n.parse::() { + (n, n) + } else { + eprintln!("{USAGE}"); + process::exit(1); + } + } + (Some(min), Some(max)) => { + if let (Ok(min), Ok(max)) = (min.parse::(), max.parse::()) { + (min, max) + } else { + eprintln!("{USAGE}"); + process::exit(1); + } + } + _ => (1, 20), + }; + let mut src = Vec::new(); io::stdin().read_to_end(&mut src).unwrap(); let mut table: HashMap<&[u8], u64> = HashMap::new(); - for n in 1..=20 { + + for n in min..=max { let (ngrams, freqs) = ngram_freqs(n, &src, &mut table); let bits = entropy(freqs); println!(