nordisk/nornir

src/clone_signals/minhash.rs

//! **MinHash + LSH** over the normalised token classes — the near-exact half
//! of signal 1 (`near_clones`).
//!
//! A function's set is its [`SHINGLE`]-token shingles of role CLASSES
//! ([`super::norm::NormBody::classes`]: `V`, `F`, `T`, keywords, kept std
//! names, punctuation, literal kinds — never a name the workspace chose).
//! [`PERMS`] independent hash functions keep their minimum over the set; the
//! fraction of equal minima estimates the Jaccard of two sets. LSH cuts the
//! signature into [`BANDS`] bands of [`ROWS`] rows: two functions are a
//! candidate when one band is equal — `1 − (1 − J^ROWS)^BANDS`, so J = 0.8 is
//! a candidate with probability 0.9996, J = 0.5 with 0.64, J = 0.3 with 0.12;
//! every candidate is then judged on its estimate, never on the band alone.

use std::collections::HashMap;

/// Hash functions in a signature.
pub const PERMS: usize = 64;
/// LSH bands × rows per band = [`PERMS`].
pub const BANDS: usize = 16;
pub const ROWS: usize = PERMS / BANDS;
/// Tokens per shingle.
pub const SHINGLE: usize = 5;
/// The estimated Jaccard at or above which two functions are near clones.
pub const NEAR_FLOOR: f32 = 0.80;

/// splitmix64: one well-mixed hash per (shingle, permutation).
fn mix(mut x: u64) -> u64 {
    x = x.wrapping_add(0x9e37_79b9_7f4a_7c15);
    x = (x ^ (x >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
    x = (x ^ (x >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb);
    x ^ (x >> 31)
}

/// The signature of a token-class sequence; `None` below one shingle.
pub fn signature(classes: &[u64]) -> Option<Vec<u32>> {
    if classes.len() < SHINGLE {
        return None;
    }
    let mut sig = vec![u32::MAX; PERMS];
    for w in classes.windows(SHINGLE) {
        let s = w.iter().fold(0x51_7cc1_b727_220au64, |h, c| mix(h ^ *c));
        for (p, slot) in sig.iter_mut().enumerate() {
            let v = (mix(s ^ (p as u64).wrapping_mul(0xa076_1d64_78bd_642f)) >> 32) as u32;
            if v < *slot {
                *slot = v;
            }
        }
    }
    Some(sig)
}

/// The estimated Jaccard of two signatures.
pub fn estimate(a: &[u32], b: &[u32]) -> f32 {
    let same = a.iter().zip(b).filter(|(x, y)| x == y).count();
    same as f32 / a.len().max(1) as f32
}

/// Signature ↔ bytes (the stored column): little-endian u32s.
pub fn to_bytes(sig: &[u32]) -> Vec<u8> {
    sig.iter().flat_map(|v| v.to_le_bytes()).collect()
}

pub fn from_bytes(b: &[u8]) -> Option<Vec<u32>> {
    (b.len() == PERMS * 4).then(|| b.chunks_exact(4).map(|c| u32::from_le_bytes([c[0], c[1], c[2], c[3]])).collect())
}

/// Every pair `(i, j, estimate)`, `i < j`, of `sigs` that shares an LSH band
/// and estimates at or above `floor`. A bucket larger than `max_bucket` is a
/// boilerplate shape every small function has: it is skipped and counted
/// (the second value), never expanded into its n² pairs.
pub fn near_pairs(sigs: &[&[u32]], floor: f32, max_bucket: usize) -> (Vec<(usize, usize, f32)>, usize) {
    let mut seen: std::collections::HashSet<(u32, u32)> = std::collections::HashSet::new();
    let mut out = Vec::new();
    let mut skipped = 0usize;
    for band in 0..BANDS {
        let mut buckets: HashMap<&[u32], Vec<u32>> = HashMap::new();
        for (i, s) in sigs.iter().enumerate() {
            buckets.entry(&s[band * ROWS..(band + 1) * ROWS]).or_default().push(i as u32);
        }
        for (_, members) in buckets {
            if members.len() < 2 {
                continue;
            }
            if members.len() > max_bucket {
                skipped += 1;
                continue;
            }
            for (x, &i) in members.iter().enumerate() {
                for &j in &members[x + 1..] {
                    if seen.insert((i, j)) {
                        let e = estimate(sigs[i as usize], sigs[j as usize]);
                        if e >= floor {
                            out.push((i as usize, j as usize, e));
                        }
                    }
                }
            }
        }
    }
    out.sort_by(|a, b| b.2.total_cmp(&a.2).then((a.0, a.1).cmp(&(b.0, b.1))));
    (out, skipped)
}

#[cfg(test)]
mod tests {
    use super::*;

    fn classes(s: &str) -> Vec<u64> {
        s.split_whitespace().map(|w| w.bytes().fold(7u64, |h, b| mix(h ^ b as u64))).collect()
    }

    #[test]
    fn the_estimate_tracks_the_jaccard_and_identical_sets_are_one() {
        let a: String = (0..200).map(|i| format!("t{} ", i % 37)).collect::<String>() + &(0..200).map(|i| format!("u{i} ")).collect::<String>();
        let sa = signature(&classes(&a)).unwrap();
        assert_eq!(estimate(&sa, &sa), 1.0);
        // one token changed in a long sequence: still near
        let b = a.replacen("u100 ", "zz ", 1);
        let sb = signature(&classes(&b)).unwrap();
        assert!(estimate(&sa, &sb) >= 0.85, "{}", estimate(&sa, &sb));
        // a different sequence: far
        let c: String = (0..400).map(|i| format!("w{i} ")).collect();
        let sc = signature(&classes(&c)).unwrap();
        assert!(estimate(&sa, &sc) < 0.2);
        assert_eq!(from_bytes(&to_bytes(&sa)).unwrap(), sa);
        let (pairs, _) = near_pairs(&[&sa, &sb, &sc], NEAR_FLOOR, 1000);
        assert_eq!(pairs.iter().map(|p| (p.0, p.1)).collect::<Vec<_>>(), [(0, 1)]);
        assert!(signature(&classes("a b")).is_none());
    }
}