r-ent 1.0.5

use std::{
    collections::BTreeMap,
    convert::Infallible,
    num::{NonZeroU8, NonZeroUsize},
    path::{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,
    /// number of bytes to process together when reading files
    #[argh(option, short = 'b', default = "NonZeroUsize::new(256).unwrap()")]
    block_size: NonZeroUsize,
    /// 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 entropy = match input_path.clone() {
        InputPath::Path(p) => measure_file_entropy(p, args.word_size.get(), args.block_size.get()),
        InputPath::Stdin => measure_entropy(&mut std::io::stdin(), 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("block size is not a multiple of word size")]
    UnalignedBlock,
    #[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 compose_word(input: &[u8], word_size: u8) -> Result<usize, WordError> {
    match std::cmp::min(word_size, input.len() as u8) {
        1 => Ok(input[0] as usize),
        2 => Ok((input[0] as usize) << 8 | (input[1] as usize)),
        3 => Ok((input[0] as usize) << 16 | (input[1] as usize) << 8 | (input[2] as usize)),
        4 => Ok((input[0] as usize) << 24
            | (input[1] as usize) << 16
            | (input[2] as usize) << 8
            | (input[3] as usize)),
        _ => Err(WordError::WordTooLarge),
    }
}

fn measure_file_entropy<P: AsRef<Path>>(
    input_file: P,
    word_size: u8,
    block_size: usize,
) -> Result<f32, WordError> {
    if !block_size.is_multiple_of(word_size as usize) {
        return Err(WordError::UnalignedBlock);
    }
    let file_len = std::fs::metadata(input_file.as_ref())?.len() as usize;
    let _file = std::fs::File::open(input_file)?;
    let mmap = unsafe { memmap2::Mmap::map(&_file)? };
    let mut last_read = 0;
    let mut counts: BTreeMap<usize, usize> = BTreeMap::new();
    while last_read < file_len {
        let block_end = std::cmp::min(last_read + block_size, file_len);
        let block = &mmap[last_read..block_end];
        last_read = block_end;
        if !block.len().is_multiple_of(word_size as usize) {
            eprintln!(
                "there are {} byte(s) ignored, word size is mismatched to length...",
                block.len() % word_size as usize
            );
        };
        // this is why block size has to be a multiple of window size
        for window in block.chunks(word_size as usize) {
            let word = compose_word(window, word_size)?;
            let entry = counts.entry(word).or_default();
            *entry = entry.saturating_add(1);
        }
    }

    Ok(calculate_entropy(&counts))
}

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),
        }
    }

    Ok(calculate_entropy(&counts))
}

fn calculate_entropy(counts: &BTreeMap<usize, usize>) -> f32 {
    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>();
    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, collections::BTreeMap, num::NonZero};

    use crate::{calculate_entropy, 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 calculate_entropies() {
        assert_eq!(
            calculate_entropy(&BTreeMap::from_iter([(0, 10), (22, 10)])),
            1.0
        );
        assert_eq!(calculate_entropy(&BTreeMap::from_iter([(0, 10),])), 0.0);
        assert_eq!(
            calculate_entropy(&BTreeMap::from_iter([(0, 10), (22, 10), (44, 20)])),
            1.5
        );
        assert_eq!(
            calculate_entropy(&BTreeMap::from_iter([
                (0, 20),
                (22, 20),
                (44, 20),
                (88, 20),
            ])),
            2.0
        );
    }

    #[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);
    }
}