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!("  {};\n", node_id));
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_{} {{\n", i));
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!(
406                    "  {} {} {};\n",
407                    source_id, edge_connector, target_id
408                ));
409            } else {
410                dot.push_str(&format!(
411                    "  {} {} {} [{}];\n",
412                    source_id,
413                    edge_connector,
414                    target_id,
415                    attributes.join(", ")
416                ));
417            }
418        }
419
420        dot.push('\n');
421        Ok(())
422    }
423
424    /// Create a label for a node
425    fn create_node_label(&self, node: &Node, options: &GraphVizOptions) -> String {
426        let mut label = node.name.clone();
427
428        if label.len() > self.config.max_label_length {
429            label.truncate(self.config.max_label_length - 3);
430            label.push_str("...");
431        }
432
433        // Add type information
434        if self.config.group_by_type {
435            label = format!("{}\n({:?})", label, node.kind);
436        }
437
438        // Add span information if requested
439        if options.show_spans {
440            label = format!(
441                "{}\n[{}..{}]",
442                label, node.span.start_byte, node.span.end_byte
443            );
444        }
445
446        label
447    }
448
449    /// Sanitize an ID for GraphViz
450    fn sanitize_id(&self, id: &str) -> String {
451        format!("node_{}", id.replace('-', "_"))
452    }
453
454    /// Escape a label for GraphViz
455    fn escape_label(&self, label: &str) -> String {
456        label
457            .replace('\\', "\\\\")
458            .replace('"', "\\\"")
459            .replace('\n', "\\n")
460            .replace('\r', "\\r")
461            .replace('\t', "\\t")
462    }
463
464    /// Export tree-sitter syntax tree to GraphViz
465    pub fn export_syntax_tree(&self, tree: &tree_sitter::Tree, source: &str) -> Result<String> {
466        let root_node = tree.root_node();
467        let mut dot = String::new();
468
469        dot.push_str(&format!("digraph {} {{\n", self.config.graph_name));
470        dot.push_str("  rankdir=\"TB\";\n");
471        dot.push_str("  node [shape=\"box\", style=\"filled\", fillcolor=\"lightblue\"];\n");
472
473        self.export_syntax_node_recursive(&mut dot, &root_node, source, 0)?;
474
475        dot.push_str("}\n");
476        Ok(dot)
477    }
478
479    /// Recursively export syntax tree nodes
480    fn export_syntax_node_recursive(
481        &self,
482        dot: &mut String,
483        node: &tree_sitter::Node,
484        source: &str,
485        depth: usize,
486    ) -> Result<()> {
487        let node_id = format!("syntax_{}_{}", depth, node.start_byte());
488        let mut label = node.kind().to_string();
489
490        // Add text content for leaf nodes
491        if node.child_count() == 0 {
492            if let Ok(text) = node.utf8_text(source.as_bytes()) {
493                if text.len() <= 20 {
494                    label = format!("{}\\n\"{}\"", label, self.escape_label(text));
495                } else {
496                    label = format!("{}\\n\"{}...\"", label, self.escape_label(&text[..17]));
497                }
498            }
499        }
500
501        dot.push_str(&format!("  {} [label=\"{}\"];\n", node_id, label));
502
503        // Add edges to children
504        for i in 0..node.child_count() {
505            if let Some(child) = node.child(i) {
506                let child_id = format!("syntax_{}_{}", depth + 1, child.start_byte());
507                dot.push_str(&format!("  {} -> {};\n", node_id, child_id));
508                self.export_syntax_node_recursive(dot, &child, source, depth + 1)?;
509            }
510        }
511
512        Ok(())
513    }
514}
515
516impl Default for GraphVizExporter {
517    fn default() -> Self {
518        Self::new()
519    }
520}
521
522impl fmt::Display for LayoutEngine {
523    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
524        match self {
525            LayoutEngine::Dot => write!(f, "dot"),
526            LayoutEngine::Neato => write!(f, "neato"),
527            LayoutEngine::Circo => write!(f, "circo"),
528            LayoutEngine::Fdp => write!(f, "fdp"),
529            LayoutEngine::Sfdp => write!(f, "sfdp"),
530            LayoutEngine::Twopi => write!(f, "twopi"),
531        }
532    }
533}
534
535#[cfg(test)]
536mod tests {
537    use super::*;
538    use codeprism_core::{EdgeKind, Language, NodeKind, Span};
539    use std::path::PathBuf;
540
541    fn create_test_node(_id: u64, name: &str, kind: NodeKind) -> Node {
542        let path = PathBuf::from("test.rs");
543        let span = Span::new(0, 10, 1, 1, 1, 10);
544        let repo_id = "test_repo";
545
546        Node {
547            id: codeprism_core::NodeId::new(repo_id, &path, &span, &kind),
548            kind,
549            name: name.to_string(),
550            file: path,
551            span,
552            lang: Language::Rust,
553            metadata: Default::default(),
554            signature: Default::default(),
555        }
556    }
557
558    fn create_test_edge(source_id: &str, target_id: &str, kind: EdgeKind) -> Edge {
559        Edge {
560            source: codeprism_core::NodeId::from_hex(source_id).unwrap(),
561            target: codeprism_core::NodeId::from_hex(target_id).unwrap(),
562            kind,
563        }
564    }
565
566    #[test]
567    fn test_graphviz_exporter_creation() {
568        let exporter = GraphVizExporter::new();
569        assert_eq!(exporter.config.graph_name, "ast_graph");
570        assert!(matches!(exporter.config.graph_type, GraphType::Directed));
571    }
572
573    #[test]
574    fn test_sanitize_id() {
575        let exporter = GraphVizExporter::new();
576        let sanitized = exporter.sanitize_id("abc-def-123");
577        assert_eq!(sanitized, "node_abc_def_123");
578    }
579
580    #[test]
581    fn test_escape_label() {
582        let exporter = GraphVizExporter::new();
583        let escaped = exporter.escape_label("test\n\"quoted\"");
584        assert_eq!(escaped, "test\\n\\\"quoted\\\"");
585    }
586
587    #[test]
588    fn test_export_simple_graph() {
589        let exporter = GraphVizExporter::new();
590        let nodes = vec![
591            create_test_node(1, "main", NodeKind::Function),
592            create_test_node(2, "helper", NodeKind::Function),
593        ];
594
595        let main_id = nodes[0].id.to_hex();
596        let helper_id = nodes[1].id.to_hex();
597
598        let edges = vec![create_test_edge(&main_id, &helper_id, EdgeKind::Calls)];
599
600        let dot = exporter.export_nodes_and_edges(&nodes, &edges).unwrap();
601        assert!(dot.contains("digraph ast_graph"));
602        assert!(dot.contains("node_"));
603        assert!(dot.contains("->"));
604    }
605}
codeprism_dev_tools/graphviz_export.rs

codeprism_dev_tools/
graphviz_export.rs
