r-ent 1.0.4

use std::{
    collections::BTreeMap, convert::Infallible, num::NonZeroU8, path::PathBuf, str::FromStr,
};

use argh::FromArgs;

#[derive(Debug, Clone, Default)]
enum InputPath {
    Path(PathBuf),
    #[default]
    Stdin,
}

impl std::fmt::Display for InputPath {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Path(p) => write!(f, "{p:?}"),
            Self::Stdin => f.write_str("<stdin>"),
        }
    }
}

impl FromStr for InputPath {
    type Err = Infallible;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if s == "-" {
            Ok(Self::Stdin)
        } else {
            PathBuf::from_str(s).map(InputPath::Path)
        }
    }
}

/// Measures entropy from groupings of bytes
#[derive(FromArgs)]
#[argh(help_triggers("--help", "-h"))]
struct Args {
    /// input file to read, or use `-` to read from standard input
    #[argh(positional)]
    input: Option<InputPath>,
    /// number of bytes to group together for calculating probability
    #[argh(option, short = 'w', default = "NonZeroU8::new(1).unwrap()")]
    word_size: NonZeroU8,
    /// output the "ideal" or maximal entropy
    #[argh(switch, short = 'i')]
    ideal: bool,
}

fn main() -> anyhow::Result<()> {
    let args: Args = argh::from_env();

    let ideal = ideal_entropy(args.word_size);
    let input_path = args.input.unwrap_or_default();
    let mut input: Box<dyn std::io::Read> = match input_path.clone() {
        InputPath::Path(p) => {
            let file = std::fs::File::open(p)?;
            Box::new(file)
        }
        InputPath::Stdin => Box::new(std::io::stdin()),
    };

    let entropy = measure_entropy(&mut input, args.word_size.get())?;
    let add_plural = entropy != 1.0;

    if args.ideal {
        println!(
            "{input_path}: {} bit{plu} of entropy out of {}",
            entropy,
            ideal,
            plu = if add_plural { "s" } else { "" }
        );
    } else {
        println!(
            "{input_path}: {} bit{plu} of entropy",
            entropy,
            plu = if add_plural { "s" } else { "" }
        );
    }

    Ok(())
}

fn read_fixed_array<const SIZE: usize>(
    input: &mut dyn std::io::Read,
) -> Result<[u8; SIZE], std::io::Error> {
    let mut arr = [0; SIZE];
    input.read_exact(&mut arr)?;
    Ok(arr)
}

