traverse_graph/
cg_dot.rs

1//! DOT (Graphviz) format export for Call Graphs.
2//!
3//! Provides functionality to convert a `CallGraph` into a DOT language string,
4//! suitable for visualization with tools like Graphviz. Includes default formatting
5//! and allows customization via closures.
6
7use crate::cg::{CallGraph, Edge, EdgeType, Node, NodeType};
8use std::collections::HashSet; // Import HashSet
9use std::fmt::Write;
10use tracing::debug;
11
12/// Configuration options for DOT export.
13#[derive(Debug, Clone, PartialEq, Eq)]
14pub struct DotExportConfig {
15    /// If true, nodes with no incoming or outgoing edges will be excluded from the output.
16    pub exclude_isolated_nodes: bool,
17}
18
19impl Default for DotExportConfig {
20    fn default() -> Self {
21        Self {
22            exclude_isolated_nodes: false, // Default is to include isolated nodes
23        }
24    }
25}
26
27pub trait ToDotLabel {
28    fn to_dot_label(&self) -> String;
29}
30
31pub trait ToDotAttributes {
32    fn to_dot_attributes(&self) -> Vec<(String, String)>;
33}
34
35impl ToDotLabel for &str {
36    fn to_dot_label(&self) -> String {
37        escape_dot_string(self)
38    }
39}
40
41impl ToDotLabel for String {
42    fn to_dot_label(&self) -> String {
43        escape_dot_string(self)
44    }
45}
46
47impl<T> ToDotAttributes for T {
48    fn to_dot_attributes(&self) -> Vec<(String, String)> {
49        Vec::new()
50    }
51}
52
53pub trait CgToDot {
54    /// Exports the graph to DOT format using default label/attribute generation.
55    ///
56    /// # Arguments
57    ///
58    /// * `name` - The name of the graph in the DOT output.
59    ///
60    /// # Returns
61    ///
62    /// A string containing the DOT representation of the graph.
63    fn to_dot(&self, name: &str, config: &DotExportConfig) -> String;
64
65    /// Exports the graph to DOT format with custom node and edge formatters.
66    ///
67    /// # Arguments
68    ///
69    /// * `name` - The name of the graph.
70    /// * `config` - Configuration options for the export.
71    /// * `node_formatter` - A closure that takes a `&Node` and returns its DOT attributes.
72    /// * `edge_formatter` - A closure that takes an `&Edge` and returns its DOT attributes.
73    ///
74    /// # Returns
75    ///
76    /// A string containing the DOT representation of the graph.
77    fn to_dot_with_formatters<NF, EF>(
78        &self,
79        name: &str,
80        config: &DotExportConfig,
81        node_formatter: NF,
82        edge_formatter: EF,
83    ) -> String
84    where
85        NF: Fn(&Node) -> Vec<(String, String)>,
86        EF: Fn(&Edge) -> Vec<(String, String)>;
87}
88
89impl CgToDot for CallGraph {
90    /// Default DOT export using `ToDotLabel` and `ToDotAttributes` implementations
91    /// for `Node` and `Edge`.
92    fn to_dot(&self, name: &str, config: &DotExportConfig) -> String {
93        self.to_dot_with_formatters(
94            name,
95            config, // Pass config
96
97            |node| {
98                let mut attrs = vec![
99                    ("label".to_string(), escape_dot_string(&node.to_dot_label())),
100                    (
101                        "tooltip".to_string(),
102                        escape_dot_string(&format!(
103                            "Type: {:?}\\nVisibility: {:?}\\nSpan: {:?}",
104                            node.node_type, node.visibility, node.span
105                        )),
106                    ),
107                    (
108                        "fillcolor".to_string(),
109                        match node.node_type {
110                            NodeType::Function => "lightblue".to_string(),
111                            NodeType::Constructor => "lightgoldenrodyellow".to_string(),
112                            NodeType::Modifier => "lightcoral".to_string(),
113                            NodeType::Library => "lightgrey".to_string(), // Keep lightgrey for Library
114                            NodeType::Interface => "lightpink".to_string(), // Use lightpink for Interface
115                            NodeType::StorageVariable => "khaki".to_string(), // Added color for StorageVariable
116                            NodeType::Evm => "gray".to_string(), // Added color for EVM
117                            NodeType::EventListener => "lightcyan".to_string(), // Added color for EventListener
118                            NodeType::RequireCondition => "orange".to_string(), // Added color for RequireCondition
119                            NodeType::IfStatement => "mediumpurple1".to_string(), // Added color for IfStatement
120                            NodeType::ThenBlock => "palegreen".to_string(),    // Added color for ThenBlock
121                            NodeType::ElseBlock => "lightsalmon".to_string(),  // Added color for ElseBlock
122                            NodeType::WhileStatement => "lightsteelblue".to_string(), // Added color for WhileStatement
123                            NodeType::WhileBlock => "lightseagreen".to_string(), // Added color for WhileBlock
124                            NodeType::ForCondition => "darkseagreen1".to_string(), // Added color for ForCondition
125                            NodeType::ForBlock => "darkolivegreen1".to_string(), // Added color for ForBlock
126                        },
127                    ),
128                ];
129                attrs.extend(node.to_dot_attributes());
130                attrs
131            },
132            |edge| {
133                let mut attrs = Vec::new();
134                match edge.edge_type {
135                    EdgeType::Call => {
136                        // Construct base tooltip
137                        let mut tooltip = format!("Call Site Span: {:?}", edge.call_site_span);
138
139                        // Construct the argument string
140                        let args_str = edge.argument_names.as_ref()
141                            .map(|args| {
142                                if args.is_empty() {
143                                    "".to_string() // No arguments, empty string inside parens
144                                } else {
145                                    // Escape each argument individually before joining
146                                    args.iter().map(|arg| escape_dot_string(arg)).collect::<Vec<_>>().join(", ")
147                                }
148                            })
149                            .unwrap_or_default(); // Default to empty string if None
150
151                        // Construct label and potentially add event info to tooltip
152                        let raw_label = if let Some(event_name) = &edge.event_name {
153                            // It's an emit-related edge
154                            write!(tooltip, "\\nEvent: {}", escape_dot_string(event_name)).unwrap();
155                            format!("emit {}({})\nSeq: {}", escape_dot_string(event_name), args_str, edge.sequence_number)
156                        } else {
157                            // It's a regular function call
158                            format!("({})\n{}", args_str, edge.sequence_number)
159                        };
160
161                        // Add sequence number to tooltip regardless
162                        write!(tooltip, "\\nSequence: {}", edge.sequence_number).unwrap();
163
164                        // Add final tooltip and label attributes
165                        attrs.push(("tooltip".to_string(), escape_dot_string(&tooltip)));
166                        attrs.push(("label".to_string(), escape_dot_string(&raw_label)));
167
168                        // Style emit edges differently? (Optional)
169                        if edge.event_name.is_some() {
170                            attrs.push(("color".to_string(), "blue".to_string()));
171                            attrs.push(("fontcolor".to_string(), "blue".to_string()));
172                        }
173                    }
174                    EdgeType::Return => {
175                        debug!("Formatting Return edge: {} -> {}", edge.source_node_id, edge.target_node_id);
176                        attrs.push((
177                            "tooltip".to_string(),
178                            escape_dot_string(&format!(
179                                "Return from function defined at {:?}\\nReturn Statement Span: {:?}\\nSequence: {}", // Added sequence to tooltip
180                                edge.call_site_span, // Span of the function definition
181                                edge.return_site_span.unwrap_or((0, 0)), // Span of the return statement
182                                edge.sequence_number // Sequence number (matches call)
183                            )),
184                        ));
185                        // Add returned value to tooltip if present
186                        if let Some(ref value) = edge.returned_value {
187                             if let Some((_, tooltip_val)) = attrs.last_mut() { // Find the tooltip attribute
188                                write!(tooltip_val, "\\nReturns: {}", escape_dot_string(value)).unwrap(); // Append to existing tooltip
189                             }
190                        }
191                        // Set label based on returned value
192                        let label = match &edge.returned_value {
193                            Some(value) => format!("ret {}", escape_dot_string(value)),
194                            None => "ret".to_string(),
195                        };
196                        attrs.push(("label".to_string(), label.clone()));
197                        // Style return edges differently
198                        attrs.push(("style".to_string(), "dashed".to_string()));
199                        attrs.push(("color".to_string(), "grey".to_string()));
200                        attrs.push(("arrowhead".to_string(), "empty".to_string()));
201
202                    }
203                    EdgeType::StorageRead => {
204                        let tooltip = format!("Read Span: {:?}", edge.call_site_span);
205                        attrs.push(("label".to_string(), "read".to_string()));
206                        attrs.push(("tooltip".to_string(), escape_dot_string(&tooltip)));
207                        attrs.push(("color".to_string(), "darkgreen".to_string()));
208                        attrs.push(("fontcolor".to_string(), "darkgreen".to_string()));
209                        attrs.push(("style".to_string(), "dotted".to_string()));
210                    }
211                    EdgeType::StorageWrite => {
212                        let tooltip = format!("Write Span: {:?}", edge.call_site_span);
213                        attrs.push(("label".to_string(), "write".to_string()));
214                        attrs.push(("tooltip".to_string(), escape_dot_string(&tooltip)));
215                        attrs.push(("color".to_string(), "darkred".to_string()));
216                        attrs.push(("fontcolor".to_string(), "darkred".to_string()));
217                        attrs.push(("style".to_string(), "bold".to_string()));
218                    }
219                    EdgeType::Require => {
220                        debug!("[DOT Require DEBUG] Formatting Require edge: {} -> {}", edge.source_node_id, edge.target_node_id);
221                        let tooltip = format!("Require Check Span: {:?}", edge.call_site_span);
222                        let args_str = edge.argument_names.as_ref()
223                            .map(|args| {
224                                if args.is_empty() {
225                                    "".to_string()
226                                } else {
227                                    args.iter().map(|arg| escape_dot_string(arg)).collect::<Vec<_>>().join(", ")
228                                }
229                            })
230                            .unwrap_or_default();
231                        let label = format!("require({})", args_str);
232                        attrs.push(("label".to_string(), escape_dot_string(&label)));
233                        attrs.push(("tooltip".to_string(), escape_dot_string(&tooltip)));
234                        attrs.push(("color".to_string(), "orange".to_string()));
235                        attrs.push(("fontcolor".to_string(), "orange".to_string()));
236                        attrs.push(("style".to_string(), "dashed".to_string()));
237                    }
238                    EdgeType::IfConditionBranch => {
239                        let condition = edge.argument_names.as_ref()
240                            .and_then(|args| args.first())
241                            .map(|arg| escape_dot_string(arg))
242                            .unwrap_or_else(|| "condition".to_string());
243                        let tooltip = format!("If Condition: {}\\nSpan: {:?}", condition, edge.call_site_span);
244                        attrs.push(("label".to_string(), format!("if ({})", condition)));
245                        attrs.push(("tooltip".to_string(), escape_dot_string(&tooltip)));
246                        attrs.push(("color".to_string(), "mediumpurple4".to_string()));
247                        attrs.push(("fontcolor".to_string(), "mediumpurple4".to_string()));
248                    }
249                    EdgeType::ThenBranch => {
250                        let tooltip = format!("Then branch taken\\nSpan: {:?}", edge.call_site_span);
251                        attrs.push(("label".to_string(), "then".to_string()));
252                        attrs.push(("tooltip".to_string(), escape_dot_string(&tooltip)));
253                        attrs.push(("color".to_string(), "green4".to_string()));
254                        attrs.push(("fontcolor".to_string(), "green4".to_string()));
255                    }
256                    EdgeType::ElseBranch => {
257                        let tooltip = format!("Else branch taken\\nSpan: {:?}", edge.call_site_span);
258                        attrs.push(("label".to_string(), "else".to_string()));
259                        attrs.push(("tooltip".to_string(), escape_dot_string(&tooltip)));
260                        attrs.push(("color".to_string(), "salmon4".to_string()));
261                        attrs.push(("fontcolor".to_string(), "salmon4".to_string()));
262                    }
263                    EdgeType::WhileConditionBranch | EdgeType::WhileBodyBranch => {
264                        // Placeholder for specific styling if needed, otherwise default edge style applies
265                    }
266                    EdgeType::ForConditionBranch | EdgeType::ForBodyBranch => {
267                        // Placeholder for specific styling if needed
268                        let label = if edge.edge_type == EdgeType::ForConditionBranch {
269                            edge.argument_names.as_ref()
270                                .and_then(|args| args.first())
271                                .map(|arg| escape_dot_string(arg))
272                                .unwrap_or_else(|| "for_cond".to_string())
273                        } else {
274                            "for_body".to_string()
275                        };
276                        attrs.push(("label".to_string(), label));
277                        attrs.push(("color".to_string(), "olivedrab".to_string()));
278                        attrs.push(("fontcolor".to_string(), "olivedrab".to_string()));
279                    }
280                }
281                // Allow overriding attributes from Edge::to_dot_attributes if needed
282                attrs.extend(edge.to_dot_attributes());
283                attrs
284            },
285        )
286    }
287
288    /// Generic DOT export allowing full customization via closures.
289    fn to_dot_with_formatters<NF, EF>(
290        &self,
291        name: &str,
292        config: &DotExportConfig, // Add config parameter
293        node_formatter: NF,
294        edge_formatter: EF,
295    ) -> String
296    where
297        NF: Fn(&Node) -> Vec<(String, String)>,
298        EF: Fn(&Edge) -> Vec<(String, String)>,
299    {
300        let mut dot_output = String::new();
301        let _ = writeln!(dot_output, "digraph \"{}\" {{", escape_dot_string(name));
302
303
304        let _ = writeln!(
305            dot_output,
306            "    graph [rankdir=LR, fontname=\"Arial\", splines=true];"
307        );
308        let _ = writeln!(
309            dot_output,
310            "    node [shape=box, style=\"rounded,filled\", fontname=\"Arial\"];"
311        );
312        let _ = writeln!(dot_output, "    edge [fontname=\"Arial\"];");
313        let _ = writeln!(dot_output);
314
315        // --- Node Filtering Logic ---
316        let connected_node_ids: Option<HashSet<usize>> = if config.exclude_isolated_nodes {
317            let mut ids = HashSet::new();
318            for edge in self.iter_edges() {
319                ids.insert(edge.source_node_id);
320                ids.insert(edge.target_node_id);
321            }
322            Some(ids)
323        } else {
324            None // No filtering needed
325        };
326
327        if config.exclude_isolated_nodes {
328            if let Some(ref connected_ids) = connected_node_ids {
329                debug!("Connected Node IDs: {:?}", connected_ids);
330            } else {
331                 debug!("Filtering active, but connected_node_ids is None (unexpected).");
332            }
333        }
334
335        for node in self.iter_nodes() {
336            // --- Apply Filtering ---
337            let is_isolated = if let Some(ref connected_ids) = connected_node_ids {
338                if !connected_ids.contains(&node.id) {
339                    debug!(
340                        "Skipping isolated node: ID={}, Name='{}', Contract='{:?}'",
341                        node.id, node.name, node.contract_name
342                    );
343                    continue; // Skip isolated node
344                }
345                false
346            } else {
347                false
348            };
349
350            if config.exclude_isolated_nodes { // Only log if filtering is active
351                 debug!(
352                     "Including {}node: ID={}, Name='{}', Contract='{:?}'",
353                     if is_isolated { "ISOLATED (ERROR?) " } else { "" }, // Highlight if it was marked isolated but not skipped
354                     node.id, node.name, node.contract_name
355                 );
356            }
357
358
359            let attrs = node_formatter(node);
360            let attrs_str = attrs
361                .iter()
362                .map(|(k, v)| format!("{}=\"{}\"", k, v))
363                .collect::<Vec<_>>()
364                .join(", ");
365            let _ = writeln!(dot_output, "    n{} [{}];", node.id, attrs_str);
366        }
367        let _ = writeln!(dot_output);
368
369        for edge in self.iter_edges() {
370            let attrs = edge_formatter(edge);
371            let attrs_str = attrs
372                .iter()
373                .map(|(k, v)| format!("{}=\"{}\"", k, v))
374                .collect::<Vec<_>>()
375                .join(", ");
376            let _ = writeln!(
377                dot_output,
378                "    n{} -> n{} [{}];",
379                edge.source_node_id, edge.target_node_id, attrs_str
380            );
381        }
382
383        let _ = writeln!(dot_output, "}}");
384        dot_output
385    }
386}
387
388/// Escapes characters in a string to be valid within a DOT label, tooltip, or attribute value.
389/// Ensures the output is enclosed in double quotes if needed (e.g., for labels/tooltips).
390pub fn escape_dot_string(s: &str) -> String {
391    s.replace('\\', "\\\\")
392        .replace('"', "\\\"")
393        .replace('\n', "\\n")
394        .replace('\r', "")
395        .replace('\t', "\\t")
396        .replace('{', "\\{")
397        .replace('}', "\\}")
398        .replace('<', "\\<")
399        .replace('>', "\\>")
400}
401
402#[cfg(test)]
403mod tests {
404    use super::*;
405    use crate::cg::{CallGraph, NodeType, Visibility};
406
407    fn create_test_graph() -> CallGraph {
408        let mut graph = CallGraph::new();
409        let n0 = graph.add_node(
410            "foo".to_string(),
411            NodeType::Function,
412            Some("ContractA".to_string()),
413            Visibility::Public,
414            (10, 20),
415        );
416        let n1 = graph.add_node(
417            "bar".to_string(),
418            NodeType::Function,
419            Some("ContractA".to_string()),
420            Visibility::Private,
421            (30, 40),
422        );
423        // Updated add_edge call to match the new signature
424        graph.add_edge(
425            n1,
426            n0,
427            EdgeType::Call, // edge_type
428            (35, 38),       // call_site_span
429            None,           // return_site_span
430            1,              // sequence_number
431            None,           // returned_value
432            None,           // argument_names
433            None,
434            None,
435        );
436        graph
437    }
438
439    #[test]
440    fn test_default_dot_export() {
441        let graph = create_test_graph();
442        let config = DotExportConfig::default(); // Use default config
443        let dot = graph.to_dot("TestDefault", &config); // Pass config
444        assert!(dot.starts_with("digraph \"TestDefault\" {"));
445        assert!(dot.contains("n0"), "Node n0 definition missing");
446        assert!(dot.contains("label=\"ContractA.foo\\n(function)\""), "Node n0 label incorrect");
447        assert!(dot.contains("tooltip=\"Type: Function\\\\nVisibility: Public\\\\nSpan: (10, 20)\""), "Node n0 tooltip incorrect");
448        assert!(dot.contains("fillcolor=\"lightblue\""), "Node n0 fillcolor incorrect");
449        assert!(dot.contains("n1"), "Node n1 definition missing");
450        assert!(dot.contains("label=\"ContractA.bar\\n(function)\""), "Node n1 label incorrect");
451        assert!(dot.contains("tooltip=\"Type: Function\\\\nVisibility: Private\\\\nSpan: (30, 40)\""), "Node n1 tooltip incorrect");
452        assert!(dot.contains("fillcolor=\"lightblue\""), "Node n1 fillcolor incorrect");
453        assert!(dot.contains("n1 -> n0"), "Edge n1 -> n0 missing");
454        // Check tooltip format (Call Site Span + Sequence)
455        assert!(dot.contains("tooltip=\"Call Site Span: (35, 38)\\\\nSequence: 1\""), "Edge tooltip incorrect"); // Updated to include sequence
456        // Check label format (Args + Sequence) - No args in this test case
457        assert!(dot.contains("label=\"()\\n1\""), "Edge label incorrect or missing sequence"); // Label still includes sequence
458        assert!(dot.ends_with("}\n"));
459    }
460
461    #[test]
462    fn test_custom_formatter_dot_export() {
463        let graph = create_test_graph();
464        let config = DotExportConfig::default(); // Use default config
465        let dot = graph.to_dot_with_formatters(
466            "TestCustom",
467            &config, // Pass config
468            |node| {
469                vec![
470                    ("label".to_string(), format!("N_{}", node.name)),
471                    ("color".to_string(), escape_dot_string("\"green\"")),
472                ]
473            },
474            |_edge| { // Prefix edge with underscore as it's unused in this custom formatter
475                vec![
476                    ("label".to_string(), escape_dot_string("\"CustomCall\"")),
477                    ("style".to_string(), escape_dot_string("\"dashed\"")),
478                ]
479            },
480        );
481
482        assert!(dot.starts_with("digraph \"TestCustom\" {"));
483        assert!(dot.contains("n0"), "Node n0 definition missing");
484        assert!(dot.contains("label=\"N_foo\""), "Node n0 label incorrect");
485        assert!(dot.contains("color=\"\\\"green\\\"\""), "Node n0 color incorrect");
486
487        assert!(dot.contains("n1"), "Node n1 definition missing");
488        assert!(dot.contains("label=\"N_bar\""), "Node n1 label incorrect");
489        assert!(dot.contains("color=\"\\\"green\\\"\""), "Node n1 color incorrect");
490
491        assert!(dot.contains("n1 -> n0"), "Edge n1 -> n0 missing");
492        assert!(dot.contains("label=\"\\\"CustomCall\\\"\""), "Edge label incorrect");
493        assert!(dot.contains("style=\"\\\"dashed\\\"\""), "Edge style incorrect");
494        assert!(dot.ends_with("}\n"));
495    }
496
497     #[test]
498    fn test_dot_escape_string_internal() {
499        assert_eq!(escape_dot_string(""), "");
500        assert_eq!(escape_dot_string("simple"), "simple");
501        assert_eq!(escape_dot_string("with \"quotes\""), "with \\\"quotes\\\"");
502        assert_eq!(escape_dot_string("new\nline"), "new\\nline");
503        assert_eq!(escape_dot_string("back\\slash"), "back\\\\slash");
504        assert_eq!(escape_dot_string("<html>"), "\\<html\\>");
505        assert_eq!(escape_dot_string("{record}"), "\\{record\\}");
506    }
507
508    #[test]
509    fn test_exclude_isolated_nodes() {
510        let mut graph = CallGraph::new();
511        let n0 = graph.add_node(
512            "connected1".to_string(),
513            NodeType::Function,
514            Some("ContractA".to_string()),
515            Visibility::Public,
516            (10, 20),
517        );
518        let n1 = graph.add_node(
519            "connected2".to_string(),
520            NodeType::Function,
521            Some("ContractA".to_string()),
522            Visibility::Private,
523            (30, 40),
524        );
525        let _n2 = graph.add_node(
526            "isolated".to_string(),
527            NodeType::Function,
528            Some("ContractB".to_string()),
529            Visibility::Public,
530            (50, 60),
531        );
532        // Add an edge to connect n0 and n1
533        graph.add_edge(
534            n0, n1, EdgeType::Call, (15, 18), None, 1, None, None, None, None,
535        );
536
537        // Test 1: Exclude isolated nodes
538        let config_exclude = DotExportConfig {
539            exclude_isolated_nodes: true,
540        };
541        let dot_excluded = graph.to_dot("TestExcludeIsolated", &config_exclude);
542
543        assert!(dot_excluded.contains("n0"), "Node n0 (connected) should be present when excluding");
544        assert!(dot_excluded.contains("n1"), "Node n1 (connected) should be present when excluding");
545        assert!(!dot_excluded.contains("n2"), "Node n2 (isolated) should NOT be present when excluding");
546        assert!(dot_excluded.contains("n0 -> n1"), "Edge n0 -> n1 should be present when excluding");
547
548        // Test 2: Include isolated nodes (default)
549        let config_include = DotExportConfig {
550            exclude_isolated_nodes: false, // Or use DotExportConfig::default()
551        };
552        let dot_included = graph.to_dot("TestIncludeIsolated", &config_include);
553
554        assert!(dot_included.contains("n0"), "Node n0 (connected) should be present when including");
555        assert!(dot_included.contains("n1"), "Node n1 (connected) should be present when including");
556        assert!(dot_included.contains("n2"), "Node n2 (isolated) should be present when including");
557        assert!(dot_included.contains("n0 -> n1"), "Edge n0 -> n1 should be present when including");
558    }
559}
560
traverse_graph/cg_dot.rs

traverse_graph/
cg_dot.rs
