1use crate::attrs::{attr_f64, is_external, num_attr};
8use code_ranker_plugin_api::{attrs::AttrValue, graph::Graph, node::NodeId};
9use std::collections::{HashMap, HashSet};
10
11pub fn annotate_hk(graph: &mut Graph, flow_kinds: &HashSet<String>) {
14 let external_ids: HashSet<&str> = graph
15 .nodes
16 .iter()
17 .filter(|n| is_external(n))
18 .map(|n| n.id.as_str())
19 .collect();
20
21 let mut fan_in: HashMap<NodeId, HashSet<NodeId>> = HashMap::new();
22 let mut fan_out: HashMap<NodeId, HashSet<NodeId>> = HashMap::new();
23 let mut fan_out_ext: HashMap<NodeId, HashSet<NodeId>> = HashMap::new();
24
25 for edge in &graph.edges {
26 if !flow_kinds.contains(&edge.kind) {
27 continue;
28 }
29 let to_external = external_ids.contains(edge.target.as_str());
30 let from_external = external_ids.contains(edge.source.as_str());
31 if to_external {
32 fan_out_ext
33 .entry(edge.source.clone())
34 .or_default()
35 .insert(edge.target.clone());
36 continue;
37 }
38 if from_external {
39 continue;
40 }
41 fan_out
42 .entry(edge.source.clone())
43 .or_default()
44 .insert(edge.target.clone());
45 fan_in
46 .entry(edge.target.clone())
47 .or_default()
48 .insert(edge.source.clone());
49 }
50
51 for node in &mut graph.nodes {
52 if is_external(node) {
53 continue;
54 }
55 let fi = fan_in.get(&node.id).map(|s| s.len()).unwrap_or(0);
56 let fo = fan_out.get(&node.id).map(|s| s.len()).unwrap_or(0);
57 let foe = fan_out_ext.get(&node.id).map(|s| s.len()).unwrap_or(0);
58 let loc = attr_f64(node, "sloc")
61 .or_else(|| attr_f64(node, "loc"))
62 .unwrap_or(0.0);
63 let hk = loc * ((fi * fo) as f64).powi(2);
64
65 set_or_clear(node, "fan_in", fi as f64);
66 set_or_clear(node, "fan_out", fo as f64);
67 set_or_clear(node, "fan_out_external", foe as f64);
68 if hk > 0.0 {
69 node.attrs.insert("hk".to_string(), num_attr(hk));
70 } else {
71 node.attrs.remove("hk");
72 }
73 }
74}
75
76fn set_or_clear(node: &mut code_ranker_plugin_api::node::Node, key: &str, v: f64) {
77 if v > 0.0 {
78 node.attrs.insert(key.to_string(), AttrValue::Int(v as i64));
79 } else {
80 node.attrs.remove(key);
81 }
82}
83
84#[cfg(test)]
85mod tests {
86 use super::*;
87 use code_ranker_plugin_api::{edge::Edge, node::Node};
88
89 fn file(id: &str, sloc: i64) -> Node {
90 let mut n = Node {
91 id: id.into(),
92 kind: "file".into(),
93 name: id.into(),
94 parent: None,
95 attrs: Default::default(),
96 };
97 n.attrs.insert("sloc".into(), AttrValue::Int(sloc));
98 n
99 }
100 fn uses(from: &str, to: &str) -> Edge {
101 Edge {
102 source: from.into(),
103 target: to.into(),
104 kind: "uses".into(),
105 line: None,
106 attrs: Default::default(),
107 }
108 }
109 fn flow() -> HashSet<String> {
110 HashSet::from(["uses".to_string()])
111 }
112
113 #[test]
114 fn hk_is_loc_times_fan_squared() {
115 let mut g = Graph {
116 nodes: vec![file("A", 4), file("B", 10), file("C", 5)],
117 edges: vec![uses("A", "B"), uses("B", "C")],
118 };
119 annotate_hk(&mut g, &flow());
120 let b = &g.nodes[1];
121 assert_eq!(attr_f64(b, "fan_in"), Some(1.0));
122 assert_eq!(attr_f64(b, "fan_out"), Some(1.0));
123 assert_eq!(attr_f64(b, "hk"), Some(10.0));
124 }
125
126 #[test]
127 fn external_target_counts_as_fan_out_external() {
128 let mut g = Graph {
129 nodes: vec![
130 file("a", 5),
131 Node {
132 id: "ext:x".into(),
133 kind: "external".into(),
134 name: "x".into(),
135 parent: None,
136 attrs: Default::default(),
137 },
138 ],
139 edges: vec![uses("a", "ext:x")],
140 };
141 annotate_hk(&mut g, &flow());
142 let a = &g.nodes[0];
143 assert_eq!(attr_f64(a, "fan_out_external"), Some(1.0));
144 assert_eq!(a.attrs.get("hk"), None, "no internal coupling → no hk");
145 }
146}