1use serde::{Deserialize, Serialize};
2use serde_json::Value;
3use std::collections::HashSet;
4
5#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
6pub struct HyperEdgeEntry {
7 pub id: String,
8 pub label: String,
9 #[serde(default)]
10 pub nodes: Vec<String>,
11 #[serde(default)]
12 pub relation: Option<String>,
13 #[serde(default)]
14 pub confidence: Option<String>,
15 #[serde(default)]
16 pub confidence_score: Option<f64>,
17 #[serde(default)]
18 pub source_file: Option<String>,
19}
20
21#[derive(Debug, Clone, Serialize, Deserialize)]
22pub struct HyperGraphData {
23 pub nodes: Vec<Value>,
24 pub edges: Vec<Value>,
25 pub hyperedges: Vec<HyperEdgeEntry>,
26}
27
28pub fn build_from_json(data: &Value) -> HyperGraphData {
29 let nodes = data["nodes"].as_array().cloned().unwrap_or_default();
30 let edges = data["edges"].as_array().cloned().unwrap_or_default();
31 let hyperedges = data
32 .get("hyperedges")
33 .and_then(|v| v.as_array())
34 .map(|arr| {
35 arr.iter()
36 .filter_map(|v| serde_json::from_value(v.clone()).ok())
37 .collect()
38 })
39 .unwrap_or_default();
40 HyperGraphData {
41 nodes,
42 edges,
43 hyperedges,
44 }
45}
46
47pub fn attach_hyperedges(existing: &mut Vec<HyperEdgeEntry>, new: &[Value]) {
48 let existing_ids: HashSet<String> = existing.iter().map(|h| h.id.clone()).collect();
49 let mut seen = existing_ids;
50 for v in new {
51 let id = match v.get("id").and_then(|i| i.as_str()) {
52 Some(id) => id.to_string(),
53 None => continue,
54 };
55 if seen.contains(&id) {
56 continue;
57 }
58 if let Ok(entry) = serde_json::from_value::<HyperEdgeEntry>(v.clone()) {
59 seen.insert(id);
60 existing.push(entry);
61 }
62 }
63}
64
65pub fn to_json(data: &HyperGraphData) -> Result<String, serde_json::Error> {
66 serde_json::to_string_pretty(data)
67}
68
69pub fn report_hyperedges_section(hyperedges: &[HyperEdgeEntry]) -> String {
70 if hyperedges.is_empty() {
71 return String::new();
72 }
73 let mut lines = vec!["## Hyperedges (group relationships)".to_string()];
74 for h in hyperedges {
75 let confidence_part = match (&h.confidence, h.confidence_score) {
76 (Some(c), Some(s)) => format!(" {} {}", c, s),
77 (Some(c), None) => format!(" {}", c),
78 _ => String::new(),
79 };
80 let nodes_part = h.nodes.join(", ");
81 lines.push(format!(
82 "- **{}**{} — {}",
83 h.label, confidence_part, nodes_part
84 ));
85 }
86 lines.join("\n")
87}
88
89#[cfg(test)]
90mod tests {
91 use super::*;
92 use serde_json::json;
93 use std::io::Write;
94 use tempfile::NamedTempFile;
95
96 fn sample_extraction() -> Value {
97 json!({
98 "nodes": [
99 {"id": "BasicAuth", "label": "BasicAuth", "file_type": "code", "source_file": "auth.py"},
100 {"id": "DigestAuth", "label": "DigestAuth", "file_type": "code", "source_file": "auth.py"},
101 {"id": "Request", "label": "Request", "file_type": "code", "source_file": "http.py"},
102 ],
103 "edges": [
104 {"source": "BasicAuth", "target": "Request", "relation": "uses", "confidence": "EXTRACTED", "confidence_score": 1.0, "source_file": "auth.py"}
105 ],
106 "hyperedges": [
107 {
108 "id": "auth_flow",
109 "label": "Auth Flow",
110 "nodes": ["BasicAuth", "DigestAuth", "Request"],
111 "relation": "participate_in",
112 "confidence": "INFERRED",
113 "confidence_score": 0.75,
114 "source_file": "auth.py"
115 }
116 ]
117 })
118 }
119
120 #[test]
121 fn test_build_from_json_stores_hyperedges() {
122 let g = build_from_json(&sample_extraction());
123 assert_eq!(g.hyperedges.len(), 1);
124 assert_eq!(g.hyperedges[0].id, "auth_flow");
125 }
126
127 #[test]
128 fn test_build_from_json_no_hyperedges() {
129 let mut data = sample_extraction();
130 data["hyperedges"] = json!([]);
131 let g = build_from_json(&data);
132 assert_eq!(g.hyperedges.len(), 0);
133 }
134
135 #[test]
136 fn test_build_from_json_missing_hyperedges_key() {
137 let data = json!({
138 "nodes": [],
139 "edges": []
140 });
141 let g = build_from_json(&data);
142 assert_eq!(g.hyperedges.len(), 0);
143 }
144
145 #[test]
146 fn test_attach_hyperedges_adds_new() {
147 let mut existing: Vec<HyperEdgeEntry> = vec![];
148 let new = vec![json!({"id": "auth_flow", "label": "Auth Flow", "nodes": ["A", "B", "C"]})];
149 attach_hyperedges(&mut existing, &new);
150 assert_eq!(existing.len(), 1);
151 }
152
153 #[test]
154 fn test_attach_hyperedges_deduplicates() {
155 let mut existing: Vec<HyperEdgeEntry> = vec![];
156 let h = json!({"id": "auth_flow", "label": "Auth Flow", "nodes": ["A", "B", "C"]});
157 attach_hyperedges(&mut existing, std::slice::from_ref(&h));
158 attach_hyperedges(&mut existing, &[h]);
159 assert_eq!(existing.len(), 1);
160 }
161
162 #[test]
163 fn test_attach_hyperedges_multiple_different_ids() {
164 let mut existing: Vec<HyperEdgeEntry> = vec![];
165 let new = vec![
166 json!({"id": "flow_a", "label": "Flow A", "nodes": ["A", "B"]}),
167 json!({"id": "flow_b", "label": "Flow B", "nodes": ["C", "D"]}),
168 ];
169 attach_hyperedges(&mut existing, &new);
170 assert_eq!(existing.len(), 2);
171 }
172
173 #[test]
174 fn test_attach_hyperedges_skips_entry_without_id() {
175 let mut existing: Vec<HyperEdgeEntry> = vec![];
176 let new = vec![json!({"label": "No ID", "nodes": ["A", "B"]})];
177 attach_hyperedges(&mut existing, &new);
178 assert_eq!(existing.len(), 0);
179 }
180
181 #[test]
182 fn test_to_json_includes_hyperedges() {
183 let g = build_from_json(&sample_extraction());
184 let json_str = to_json(&g).unwrap();
185 let parsed: Value = serde_json::from_str(&json_str).unwrap();
186 assert!(parsed.get("hyperedges").is_some());
187 let hyperedges = parsed["hyperedges"].as_array().unwrap();
188 assert_eq!(hyperedges.len(), 1);
189 assert_eq!(hyperedges[0]["id"], "auth_flow");
190 }
191
192 #[test]
193 fn test_to_json_hyperedges_empty_when_none() {
194 let mut data = sample_extraction();
195 data["hyperedges"] = json!([]);
196 let g = build_from_json(&data);
197 let json_str = to_json(&g).unwrap();
198 let parsed: Value = serde_json::from_str(&json_str).unwrap();
199 assert!(parsed.get("hyperedges").is_some());
200 assert_eq!(parsed["hyperedges"].as_array().unwrap().len(), 0);
201 }
202
203 #[test]
204 fn test_hyperedges_roundtrip_via_json_file() {
205 let g = build_from_json(&sample_extraction());
206 let json_str = to_json(&g).unwrap();
207
208 let mut tmp = NamedTempFile::new().unwrap();
209 tmp.write_all(json_str.as_bytes()).unwrap();
210 let path = tmp.path();
211
212 let loaded: Value = serde_json::from_str(&std::fs::read_to_string(path).unwrap()).unwrap();
213 let g2 = build_from_json(&loaded);
214 assert!(!g2.hyperedges.is_empty());
215 assert_eq!(g2.hyperedges[0].id, "auth_flow");
216 }
217
218 #[test]
219 fn test_report_includes_hyperedges_section() {
220 let g = build_from_json(&sample_extraction());
221 let report = report_hyperedges_section(&g.hyperedges);
222 assert!(report.contains("## Hyperedges (group relationships)"));
223 assert!(report.contains("Auth Flow"));
224 assert!(report.contains("INFERRED"));
225 assert!(report.contains("0.75"));
226 }
227
228 #[test]
229 fn test_report_includes_hyperedge_node_list() {
230 let g = build_from_json(&sample_extraction());
231 let report = report_hyperedges_section(&g.hyperedges);
232 assert!(report.contains("BasicAuth"));
233 assert!(report.contains("DigestAuth"));
234 }
235
236 #[test]
237 fn test_report_skips_hyperedges_section_when_empty() {
238 let report = report_hyperedges_section(&[]);
239 assert!(!report.contains("## Hyperedges"));
240 }
241
242 #[test]
243 fn test_report_skips_hyperedges_section_when_key_missing() {
244 let data = json!({"nodes": [], "edges": []});
245 let g = build_from_json(&data);
246 let report = report_hyperedges_section(&g.hyperedges);
247 assert!(!report.contains("## Hyperedges"));
248 }
249}