codeprism_dev_tools/
graphviz_export.rs

1//! GraphViz export utilities for AST visualization
2
3use anyhow::Result;
4use codeprism_core::{Edge, EdgeKind, Node, NodeKind};
5use serde::{Deserialize, Serialize};
6use std::collections::HashMap;
7use std::fmt;
8
9/// GraphViz exporter for generating DOT format graphs
10#[derive(Debug, Clone)]
11pub struct GraphVizExporter {
12    config: GraphVizConfig,
13}
14
15/// Configuration for GraphViz export
16#[derive(Debug, Clone, Serialize, Deserialize)]
17pub struct GraphVizConfig {
18    pub graph_name: String,
19    pub graph_type: GraphType,
20    pub node_options: NodeStyle,
21    pub edge_options: EdgeStyle,
22    pub layout_engine: LayoutEngine,
23    pub include_node_labels: bool,
24    pub include_edge_labels: bool,
25    pub max_label_length: usize,
26    pub group_by_type: bool,
27}
28
29impl Default for GraphVizConfig {
30    fn default() -> Self {
31        Self {
32            graph_name: "ast_graph".to_string(),
33            graph_type: GraphType::Directed,
34            node_options: NodeStyle::default(),
35            edge_options: EdgeStyle::default(),
36            layout_engine: LayoutEngine::Dot,
37            include_node_labels: true,
38            include_edge_labels: true,
39            max_label_length: 20,
40            group_by_type: false,
41        }
42    }
43}
44
45/// Graph type for GraphViz
46#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
47pub enum GraphType {
48    Directed,
49    Undirected,
50}
51
52/// Layout engine options
53#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
54pub enum LayoutEngine {
55    Dot,
56    Neato,
57    Circo,
58    Fdp,
59    Sfdp,
60    Twopi,
61}
62
63/// Node styling options
64#[derive(Debug, Clone, Serialize, Deserialize)]
65pub struct NodeStyle {
66    pub shape: String,
67    pub color: String,
68    pub fillcolor: String,
69    pub style: String,
70    pub fontname: String,
71    pub fontsize: u32,
72    pub node_type_colors: HashMap<String, String>,
73}
74
75impl Default for NodeStyle {
76    fn default() -> Self {
77        let mut node_type_colors = HashMap::new();
78        node_type_colors.insert("Function".to_string(), "#e1f5fe".to_string());
79        node_type_colors.insert("Class".to_string(), "#f3e5f5".to_string());
80        node_type_colors.insert("Variable".to_string(), "#fff3e0".to_string());
81        node_type_colors.insert("Import".to_string(), "#e8f5e8".to_string());
82
83        Self {
84            shape: "box".to_string(),
85            color: "black".to_string(),
86            fillcolor: "#f5f5f5".to_string(),
87            style: "filled".to_string(),
88            fontname: "Arial".to_string(),
89            fontsize: 10,
90            node_type_colors,
91        }
92    }
93}
94
95/// Edge styling options
96#[derive(Debug, Clone, Serialize, Deserialize)]
97pub struct EdgeStyle {
98    pub color: String,
99    pub style: String,
100    pub arrowhead: String,
101    pub fontname: String,
102    pub fontsize: u32,
103    pub edge_type_colors: HashMap<String, String>,
104}
105
106impl Default for EdgeStyle {
107    fn default() -> Self {
108        let mut edge_type_colors = HashMap::new();
109        edge_type_colors.insert("Calls".to_string(), "#2196f3".to_string());
110        edge_type_colors.insert("Imports".to_string(), "#4caf50".to_string());
111        edge_type_colors.insert("Reads".to_string(), "#ff9800".to_string());
112        edge_type_colors.insert("Writes".to_string(), "#f44336".to_string());
113
114        Self {
115            color: "gray".to_string(),
116            style: "solid".to_string(),
117            arrowhead: "normal".to_string(),
118            fontname: "Arial".to_string(),
119            fontsize: 8,
120            edge_type_colors,
121        }
122    }
123}
124
125/// GraphViz options for specific exports
126#[derive(Debug, Clone, Default)]
127pub struct GraphVizOptions {
128    pub title: Option<String>,
129    pub subtitle: Option<String>,
130    pub highlight_nodes: Vec<String>,
131    pub highlight_edges: Vec<String>,
132    pub filter_node_types: Option<Vec<NodeKind>>,
133    pub filter_edge_types: Option<Vec<EdgeKind>>,
134    pub cluster_by_file: bool,
135    pub show_spans: bool,
136}
137
138impl GraphVizExporter {
139    /// Create a new GraphViz exporter with default configuration
140    pub fn new() -> Self {
141        Self {
142            config: GraphVizConfig::default(),
143        }
144    }
145
146    /// Create a GraphViz exporter with custom configuration
147    pub fn with_config(config: GraphVizConfig) -> Self {
148        Self { config }
149    }
150
151    /// Export nodes and edges to GraphViz DOT format
152    pub fn export_nodes_and_edges(&self, nodes: &[Node], edges: &[Edge]) -> Result<String> {
153        self.export_with_options(nodes, edges, &GraphVizOptions::default())
154    }
155
156    /// Export with custom options
157    pub fn export_with_options(
158        &self,
159        nodes: &[Node],
160        edges: &[Edge],
161        options: &GraphVizOptions,
162    ) -> Result<String> {
163        let mut dot = String::new();
164
165        // Graph header
166        let graph_keyword = match self.config.graph_type {
167            GraphType::Directed => "digraph",
168            GraphType::Undirected => "graph",
169        };
170
171        dot.push_str(&format!(
172            "{} {} {{\n",
173            graph_keyword, self.config.graph_name
174        ));
175
176        // Graph attributes
177        self.write_graph_attributes(&mut dot, options);
178
179        // Filter nodes if requested
180        let filtered_nodes: Vec<_> = if let Some(ref filter_types) = options.filter_node_types {
181            nodes
182                .iter()
183                .filter(|n| filter_types.contains(&n.kind))
184                .collect()
185        } else {
186            nodes.iter().collect()
187        };
188
189        // Group nodes by file if clustering is enabled
190        if options.cluster_by_file {
191            self.write_clustered_nodes(&mut dot, &filtered_nodes, options)?;
192        } else {
193            self.write_nodes(&mut dot, &filtered_nodes, options)?;
194        }
195
196        // Filter and write edges
197        let filtered_edges: Vec<_> = if let Some(ref filter_types) = options.filter_edge_types {
198            edges
199                .iter()
200                .filter(|e| filter_types.contains(&e.kind))
201                .collect()
202        } else {
203            edges.iter().collect()
204        };
205
206        self.write_edges(&mut dot, &filtered_edges, &filtered_nodes, options)?;
207
208        dot.push_str("}\n");
209
210        Ok(dot)
211    }
212
213    /// Write graph-level attributes
214    fn write_graph_attributes(&self, dot: &mut String, options: &GraphVizOptions) {
215        dot.push_str("  // Graph attributes\n");
216        dot.push_str(&format!("  layout=\"{:?}\";\n", self.config.layout_engine).to_lowercase());
217        dot.push_str("  rankdir=\"TB\";\n");
218        dot.push_str("  splines=\"ortho\";\n");
219        dot.push_str("  nodesep=\"0.5\";\n");
220        dot.push_str("  ranksep=\"1.0\";\n");
221
222        // Add title if provided
223        if let Some(ref title) = options.title {
224            dot.push_str(&format!("  label=\"{}\";\n", self.escape_label(title)));
225            dot.push_str("  fontsize=\"16\";\n");
226            dot.push_str("  fontname=\"Arial Bold\";\n");
227        }
228
229        // Default node and edge attributes
230        dot.push_str("  // Default node attributes\n");
231        dot.push_str(&format!("  node [shape=\"{}\", style=\"{}\", fillcolor=\"{}\", color=\"{}\", fontname=\"{}\", fontsize=\"{}\"];\n",
232            self.config.node_options.shape,
233            self.config.node_options.style,
234            self.config.node_options.fillcolor,
235            self.config.node_options.color,
236            self.config.node_options.fontname,
237            self.config.node_options.fontsize
238        ));
239
240        dot.push_str("  // Default edge attributes\n");
241        dot.push_str(&format!("  edge [color=\"{}\", style=\"{}\", arrowhead=\"{}\", fontname=\"{}\", fontsize=\"{}\"];\n",
242            self.config.edge_options.color,
243            self.config.edge_options.style,
244            self.config.edge_options.arrowhead,
245            self.config.edge_options.fontname,
246            self.config.edge_options.fontsize
247        ));
248
249        dot.push('\n');
250    }
251
252    /// Write nodes to DOT format
253    fn write_nodes(
254        &self,
255        dot: &mut String,
256        nodes: &[&Node],
257        options: &GraphVizOptions,
258    ) -> Result<()> {
259        dot.push_str("  // Nodes\n");
260
261        for node in nodes {
262            let node_id = self.sanitize_id(&node.id.to_hex());
263            let mut attributes = Vec::new();
264
265            // Node label
266            if self.config.include_node_labels {
267                let label = self.create_node_label(node, options);
268                attributes.push(format!("label=\"{}\"", self.escape_label(&label)));
269            }
270
271            // Node color based on type
272            let node_type_str = format!("{:?}", node.kind);
273            if let Some(color) = self
274                .config
275                .node_options
276                .node_type_colors
277                .get(&node_type_str)
278            {
279                attributes.push(format!("fillcolor=\"{color}\""));
280            }
281
282            // Highlight if requested
283            if options.highlight_nodes.contains(&node.id.to_hex()) {
284                attributes.push("penwidth=\"3\"".to_string());
285                attributes.push("color=\"red\"".to_string());
286            }
287
288            // Write node
289            if attributes.is_empty() {
290                dot.push_str(&format!("  {node_id};\n"));
291            } else {
292                dot.push_str(&format!("  {} [{}];\n", node_id, attributes.join(", ")));
293            }
294        }
295
296        dot.push('\n');
297        Ok(())
298    }
299
300    /// Write nodes grouped by file (clusters)
301    fn write_clustered_nodes(
302        &self,
303        dot: &mut String,
304        nodes: &[&Node],
305        options: &GraphVizOptions,
306    ) -> Result<()> {
307        let mut file_groups: HashMap<_, Vec<_>> = HashMap::new();
308
309        for node in nodes {
310            let file_path = node.file.to_string_lossy();
311            file_groups
312                .entry(file_path.to_string())
313                .or_default()
314                .push(*node);
315        }
316
317        dot.push_str("  // Clustered nodes by file\n");
318
319        for (i, (file_path, file_nodes)) in file_groups.iter().enumerate() {
320            dot.push_str(&format!("  subgraph cluster_{i} {{\n"));
321            dot.push_str(&format!(
322                "    label=\"{}\";\n",
323                self.escape_label(file_path)
324            ));
325            dot.push_str("    style=\"filled\";\n");
326            dot.push_str("    fillcolor=\"#f0f0f0\";\n");
327            dot.push_str("    color=\"gray\";\n");
328
329            for node in file_nodes {
330                let node_id = self.sanitize_id(&node.id.to_hex());
331                let label = if self.config.include_node_labels {
332                    self.create_node_label(node, options)
333                } else {
334                    node.name.clone()
335                };
336
337                dot.push_str(&format!(
338                    "    {} [label=\"{}\"];\n",
339                    node_id,
340                    self.escape_label(&label)
341                ));
342            }
343
344            dot.push_str("  }\n");
345        }
346
347        dot.push('\n');
348        Ok(())
349    }
350
351    /// Write edges to DOT format
352    fn write_edges(
353        &self,
354        dot: &mut String,
355        edges: &[&Edge],
356        nodes: &[&Node],
357        options: &GraphVizOptions,
358    ) -> Result<()> {
359        let node_ids: std::collections::HashSet<_> = nodes.iter().map(|n| &n.id).collect();
360
361        dot.push_str("  // Edges\n");
362
363        let edge_connector = match self.config.graph_type {
364            GraphType::Directed => "->",
365            GraphType::Undirected => "--",
366        };
367
368        for edge in edges {
369            // Only include edges between filtered nodes
370            if !node_ids.contains(&edge.source) || !node_ids.contains(&edge.target) {
371                continue;
372            }
373
374            let source_id = self.sanitize_id(&edge.source.to_hex());
375            let target_id = self.sanitize_id(&edge.target.to_hex());
376
377            let mut attributes = Vec::new();
378
379            // Edge label
380            if self.config.include_edge_labels {
381                let label = format!("{:?}", edge.kind);
382                attributes.push(format!("label=\"{}\"", self.escape_label(&label)));
383            }
384
385            // Edge color based on type
386            let edge_type_str = format!("{:?}", edge.kind);
387            if let Some(color) = self
388                .config
389                .edge_options
390                .edge_type_colors
391                .get(&edge_type_str)
392            {
393                attributes.push(format!("color=\"{color}\""));
394            }
395
396            // Highlight if requested
397            let edge_id = format!("{}->{}", edge.source.to_hex(), edge.target.to_hex());
398            if options.highlight_edges.contains(&edge_id) {
399                attributes.push("penwidth=\"3\"".to_string());
400                attributes.push("color=\"red\"".to_string());
401            }
402
403            // Write edge
404            if attributes.is_empty() {
405                dot.push_str(&format!("  {source_id} {edge_connector} {target_id};\n"));
406            } else {
407                dot.push_str(&format!(
408                    "  {} {} {} [{}];\n",
409                    source_id,
410                    edge_connector,
411                    target_id,
412                    attributes.join(", ")
413                ));
414            }
415        }
416
417        dot.push('\n');
418        Ok(())
419    }
420
421    /// Create a label for a node
422    fn create_node_label(&self, node: &Node, options: &GraphVizOptions) -> String {
423        let mut label = node.name.clone();
424
425        if label.len() > self.config.max_label_length {
426            label.truncate(self.config.max_label_length - 3);
427            label.push_str("...");
428        }
429
430        // Add type information
431        if self.config.group_by_type {
432            label = format!("{}\n({:?})", label, node.kind);
433        }
434
435        // Add span information if requested
436        if options.show_spans {
437            label = format!(
438                "{}\n[{}..{}]",
439                label, node.span.start_byte, node.span.end_byte
440            );
441        }
442
443        label
444    }
445
446    /// Sanitize an ID for GraphViz
447    fn sanitize_id(&self, id: &str) -> String {
448        format!("node_{}", id.replace('-', "_"))
449    }
450
451    /// Escape a label for GraphViz
452    fn escape_label(&self, label: &str) -> String {
453        label
454            .replace('\\', "\\\\")
455            .replace('"', "\\\"")
456            .replace('\n', "\\n")
457            .replace('\r', "\\r")
458            .replace('\t', "\\t")
459    }
460
461    /// Export tree-sitter syntax tree to GraphViz
462    pub fn export_syntax_tree(&self, tree: &tree_sitter::Tree, source: &str) -> Result<String> {
463        let root_node = tree.root_node();
464        let mut dot = String::new();
465
466        dot.push_str(&format!("digraph {} {{\n", self.config.graph_name));
467        dot.push_str("  rankdir=\"TB\";\n");
468        dot.push_str("  node [shape=\"box\", style=\"filled\", fillcolor=\"lightblue\"];\n");
469
470        self.export_syntax_node_recursive(&mut dot, &root_node, source, 0)?;
471
472        dot.push_str("}\n");
473        Ok(dot)
474    }
475
476    /// Recursively export syntax tree nodes
477    fn export_syntax_node_recursive(
478        &self,
479        dot: &mut String,
480        node: &tree_sitter::Node,
481        source: &str,
482        depth: usize,
483    ) -> Result<()> {
484        let node_id = format!("syntax_{}_{}", depth, node.start_byte());
485        let mut label = node.kind().to_string();
486
487        // Add text content for leaf nodes
488        if node.child_count() == 0 {
489            if let Ok(text) = node.utf8_text(source.as_bytes()) {
490                if text.len() <= 20 {
491                    label = format!("{}\\n\"{}\"", label, self.escape_label(text));
492                } else {
493                    label = format!("{}\\n\"{}...\"", label, self.escape_label(&text[..17]));
494                }
495            }
496        }
497
498        dot.push_str(&format!("  {node_id} [label=\"{label}\"];\n"));
499
500        // Add edges to children
501        for i in 0..node.child_count() {
502            if let Some(child) = node.child(i) {
503                let child_id = format!("syntax_{}_{}", depth + 1, child.start_byte());
504                dot.push_str(&format!("  {node_id} -> {child_id};\n"));
505                self.export_syntax_node_recursive(dot, &child, source, depth + 1)?;
506            }
507        }
508
509        Ok(())
510    }
511}
512
513impl Default for GraphVizExporter {
514    fn default() -> Self {
515        Self::new()
516    }
517}
518
519impl fmt::Display for LayoutEngine {
520    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
521        match self {
522            LayoutEngine::Dot => write!(f, "dot"),
523            LayoutEngine::Neato => write!(f, "neato"),
524            LayoutEngine::Circo => write!(f, "circo"),
525            LayoutEngine::Fdp => write!(f, "fdp"),
526            LayoutEngine::Sfdp => write!(f, "sfdp"),
527            LayoutEngine::Twopi => write!(f, "twopi"),
528        }
529    }
530}
531
532#[cfg(test)]
533mod tests {
534    use super::*;
535    use codeprism_core::{EdgeKind, Language, NodeKind, Span};
536    use std::path::PathBuf;
537
538    fn create_test_node(_id: u64, name: &str, kind: NodeKind) -> Node {
539        let path = PathBuf::from("test.rs");
540        let span = Span::new(0, 10, 1, 1, 1, 10);
541        let repo_id = "test_repo";
542
543        Node {
544            id: codeprism_core::NodeId::new(repo_id, &path, &span, &kind),
545            kind,
546            name: name.to_string(),
547            file: path,
548            span,
549            lang: Language::Rust,
550            metadata: Default::default(),
551            signature: Default::default(),
552        }
553    }
554
555    fn create_test_edge(source_id: &str, target_id: &str, kind: EdgeKind) -> Edge {
556        Edge {
557            source: codeprism_core::NodeId::from_hex(source_id).unwrap(),
558            target: codeprism_core::NodeId::from_hex(target_id).unwrap(),
559            kind,
560        }
561    }
562
563    #[test]
564    fn test_graphviz_exporter_creation() {
565        let exporter = GraphVizExporter::new();
566        assert_eq!(exporter.config.graph_name, "ast_graph");
567        assert!(matches!(exporter.config.graph_type, GraphType::Directed));
568    }
569
570    #[test]
571    fn test_sanitize_id() {
572        let exporter = GraphVizExporter::new();
573        let sanitized = exporter.sanitize_id("abc-def-123");
574        assert_eq!(sanitized, "node_abc_def_123");
575    }
576
577    #[test]
578    fn test_escape_label() {
579        let exporter = GraphVizExporter::new();
580        let escaped = exporter.escape_label("test\n\"quoted\"");
581        assert_eq!(escaped, "test\\n\\\"quoted\\\"");
582    }
583
584    #[test]
585    fn test_export_simple_graph() {
586        let exporter = GraphVizExporter::new();
587        let nodes = vec![
588            create_test_node(1, "main", NodeKind::Function),
589            create_test_node(2, "helper", NodeKind::Function),
590        ];
591
592        let main_id = nodes[0].id.to_hex();
593        let helper_id = nodes[1].id.to_hex();
594
595        let edges = vec![create_test_edge(&main_id, &helper_id, EdgeKind::Calls)];
596
597        let dot = exporter.export_nodes_and_edges(&nodes, &edges).unwrap();
598        assert!(dot.contains("digraph ast_graph"));
599        assert!(dot.contains("node_"));
600        assert!(dot.contains("->"));
601    }
602}
codeprism_dev_tools/graphviz_export.rs

codeprism_dev_tools/
graphviz_export.rs
