string-analyze 0.1.0

use std::cell::RefCell;

#[inline(always)]
fn calc_entropy<'a, I>(counts: I, total_grams: f64) -> f64
where
    I: Iterator<Item = &'a u32>,
{
    let sum_c_log2_c: f64 = counts
        .filter(|&&c| c > 0)
        .map(|&c| {
            let c_f64 = c as f64;
            c_f64 * c_f64.log2()
        })
        .sum();
    total_grams.log2() - (sum_c_log2_c / total_grams)
}

pub fn entropy(bytes:  &[u8]) -> f64 {
    if bytes.is_empty() {
        return 0.0;
    }

    let mut counts = [0u32; 256];
    for &b in bytes {
        counts[b as usize] += 1;
    }
    calc_entropy(counts.iter(), bytes.len() as f64)
}

pub fn delta_entropy(bytes: &[u8]) -> f64 {
    let len = bytes.len();

    if len < 2 {
        return 0.0;
    }

    let mut counts = [0u32; 256];
    for window in bytes.windows(2) {
        let delta = window[1].wrapping_sub(window[0]);
        counts[delta as usize] += 1;
    }

    calc_entropy(counts.iter(), (len - 1) as f64)
}
pub fn gram_entropy2(bytes: &[u8]) -> f64 {
    thread_local! {
        static GRAM_BUFFER: RefCell<(Vec<u32>, Vec<u16>)> = RefCell::new((vec![0; 65536], Vec::with_capacity(4096)));
    }
    let len = bytes.len();

    if len < 2 {
        return 0.0;
    }

    GRAM_BUFFER.with(|buf| {
        let (counts, visited) = &mut *buf.borrow_mut();
        let total_grams = (len - 1) as f64;

        for w in bytes.windows(2) {
            let idx = ((w[0] as usize) << 8) | (w[1] as usize);
            if counts[idx] == 0 {
                visited.push(idx as u16);
            }
            counts[idx] += 1;
        }

        let mut sum_c_log2_c = 0.0;
        for &idx in visited.iter() {
            let u_idx = idx as usize;
            let c = counts[u_idx];
            counts[u_idx] = 0;

            let c_f64 = c as f64;
            sum_c_log2_c += c_f64 * c_f64.log2();
        }

        visited.clear();

        total_grams.log2() - (sum_c_log2_c / total_grams)
    })
}


pub fn composite_entropy(bytes: &[u8]) -> f64 {
    let len = bytes.len();
    if len < 4 {
        return 0.0;
    }

    let e1 = entropy(bytes);
    let e2 = delta_entropy(bytes);
    let e3 = gram_entropy2(bytes);

    const EPS: f64 = 1e-6;
    let n1 = (e1 / 8.0).clamp(EPS, 1.0);
    let n2 = (e2 / 8.0).clamp(EPS, 1.0);

    let n3_raw = (e3 / 16.0).clamp(EPS, 1.0);
    // y = -4x^5 + 15x^4 - 20x^3 + 10x^2
    let n3 = {
        let x = n3_raw;
        (x * x * (10.0 + x * (-20.0 + x * (15.0 - 4.0 * x)))).clamp(EPS, 1.0)
    };

    // #[cfg(debug_assertions)]
    // dbg!(n1, n2, n3);

    const W1: f64 = 0.6;
    const W2: f64 = 0.3;
    const W3: f64 = 0.1;
    const P: f64 = -1.6;

    let mean_pow = W1 * n1.powf(P) + W2 * n2.powf(P) + W3 * n3.powf(P);
    let combined = mean_pow.powf(1.0 / P);

    (combined * 10.0).clamp(0.0, 10.0)
}


#[cfg(test)]
mod tests {
    use rand::{random, rng, RngExt};
    use super::*;