#[derive(thiserror::Error, Debug)]
enum WordError {
    #[error("unsupported word size")]
    WordTooLarge,
    #[error(transparent)]
    Io(#[from] std::io::Error),
}

fn read_word(input: &mut dyn std::io::Read, word_size: u8) -> Result<usize, WordError> {
    match word_size {
        1 => Ok(read_fixed_array::<1>(input)?[0] as usize),
        2 => {
            let [hi, lo] = read_fixed_array::<2>(input)?;
            Ok((hi as usize) << 8 | (lo as usize))
        }
        3 => {
            let [hi, mid, lo] = read_fixed_array(input)?;
            Ok((hi as usize) << 16 | (mid as usize) << 8 | (lo as usize))
        }
        4 => {
            let [hi, hmid, lmid, lo] = read_fixed_array(input)?;
            Ok((hi as usize) << 24 | (hmid as usize) << 16 | (lmid as usize) << 8 | (lo as usize))
        }
        _ => Err(WordError::WordTooLarge),
    }
}

fn measure_entropy(input: &mut dyn std::io::Read, word_size: u8) -> Result<f32, WordError> {
    let mut counts: BTreeMap<usize, usize> = BTreeMap::new();
    loop {
        match read_word(input, word_size) {
            Ok(word) => {
                let entry = counts.entry(word).or_default();
                *entry = entry.saturating_add(1);
            }
            Err(WordError::Io(io_err)) if io_err.kind() == std::io::ErrorKind::UnexpectedEof => {
                // couldn't fill the word buffer. word size is wrong
                let mut v = Vec::new();
                input.read_to_end(&mut v)?;
                if !v.is_empty() {
                    eprintln!(
                        "there are {} bytes ignored, word size is mismatched...",
                        v.len()
                    );
                }
                break;
            }
            Err(we) => return Err(we),
        }
    }

    let total = counts.values().copied().sum::<usize>();
    let probabilities: Vec<_> = counts
        .values()
        .copied()
        .map(|cnt| {
            if total != 0 {
                (cnt as f32) / (total as f32)
            } else {
                0.0
            }
        })
        .collect();

    let neg_entropy = probabilities
        .into_iter()
        .map(|p| if p != 0.0 { p * p.log2() } else { 0.0 })
        .sum::<f32>();
    Ok(if neg_entropy == 0.0 {
        0.0
    } else {
        -neg_entropy
    })
}

/// "Ideal" entropy is the entropy where surprisal is maximal, which means that the symbols
/// of the alphabet form a uniform probability distribution.
fn ideal_entropy(word_size: NonZeroU8) -> f32 {
    let word_size: u32 = word_size.get() as u32;
    // the +1 inside wordsize count is so that we can always add 1 to the per-elem count
    // given a length, how many times should a particular element be repeated? (hi prediction, ceiling)
    let per_elem = 1.0 / (256.0f32.powi(word_size as i32));

    -(per_elem * per_elem.log2()) * 256usize.pow(word_size) as f32
}

#[cfg(test)]
mod tests {
    use std::{assert_matches, num::NonZero};

    use crate::{ideal_entropy, measure_entropy};

    #[test]
    fn test_zero_entropy() {
        let mut v = std::io::Cursor::new(vec![22; 256]);
        let entropy = measure_entropy(&mut v, 1);
        // item doesn't vary at all; no entropy
        assert_matches!(entropy, Ok(f) if f == 0.0);
    }

    #[test]
    fn test_low_entropy() {
        let mut v = std::io::Cursor::new(std::array::from_fn::<_, 200, _>(|i| {
            if i < 100 { 22 } else { 0 }
        }));
        let entropy = measure_entropy(&mut v, 1);
        // item varies half the time; 1-bit entropy
        assert_matches!(entropy, Ok(f) if f == 1.0);
    }

    #[test]
    fn test_high_entropy() {
        let mut v = std::io::Cursor::new(std::array::from_fn::<_, 256, _>(|i| i as u8));
        let entropy = measure_entropy(&mut v, 1);
        // item varies as much as possible; full entropy
        assert_matches!(entropy, Ok(f) if f == 8.0);

        let mut v = std::io::Cursor::new(
            (0..=u16::MAX)
                .into_iter()
                .flat_map(|n| n.to_be_bytes())
                .collect::<Vec<_>>(),
        );
        let entropy = measure_entropy(&mut v, 2);
        // item varies as much as possible; full entropy (16-bit)
        assert_matches!(entropy, Ok(f) if f == 16.0)
    }

    #[test]
    fn test_arbitrary_string() {
        let mut v = std::io::Cursor::new(b"0000\n");
        let entropy = measure_entropy(&mut v, 1);
        assert_matches!(entropy, Ok(f) if f == 0.72192806)
    }

    #[test]
    fn test_ideal_entropy() {
        fn non_zero(i: u8) -> NonZero<u8> {
            NonZero::new(i).unwrap()
        }
        assert_eq!(ideal_entropy(non_zero(1)), 8.0);
        assert_eq!(ideal_entropy(non_zero(2)), 16.0);
        assert_eq!(ideal_entropy(non_zero(3)), 24.0);
        assert_eq!(ideal_entropy(non_zero(4)), 32.0);
        assert_eq!(ideal_entropy(non_zero(5)), 40.0);
        assert_eq!(ideal_entropy(non_zero(6)), 48.0);
    }
}