srcgraph_metrics/

entropy.rs

//! Shannon entropy over per-class identifier-token frequencies.
//!
//! For each class, the visiting-tool extractor emits a `methodTokens` blob:
//!
//!   `{"methods": [{"name": str, "tokens": [str]}, …]}`
//!
//! We pool every token across every method in the class, build a frequency
//! histogram, and compute
//!
//!   H(X) = -Σᵢ pᵢ · log₂(pᵢ)            (bits; 0·log 0 ≡ 0 by convention)
//!
//! Low entropy → vocabulary is concentrated on a few terms (focused class);
//! high entropy → vocabulary is spread across many terms (sprawling class).
//! A class with ≤ 1 distinct token has H = 0 (no uncertainty).
//!
//! See `DESIGN.md` (Phase 2) at the workspace root.

use petgraph::Graph;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

use srcgraph_core::{ClassNode, EdgeKind};

/// Per-class entropy readout.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TokenEntropy {
    pub class_id: String,
    /// `None` when `methodTokens` is absent or unparseable.
    pub entropy: Option<f64>,
    /// Number of *distinct* tokens that fed into the histogram (vocabulary size).
    pub distinct: usize,
    /// Total token occurrences across all methods.
    pub total: usize,
}

/// Shannon entropy of a probability distribution (bits).
/// Zero-probability entries are skipped (0·log 0 ≡ 0).
pub fn shannon_entropy(probs: &[f64]) -> f64 {
    let mut h = 0.0;
    for &p in probs {
        if p > 0.0 {
            h -= p * p.log2();
        }
    }
    h
}

/// Shannon entropy over a token-count histogram.
pub fn entropy_from_counts(counts: &HashMap<String, u64>) -> f64 {
    let total: u64 = counts.values().sum();
    if total == 0 {
        return 0.0;
    }
    let total_f = total as f64;
    let probs: Vec<f64> = counts.values().map(|c| *c as f64 / total_f).collect();
    shannon_entropy(&probs)
}

/// Pool every `tokens` array across every method in a `methodTokens` blob into
/// a single frequency histogram.
pub fn token_histogram(blob: &serde_json::Value) -> Option<HashMap<String, u64>> {
    let methods = blob.get("methods")?.as_array()?;
    let mut counts: HashMap<String, u64> = HashMap::new();
    for m in methods {
        let Some(tokens) = m.get("tokens").and_then(|v| v.as_array()) else {
            continue;
        };
        for t in tokens {
            if let Some(s) = t.as_str() {
                *counts.entry(s.to_owned()).or_insert(0) += 1;
            }
        }
    }
    Some(counts)
}

/// Compute per-class identifier-token entropy for every node.
pub fn compute_entropy<N, E>(graph: &Graph<N, E>) -> Vec<TokenEntropy>
where
    N: ClassNode,
    E: EdgeKind,
{
    graph
        .node_indices()
        .map(|nx| {
            let node = &graph[nx];
            let class_id = node.id().to_owned();
            let Some(blob) = node.method_tokens() else {
                return TokenEntropy {
                    class_id,
                    entropy: None,
                    distinct: 0,
                    total: 0,
                };
            };
            let Some(counts) = token_histogram(blob) else {
                return TokenEntropy {
                    class_id,
                    entropy: None,
                    distinct: 0,
                    total: 0,
                };
            };
            let total: u64 = counts.values().sum();
            let distinct = counts.len();
            let entropy = if distinct <= 1 {
                Some(0.0)
            } else {
                Some(entropy_from_counts(&counts))
            };
            TokenEntropy {
                class_id,
                entropy,
                distinct,
                total: total as usize,
            }
        })
        .collect()
}

#[cfg(test)]
mod tests {
    use super::*;
    use srcgraph_core::{EdgeType, OwnedClassNode, OwnedGraph};
    use petgraph::Graph;
    use serde_json::json;

    fn class(id: &str, tokens: Option<serde_json::Value>) -> OwnedClassNode {
        OwnedClassNode {
            id: id.to_owned(),
            name: id.to_owned(),
            namespace: "test".to_owned(),
            line_count: 10,
            method_count: 1,
            halstead_eta1: 0,
            halstead_eta2: 0,
            halstead_n1: 0,
            halstead_n2: 0,
            method_connectivity: None,
            method_fingerprints: None,
            method_tokens: tokens,
            call_sequences: None,
            cyclomatic_complexity: None,
            path_conditions: None,
            invariants: None,
            error_messages: None,
            magic_numbers: None,
            dead_code: None,
            tenant_branches: None,
            state_transitions: None,
        }
    }

    #[test]
    fn shannon_of_uniform_two_outcomes_is_one_bit() {
        assert!((shannon_entropy(&[0.5, 0.5]) - 1.0).abs() < 1e-12);
    }

    #[test]
    fn shannon_of_certain_distribution_is_zero() {
        assert_eq!(shannon_entropy(&[1.0, 0.0, 0.0]), 0.0);
    }

    #[test]
    fn shannon_of_uniform_four_is_two_bits() {
        let p = vec![0.25; 4];
        assert!((shannon_entropy(&p) - 2.0).abs() < 1e-12);
    }

    #[test]
    fn entropy_from_counts_matches_uniform_formula() {
        // Four tokens, each appearing 3 times — uniform over 4 → H = log2(4) = 2.
        let mut counts: HashMap<String, u64> = HashMap::new();
        counts.insert("a".into(), 3);
        counts.insert("b".into(), 3);
        counts.insert("c".into(), 3);
        counts.insert("d".into(), 3);
        assert!((entropy_from_counts(&counts) - 2.0).abs() < 1e-12);
    }

    #[test]
    fn token_histogram_pools_across_methods() {
        let blob = json!({"methods": [
            {"name": "foo", "tokens": ["x", "y", "x"]},
            {"name": "bar", "tokens": ["y", "z"]}
        ]});
        let h = token_histogram(&blob).unwrap();
        assert_eq!(h["x"], 2);
        assert_eq!(h["y"], 2);
        assert_eq!(h["z"], 1);
    }

    #[test]
    fn single_distinct_token_yields_zero_entropy() {
        let blob = json!({"methods": [
            {"name": "foo", "tokens": ["x", "x", "x"]}
        ]});
        let mut g: OwnedGraph = Graph::new();
        g.add_node(class("A", Some(blob)));
        let rows = compute_entropy(&g);
        assert_eq!(rows[0].distinct, 1);
        assert_eq!(rows[0].entropy, Some(0.0));
        assert_eq!(rows[0].total, 3);
    }

    #[test]
    fn missing_blob_yields_none() {
        let mut g: OwnedGraph = Graph::new();
        let a = g.add_node(class("A", None));
        let b = g.add_node(class("B", Some(json!({"methods": [{"name": "m", "tokens": ["a", "b"]}]}))));
        g.add_edge(a, b, EdgeType::MethodCall);

        let rows = compute_entropy(&g);
        let ra = rows.iter().find(|r| r.class_id == "A").unwrap();
        let rb = rows.iter().find(|r| r.class_id == "B").unwrap();
        assert!(ra.entropy.is_none());
        assert_eq!(rb.distinct, 2);
        assert!((rb.entropy.unwrap() - 1.0).abs() < 1e-12);
    }

    #[test]
    fn unparseable_blob_yields_none() {
        // Wrong shape — `methods` is a string, not an array.
        let blob = json!({"methods": "not an array"});
        let mut g: OwnedGraph = Graph::new();
        g.add_node(class("A", Some(blob)));
        let rows = compute_entropy(&g);
        assert!(rows[0].entropy.is_none());
    }
}