    #[test]
    fn test_entropy() {
        let input = &[
            "unittests src/lib.rs (target/debug/deps/string_analyze-551ebe6d3e6cc1ac)",
            "0123456789ABCDEF0123456789abcdef0123456789ABCDEF0123456789abcdef",
            "96ef6443c6effd748c7202c04f0577f92761b821db19fdfea2a5c2364b439209",
            "horizontal-scroll-mode",
            "112233445566778899aabbccddeeffxxyyzz112233445566778899aabbccddeeffxxyyzz",
            "aGVsbG8=",
            "MDEyMzQ1Njc4OUFNREV5TXpRMU5qYzRPVUZDUTBSRlJqQXhNak09QkNERUYwMTIz",
            "mTyqm7wjODkrNLcWl0eqO8K8gc1BPk1GNLgUpI==",
            "12~3",
            "1[82~3",
            " ",
            "ccVoejdMVQnfrzpgIunpQchW6WXcCnuX7leuRhzY3Ripc26KwOtCqxBRcjmHbN9PiaesCgHTi9iF",
            r#"M_2cL`Eqq:;|z_hiRTvhW@RLyq&d+UPXf6$ZarHI+9AiM[ZcQ"8egZ4m3Ur>J*AJE>kitQ@5Exo!C3$z.c[E#{*Rq5D"#,
            "豫章故郡，洪都新府。星分翼軫，地接衡廬。襟三江而帶五湖，控蠻荊而引甌越。物華天寶，龍光射牛斗之墟；人傑地靈，徐孺下陳蕃之榻。雄州霧列，俊彩星馳。臺隍枕夷夏之交，賓主盡東南之美。都督閻公之雅望，棨戟遙臨；宇文新州之懿範，襜帷暫駐。十旬休假，勝友如雲。千里逢迎，高朋滿座。騰蛟起鳳，孟學士之詞宗；紫電青霜，王將軍之武庫。家君作宰，路出名區。童子何知？躬逢勝餞。",
        ];

        for s in input {
            println!("~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~");
            println!("[entropy          ] [{:<5.3}]\t{s:.128}", entropy(s.as_bytes()));
            println!("[delta_entropy    ] [{:<5.3}]\t{s:.128}", delta_entropy(s.as_bytes()));
            println!("[gram_entropy2    ] [{:<5.3}]\t{s:.128}", gram_entropy2(s.as_bytes()));
            println!("[composite_entropy] [{:<5.3}]\t{s:.128}", composite_entropy(s.as_bytes()));
        }
        let input = &mut vec![
            b"\xd6\x23\x15\x9e\xbd\xe2\x90\xf9\xaa\xa0\x2e\xa0\x80\x9a\xa6\xf3\xcd\x31\xaa\x0d\x04\x6f\x51\x9c\xf1\x34\xcd\xef\x41\x29\xa4\x28\x44\x0c\xf7\x2d\xbb\x3d\x69\xf2\x03\xff\x9d\x54\x95\x25\x2d\x83\xd2\x80\x8d\x44\xef\xef\xf5\x6a\xf8\xc3\x61\x99\x9c\xe9\x2c\xec\x23\x3b\xea\xf3\x43\x1f\x57\xbd\x45\xce\x0e\x96\x4c\xb7\xbf\x28\x08\x16\x77\x53\x83\x59\x16\xd3\xac".as_slice()
        ];
        let x256 = random::<[u8;256]>();
        input.push(x256.as_slice());
        let x512 = random::<[u8;512]>();
        input.push(x512.as_slice());

        let x1024 = random::<[u8;1024]>();
        input.push(x1024.as_slice());

        let mut x4096 = vec![0u8;4096];
        rng().fill(&mut x4096);
        input.push(x4096.as_slice());

        for b in input {
            println!("~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~");
            println!("[entropy          ] [{:<5.3}]\t{:.128}", entropy(b), format!("{b:?}"));
            println!("[delta_entropy    ] [{:<5.3}]\t{:.128}", delta_entropy(b), format!("{b:?}"));
            println!("[gram_entropy2    ] [{:<5.3}]\t{:.128}", gram_entropy2(b), format!("{b:?}"));
            println!("[composite_entropy] [{:<5.3}]\t{:.128}", composite_entropy(b), format!("{b:?}"));
        }
    }
}