obikmer/src/obikentropy/build.rs

use std::fs;
use std::path::PathBuf;

const K_MAX: usize = 32;
const WS_MAX: usize = 6;

fn normalize_circular(kmer: u64, ws: usize) -> u64 {
    let mask = (1u64 << (ws * 2)) - 1;
    let mut canonical = kmer & mask;
    let mut current = canonical;
    for _ in 0..ws - 1 {
        let top = (current >> ((ws - 1) * 2)) & 3;
        current = ((current << 2) | top) & mask;
        if current < canonical {
            canonical = current;
        }
    }
    canonical
}

fn revcomp_raw(x: u64, k: usize) -> u64 {
    let x = !x;
    let x = x.swap_bytes();
    let x = ((x >> 4) & 0x0F0F0F0F0F0F0F0F) | ((x & 0x0F0F0F0F0F0F0F0F) << 4);
    let x = ((x >> 2) & 0x3333333333333333) | ((x & 0x3333333333333333) << 2);
    x << (64 - 2 * k)
}

fn build_normalized_kmer(k: usize) -> Vec<u64> {
    let n = 1usize << (k * 2);
    let shift = 64 - k * 2;
    let mut result = vec![0u64; n];
    for i in 0..n {
        let la = (i as u64) << shift;
        let ra = i as u64;
        let rc_ra = revcomp_raw(la, k) >> shift;
        let circ = normalize_circular(ra, k);
        let circ_rc = normalize_circular(rc_ra, k);
        result[i] = if circ < circ_rc { circ } else { circ_rc };
    }
    result
}

fn build_ln_class(norm: &[u64]) -> Vec<f64> {
    let n = norm.len();
    let mut sizes = vec![0u32; n];
    for &c in norm {
        sizes[c as usize] += 1;
    }
    norm.iter()
        .map(|&c| {
            let s = sizes[c as usize];
            if s > 0 { (s as f64).ln() } else { 0.0 }
        })
        .collect()
}

fn build_n_log_n() -> [f64; K_MAX + 1] {
    let mut t = [0.0f64; K_MAX + 1];
    for n in 1..=K_MAX {
        t[n] = (n as f64) * (n as f64).ln();
    }
    t
}

fn build_emax() -> [[f64; WS_MAX + 1]; K_MAX + 1] {
    let mut t = [[0.0f64; WS_MAX + 1]; K_MAX + 1];
    for k in 2..=K_MAX {
        for ws in 1..=WS_MAX.min(k - 1) {
            let n_raw = 1usize << (ws * 2);
            let nwords = k - ws + 1;
            let c = nwords / n_raw;
            let r = nwords % n_raw;
            let nf = nwords as f64;
            let t1 = if c == 0 || n_raw == r {
                0.0
            } else {
                let f1 = c as f64 / nf;
                (n_raw - r) as f64 * f1 * f1.ln()
            };
            let t2 = if r == 0 {
                0.0
            } else {
                let f2 = (c + 1) as f64 / nf;
                r as f64 * f2 * f2.ln()
            };
            t[k][ws] = -(t1 + t2);
        }
    }
    t
}

fn build_log_nwords() -> [[f64; WS_MAX + 1]; K_MAX + 1] {
    let mut t = [[0.0f64; WS_MAX + 1]; K_MAX + 1];
    for k in 2..=K_MAX {
        for ws in 1..=WS_MAX.min(k - 1) {
            t[k][ws] = ((k - ws + 1) as f64).ln();
        }
    }
    t
}

fn emit_f64_1d(out: &mut String, name: &str, n: usize, values: &[f64]) {
    out.push_str(&format!("pub(crate) const {name}: [f64; {n}] = [\n"));
    for v in values {
        out.push_str(&format!("    {v:?},\n"));
    }
    out.push_str("];\n");
}

fn emit_f64_2d(out: &mut String, name: &str, rows: usize, cols: usize, values: &[[f64; WS_MAX + 1]]) {
    out.push_str(&format!("pub(crate) const {name}: [[f64; {cols}]; {rows}] = [\n"));
    for row in values {
        out.push_str("    [");
        for (i, v) in row.iter().enumerate() {
            if i > 0 { out.push_str(", "); }
            out.push_str(&format!("{v:?}"));
        }
        out.push_str("],\n");
    }
    out.push_str("];\n");
}

fn main() {
    let out_dir = PathBuf::from(std::env::var("OUT_DIR").unwrap());
    let mut out = String::new();

    for k in 1..=6usize {
        let n = 1usize << (k * 2);
        let norm = build_normalized_kmer(k);
        let ln_class = build_ln_class(&norm);
        emit_f64_1d(&mut out, &format!("LN_CLASS{k}"), n, &ln_class);
    }

    let n_log_n = build_n_log_n();
    emit_f64_1d(&mut out, "N_LOG_N", K_MAX + 1, &n_log_n);

    let emax = build_emax();
    emit_f64_2d(&mut out, "EMAX", K_MAX + 1, WS_MAX + 1, &emax);

    let log_nwords = build_log_nwords();
    emit_f64_2d(&mut out, "LOG_NWORDS", K_MAX + 1, WS_MAX + 1, &log_nwords);

    fs::write(out_dir.join("ln_class_tables.rs"), out).unwrap();
}