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