Skip to main content

_diffctx/
graph_export.rs

1use std::path::{Path, PathBuf};
2use std::sync::Arc;
3
4use rustc_hash::FxHashMap;
5use serde::Serialize;
6
7use crate::graph::{EdgeCategory, Graph};
8use crate::project_graph::ProjectGraph;
9use crate::types::{Fragment, FragmentId};
10
11pub struct ProjectGraphView<'a> {
12    pub graph: &'a Graph,
13    pub fragments: &'a FxHashMap<FragmentId, Fragment>,
14    pub root_dir: Option<&'a Path>,
15}
16
17pub fn view_from_project_graph(pg: &ProjectGraph) -> (FxHashMap<FragmentId, Fragment>, &Path) {
18    let map: FxHashMap<FragmentId, Fragment> = pg
19        .fragments
20        .iter()
21        .map(|f| (f.id.clone(), f.clone()))
22        .collect();
23    (map, pg.root_dir.as_path())
24}
25
26#[derive(Debug, Clone, Serialize)]
27pub struct NodeRecord {
28    pub id: String,
29    pub label: String,
30    pub path: String,
31    pub lines: String,
32    pub kind: String,
33    pub symbol: String,
34    pub token_count: u32,
35}
36
37#[derive(Debug, Clone, Serialize)]
38pub struct EdgeRecord {
39    pub source: String,
40    pub source_symbol: String,
41    pub target: String,
42    pub target_symbol: String,
43    pub weight: f64,
44    pub category: String,
45}
46
47#[derive(Debug, Clone, Serialize)]
48pub struct GraphDocument {
49    pub name: String,
50    #[serde(rename = "type")]
51    pub doc_type: String,
52    pub node_count: usize,
53    pub edge_count: usize,
54    pub nodes: Vec<NodeRecord>,
55    pub edges: Vec<EdgeRecord>,
56}
57
58#[derive(Debug, Clone, Serialize)]
59pub struct GraphSummary {
60    pub node_count: usize,
61    pub edge_count: usize,
62    pub file_count: usize,
63    pub density: f64,
64    pub edge_type_counts: FxHashMap<String, usize>,
65    pub top_in_degree: Vec<TopInDegreeEntry>,
66}
67
68#[derive(Debug, Clone, Serialize)]
69pub struct TopInDegreeEntry {
70    pub label: String,
71    pub in_degree: usize,
72}
73
74fn relative_path(path: &str, root: Option<&Path>) -> String {
75    let Some(root) = root else {
76        return path.to_string();
77    };
78    let p = PathBuf::from(path);
79    match p.strip_prefix(root) {
80        Ok(rel) => rel.to_string_lossy().replace('\\', "/"),
81        Err(_) => path.to_string(),
82    }
83}
84
85fn file_name(path: &str) -> &str {
86    Path::new(path)
87        .file_name()
88        .and_then(|s| s.to_str())
89        .unwrap_or(path)
90}
91
92fn round4(value: f64) -> f64 {
93    if value.is_finite() {
94        (value * 10_000.0).round() / 10_000.0
95    } else {
96        value
97    }
98}
99
100fn node_record(view: &ProjectGraphView<'_>, fid: &FragmentId, frag: &Fragment) -> NodeRecord {
101    let rel_path = relative_path(&fid.path, view.root_dir);
102    let loc = format!(
103        "{}:{}-{}",
104        file_name(&fid.path),
105        fid.start_line,
106        fid.end_line
107    );
108    let label = match frag.symbol_name.as_deref() {
109        Some(name) if !name.is_empty() => format!("{name} ({loc})"),
110        _ => loc.clone(),
111    };
112    NodeRecord {
113        id: format!("{}:{}-{}", rel_path, fid.start_line, fid.end_line),
114        label,
115        path: rel_path,
116        lines: format!("{}-{}", fid.start_line, fid.end_line),
117        kind: frag.kind.as_str().to_string(),
118        symbol: frag.symbol_name.clone().unwrap_or_default(),
119        token_count: frag.token_count,
120    }
121}
122
123fn edge_record(
124    view: &ProjectGraphView<'_>,
125    src: &FragmentId,
126    dst: &FragmentId,
127    weight: f64,
128    category: EdgeCategory,
129) -> EdgeRecord {
130    let src_symbol = view
131        .fragments
132        .get(src)
133        .and_then(|f| f.symbol_name.clone())
134        .unwrap_or_default();
135    let dst_symbol = view
136        .fragments
137        .get(dst)
138        .and_then(|f| f.symbol_name.clone())
139        .unwrap_or_default();
140    EdgeRecord {
141        source: format!(
142            "{}:{}-{}",
143            relative_path(&src.path, view.root_dir),
144            src.start_line,
145            src.end_line
146        ),
147        source_symbol: src_symbol,
148        target: format!(
149            "{}:{}-{}",
150            relative_path(&dst.path, view.root_dir),
151            dst.start_line,
152            dst.end_line
153        ),
154        target_symbol: dst_symbol,
155        weight: round4(weight),
156        category: category.as_str().to_string(),
157    }
158}
159
160fn collect_sorted_nodes<'a>(
161    fragments: &'a FxHashMap<FragmentId, Fragment>,
162) -> Vec<(&'a FragmentId, &'a Fragment)> {
163    let mut entries: Vec<(&FragmentId, &Fragment)> = fragments.iter().collect();
164    entries.sort_by(|a, b| {
165        a.0.path
166            .as_ref()
167            .cmp(b.0.path.as_ref())
168            .then(a.0.start_line.cmp(&b.0.start_line))
169            .then(a.0.end_line.cmp(&b.0.end_line))
170    });
171    entries
172}
173
174fn collect_sorted_edges(graph: &Graph) -> Vec<(FragmentId, FragmentId, f64, EdgeCategory)> {
175    let mut entries: Vec<(FragmentId, FragmentId, f64, EdgeCategory)> = Vec::new();
176    graph.for_each_categorized_edge(|src, dst, cat| {
177        let weight = graph.forward_edge_weight(src, dst).unwrap_or(0.0);
178        entries.push((src.clone(), dst.clone(), weight, cat));
179    });
180    entries.sort_by(|a, b| {
181        let key_a = format!("({}, {})", a.0, a.1);
182        let key_b = format!("({}, {})", b.0, b.1);
183        key_a.cmp(&key_b)
184    });
185    entries
186}
187
188pub fn graph_to_document(view: &ProjectGraphView<'_>) -> GraphDocument {
189    let root_name = view
190        .root_dir
191        .and_then(|p| p.file_name())
192        .and_then(|s| s.to_str())
193        .unwrap_or("unknown")
194        .to_string();
195
196    let sorted_nodes = collect_sorted_nodes(view.fragments);
197    let nodes: Vec<NodeRecord> = sorted_nodes
198        .iter()
199        .map(|(fid, frag)| node_record(view, fid, frag))
200        .collect();
201
202    let sorted_edges = collect_sorted_edges(view.graph);
203    let edges: Vec<EdgeRecord> = sorted_edges
204        .iter()
205        .map(|(src, dst, weight, cat)| edge_record(view, src, dst, *weight, *cat))
206        .collect();
207
208    GraphDocument {
209        name: root_name,
210        doc_type: "project_graph".to_string(),
211        node_count: view.graph.node_count(),
212        edge_count: view.graph.edge_count(),
213        nodes,
214        edges,
215    }
216}
217
218pub fn graph_to_json_string(view: &ProjectGraphView<'_>) -> Result<String, serde_json::Error> {
219    let doc = graph_to_document(view);
220    let mut out = serde_json::to_string_pretty(&doc)?;
221    out.push('\n');
222    Ok(out)
223}
224
225fn escape_graphml(text: &str) -> String {
226    let mut out = String::with_capacity(text.len());
227    for ch in text.chars() {
228        match ch {
229            '&' => out.push_str("&amp;"),
230            '<' => out.push_str("&lt;"),
231            '>' => out.push_str("&gt;"),
232            '"' => out.push_str("&quot;"),
233            '\'' => out.push_str("&apos;"),
234            '\r' => out.push_str("&#xD;"),
235            other => out.push(other),
236        }
237    }
238    out
239}
240
241pub fn graph_to_graphml_string(view: &ProjectGraphView<'_>) -> String {
242    let doc = graph_to_document(view);
243    let mut out = String::new();
244
245    out.push_str("<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n");
246    out.push_str("<graphml xmlns=\"http://graphml.graphdrawing.org/graphml\"\n");
247    out.push_str("         xmlns:xsi=\"http://www.w3.org/2001/XMLSchema-instance\"\n");
248    out.push_str(
249        "         xsi:schemaLocation=\"http://graphml.graphdrawing.org/graphml \
250http://graphml.graphdrawing.org/dtds/graphml.dtd\">\n",
251    );
252
253    out.push_str("  <key id=\"d_path\" for=\"node\" attr.name=\"path\" attr.type=\"string\"/>\n");
254    out.push_str("  <key id=\"d_lines\" for=\"node\" attr.name=\"lines\" attr.type=\"string\"/>\n");
255    out.push_str("  <key id=\"d_kind\" for=\"node\" attr.name=\"kind\" attr.type=\"string\"/>\n");
256    out.push_str(
257        "  <key id=\"d_symbol\" for=\"node\" attr.name=\"symbol\" attr.type=\"string\"/>\n",
258    );
259    out.push_str(
260        "  <key id=\"d_tokens\" for=\"node\" attr.name=\"token_count\" attr.type=\"int\"/>\n",
261    );
262    out.push_str(
263        "  <key id=\"d_weight\" for=\"edge\" attr.name=\"weight\" attr.type=\"double\"/>\n",
264    );
265    out.push_str(
266        "  <key id=\"d_category\" for=\"edge\" attr.name=\"category\" attr.type=\"string\"/>\n",
267    );
268    out.push_str(&format!(
269        "  <graph id=\"{}\" edgedefault=\"directed\">\n",
270        escape_graphml(&doc.name)
271    ));
272
273    for node in &doc.nodes {
274        let nid = escape_graphml(&node.id);
275        out.push_str(&format!("    <node id=\"{nid}\">\n"));
276        out.push_str(&format!(
277            "      <data key=\"d_path\">{}</data>\n",
278            escape_graphml(&node.path)
279        ));
280        out.push_str(&format!(
281            "      <data key=\"d_lines\">{}</data>\n",
282            escape_graphml(&node.lines)
283        ));
284        out.push_str(&format!(
285            "      <data key=\"d_kind\">{}</data>\n",
286            escape_graphml(&node.kind)
287        ));
288        if !node.symbol.is_empty() {
289            out.push_str(&format!(
290                "      <data key=\"d_symbol\">{}</data>\n",
291                escape_graphml(&node.symbol)
292            ));
293        }
294        out.push_str(&format!(
295            "      <data key=\"d_tokens\">{}</data>\n",
296            node.token_count
297        ));
298        out.push_str("    </node>\n");
299    }
300
301    for (i, edge) in doc.edges.iter().enumerate() {
302        let src = escape_graphml(&edge.source);
303        let tgt = escape_graphml(&edge.target);
304        out.push_str(&format!(
305            "    <edge id=\"e{i}\" source=\"{src}\" target=\"{tgt}\">\n"
306        ));
307        out.push_str(&format!(
308            "      <data key=\"d_weight\">{}</data>\n",
309            edge.weight
310        ));
311        out.push_str(&format!(
312            "      <data key=\"d_category\">{}</data>\n",
313            escape_graphml(&edge.category)
314        ));
315        out.push_str("    </edge>\n");
316    }
317
318    out.push_str("  </graph>\n");
319    out.push_str("</graphml>\n");
320    out
321}
322
323fn collect_files(fragments: &FxHashMap<FragmentId, Fragment>) -> usize {
324    let mut paths: rustc_hash::FxHashSet<Arc<str>> = rustc_hash::FxHashSet::default();
325    for fid in fragments.keys() {
326        paths.insert(fid.path.clone());
327    }
328    paths.len()
329}
330
331fn compute_in_degree(graph: &Graph) -> FxHashMap<FragmentId, usize> {
332    let mut counts: FxHashMap<FragmentId, usize> = FxHashMap::default();
333    graph.for_each_categorized_edge(|_src, dst, _cat| {
334        *counts.entry(dst.clone()).or_insert(0) += 1;
335    });
336    counts
337}
338
339pub fn graph_summary(view: &ProjectGraphView<'_>, top_n: usize) -> GraphSummary {
340    let node_count = view.graph.node_count();
341    let edge_count = view.graph.edge_count();
342    let file_count = collect_files(view.fragments);
343
344    let density = if node_count > 1 {
345        edge_count as f64 / (node_count * (node_count - 1)) as f64
346    } else {
347        0.0
348    };
349
350    let mut type_counts: FxHashMap<String, usize> = FxHashMap::default();
351    view.graph.for_each_categorized_edge(|_src, _dst, cat| {
352        *type_counts.entry(cat.as_str().to_string()).or_insert(0) += 1;
353    });
354
355    let in_deg = compute_in_degree(view.graph);
356    let mut sorted: Vec<(FragmentId, usize)> = in_deg.into_iter().collect();
357    sorted.sort_by(|a, b| {
358        b.1.cmp(&a.1)
359            .then_with(|| a.0.path.as_ref().cmp(b.0.path.as_ref()))
360            .then_with(|| a.0.start_line.cmp(&b.0.start_line))
361    });
362
363    let top_in_degree: Vec<TopInDegreeEntry> = sorted
364        .into_iter()
365        .take(top_n)
366        .map(|(fid, deg)| {
367            let frag = view.fragments.get(&fid);
368            let label = match frag.and_then(|f| f.symbol_name.as_deref()) {
369                Some(name) if !name.is_empty() => name.to_string(),
370                _ => format!("{}:{}", file_name(&fid.path), fid.start_line),
371            };
372            TopInDegreeEntry {
373                label,
374                in_degree: deg,
375            }
376        })
377        .collect();
378
379    GraphSummary {
380        node_count,
381        edge_count,
382        file_count,
383        density,
384        edge_type_counts: type_counts,
385        top_in_degree,
386    }
387}
388
389#[cfg(test)]
390mod tests {
391    use super::*;
392    use crate::graph::build_graph;
393    use crate::types::{Fragment, FragmentId, FragmentKind};
394    use rustc_hash::{FxHashMap, FxHashSet};
395    use std::sync::Arc;
396
397    fn make_fragment(path: &str, start: u32, end: u32, symbol: Option<&str>) -> Fragment {
398        Fragment {
399            id: FragmentId::new(Arc::from(path), start, end),
400            kind: FragmentKind::Function,
401            content: Arc::from("body"),
402            identifiers: FxHashSet::default(),
403            token_count: 10,
404            symbol_name: symbol.map(|s| s.to_string()),
405        }
406    }
407
408    fn fragments_map(frags: &[Fragment]) -> FxHashMap<FragmentId, Fragment> {
409        frags.iter().map(|f| (f.id.clone(), f.clone())).collect()
410    }
411
412    #[test]
413    fn empty_graph_emits_valid_outputs() {
414        let graph = build_graph(&[], FxHashMap::default(), FxHashMap::default());
415        let frags = FxHashMap::default();
416        let view = ProjectGraphView {
417            graph: &graph,
418            fragments: &frags,
419            root_dir: None,
420        };
421
422        let json = graph_to_json_string(&view).expect("json ok");
423        assert!(json.contains("\"node_count\": 0"));
424        assert!(json.contains("\"edge_count\": 0"));
425        assert!(json.contains("\"nodes\": []"));
426        assert!(json.contains("\"edges\": []"));
427        assert!(json.ends_with('\n'));
428
429        let xml = graph_to_graphml_string(&view);
430        assert!(xml.starts_with("<?xml version=\"1.0\" encoding=\"UTF-8\"?>"));
431        assert!(xml.contains("<graph id=\"unknown\" edgedefault=\"directed\">"));
432        assert!(xml.contains("</graphml>"));
433        assert!(!xml.contains("<node "));
434        assert!(!xml.contains("<edge "));
435
436        let summary = graph_summary(&view, 5);
437        assert_eq!(summary.node_count, 0);
438        assert_eq!(summary.edge_count, 0);
439        assert_eq!(summary.file_count, 0);
440        assert!((summary.density - 0.0).abs() < 1e-12);
441        assert!(summary.top_in_degree.is_empty());
442    }
443
444    #[test]
445    fn json_round_trip_preserves_schema() {
446        let f1 = make_fragment("src/a.rs", 1, 10, Some("foo"));
447        let f2 = make_fragment("src/b.rs", 5, 20, None);
448        let f3 = make_fragment("src/c.rs", 1, 30, Some("baz<T>"));
449
450        let frags = vec![f1.clone(), f2.clone(), f3.clone()];
451        let mut edges: FxHashMap<(FragmentId, FragmentId), f64> = FxHashMap::default();
452        let mut cats: FxHashMap<(FragmentId, FragmentId), EdgeCategory> = FxHashMap::default();
453        edges.insert((f1.id.clone(), f2.id.clone()), 0.55);
454        cats.insert((f1.id.clone(), f2.id.clone()), EdgeCategory::Semantic);
455        edges.insert((f2.id.clone(), f3.id.clone()), 0.25);
456        cats.insert((f2.id.clone(), f3.id.clone()), EdgeCategory::Structural);
457
458        let graph = build_graph(&frags, edges, cats);
459        let frag_map = fragments_map(&frags);
460        let view = ProjectGraphView {
461            graph: &graph,
462            fragments: &frag_map,
463            root_dir: None,
464        };
465
466        let json = graph_to_json_string(&view).expect("json ok");
467        let parsed: serde_json::Value = serde_json::from_str(&json).expect("valid json");
468
469        assert_eq!(parsed["type"], "project_graph");
470        assert_eq!(parsed["node_count"], 3);
471        assert_eq!(parsed["edge_count"], 2);
472        let nodes = parsed["nodes"].as_array().expect("nodes array");
473        assert_eq!(nodes.len(), 3);
474        assert_eq!(nodes[0]["path"], "src/a.rs");
475        assert_eq!(nodes[0]["symbol"], "foo");
476        assert_eq!(nodes[0]["lines"], "1-10");
477        assert_eq!(nodes[0]["kind"], "function");
478        assert_eq!(nodes[1]["symbol"], "");
479
480        let edges_out = parsed["edges"].as_array().expect("edges array");
481        assert_eq!(edges_out.len(), 2);
482        assert!(
483            edges_out
484                .iter()
485                .any(|e| e["category"] == "semantic" && e["weight"].as_f64() == Some(0.55))
486        );
487    }
488
489    #[test]
490    fn summary_counts_correct_for_small_graph() {
491        let f1 = make_fragment("src/a.rs", 1, 10, Some("alpha"));
492        let f2 = make_fragment("src/a.rs", 20, 30, Some("beta"));
493        let f3 = make_fragment("src/b.rs", 1, 5, None);
494
495        let frags = vec![f1.clone(), f2.clone(), f3.clone()];
496        let mut edges: FxHashMap<(FragmentId, FragmentId), f64> = FxHashMap::default();
497        let mut cats: FxHashMap<(FragmentId, FragmentId), EdgeCategory> = FxHashMap::default();
498        edges.insert((f1.id.clone(), f3.id.clone()), 0.4);
499        cats.insert((f1.id.clone(), f3.id.clone()), EdgeCategory::Semantic);
500        edges.insert((f2.id.clone(), f3.id.clone()), 0.3);
501        cats.insert((f2.id.clone(), f3.id.clone()), EdgeCategory::Semantic);
502
503        let graph = build_graph(&frags, edges, cats);
504        let frag_map = fragments_map(&frags);
505        let view = ProjectGraphView {
506            graph: &graph,
507            fragments: &frag_map,
508            root_dir: None,
509        };
510
511        let summary = graph_summary(&view, 5);
512        assert_eq!(summary.node_count, 3);
513        assert_eq!(summary.edge_count, 2);
514        assert_eq!(summary.file_count, 2);
515        let expected_density = 2.0 / (3.0 * 2.0);
516        assert!((summary.density - expected_density).abs() < 1e-9);
517        assert_eq!(summary.edge_type_counts.get("semantic"), Some(&2));
518        assert_eq!(summary.top_in_degree.len(), 1);
519        assert_eq!(summary.top_in_degree[0].in_degree, 2);
520    }
521
522    #[test]
523    fn graphml_escapes_special_characters() {
524        let mut frag = make_fragment("src/<weird>.rs", 1, 10, Some("ham&eggs<\"x\"'>"));
525        frag.token_count = 7;
526        let frags = vec![frag.clone()];
527        let frag_map = fragments_map(&frags);
528
529        let graph = build_graph(&frags, FxHashMap::default(), FxHashMap::default());
530        let view = ProjectGraphView {
531            graph: &graph,
532            fragments: &frag_map,
533            root_dir: None,
534        };
535
536        let xml = graph_to_graphml_string(&view);
537        assert!(xml.contains("ham&amp;eggs&lt;&quot;x&quot;&apos;&gt;"));
538        assert!(xml.contains("src/&lt;weird&gt;.rs"));
539        assert!(!xml.contains("ham&eggs"));
540        assert!(xml.contains("<data key=\"d_tokens\">7</data>"));
541    }
542}