use petgraph::Graph;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use srcgraph_core::{ClassNode, EdgeKind};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TokenEntropy {
pub class_id: String,
pub entropy: Option<f64>,
pub distinct: usize,
pub total: usize,
}
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
}
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)
}
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)
}
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() {
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() {
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());
}
}