r-ent 1.0.3

use std::{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,
}

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

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

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

    if let Some(ideal) = ideal {
        println!(
            "{input_path}: {} out of {} bit{plu} of entropy",
            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: Vec<usize> = vec![0; 256usize.pow(word_size as u32)];
    loop {
        match read_word(input, word_size) {
            Ok(word) => {
                counts[word] = counts[word].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.iter().copied().sum::<usize>();
    let probabilities: Vec<_> = counts
        .iter()
        .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
    })
}

fn ideal_entropy(len: u64, word_size: u8) -> f32 {
    // 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_hi = len / (256u64.pow(word_size as u32) + 1) + 1;

    let prob_hi = per_elem_hi as f32 / len as f32;

    -(0..len)
        .into_iter()
        .map(|_| prob_hi * prob_hi.log2())
        .sum::<f32>()
}

#[cfg(test)]
mod tests {
    use std::assert_matches;

    use crate::{ideal_entropy, measure_entropy};

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

    #[test]
    fn low_entropy() {
        let mut v = std::io::Cursor::new(std::array::from_fn::<_, 200, _>(|i| {
            if i < 100 { 22 } else { 0 }
        }));
        dbg!(ideal_entropy(256, 1));
        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 high_entropy() {
        let mut v = std::io::Cursor::new(std::array::from_fn::<_, 256, _>(|i| i as u8));
        dbg!(ideal_entropy(256, 1));
        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<_>>(),
        );
        dbg!(ideal_entropy(65_536, 2));
        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 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)
    }
}