Skip to main content

radiate_gp/collections/
dot.rs

1use crate::{Direction, Graph, Node, NodeType, Tree, TreeNode};
2use std::fmt::{Debug, Display};
3
4pub trait ToDot {
5    fn to_dot(&self) -> String;
6}
7
8impl<T> ToDot for Graph<T>
9where
10    T: Debug,
11{
12    fn to_dot(&self) -> String {
13        let mut dot = String::new();
14
15        dot += "digraph G {\n";
16
17        for (i, node) in self.iter().enumerate() {
18            let node_label = match node.node_type() {
19                NodeType::Input => format!(
20                    "{} [label=\"{:?}\", shape=box, color=blue];",
21                    i,
22                    node.value()
23                ),
24                NodeType::Output => {
25                    format!(
26                        "{} [label=\"{:?}\", shape=box, color=green];",
27                        i,
28                        node.value()
29                    )
30                }
31                NodeType::Vertex => {
32                    format!(
33                        "{} [label=\"{:?}\", shape=circle, color=orange];",
34                        i,
35                        node.value()
36                    )
37                }
38                NodeType::Edge => {
39                    format!(
40                        "{} [label=\"{:?}\", shape=diamond, color=gray];",
41                        i,
42                        node.value()
43                    )
44                }
45                _ => continue,
46            };
47            dot += &node_label;
48            dot += "\n";
49        }
50
51        for (i, node) in self.iter().enumerate() {
52            for incoming in node.incoming() {
53                let edge_style = if self[*incoming].direction() == Direction::Forward {
54                    " [style=solid]"
55                } else {
56                    " [style=dashed]"
57                };
58
59                dot += &format!("  {} -> {}{};\n", incoming, i, edge_style);
60            }
61        }
62
63        dot += "}\n";
64
65        dot
66    }
67}
68
69impl<T> ToDot for Tree<T>
70where
71    T: Display,
72{
73    fn to_dot(&self) -> String {
74        let mut dot = String::new();
75        dot.push_str("digraph G {\n");
76
77        dot_recursive(self.root().unwrap(), 0, &mut dot);
78
79        dot.push_str("}\n");
80        dot
81    }
82}
83
84fn dot_recursive<T: Display>(node: &TreeNode<T>, id: usize, dot: &mut String) -> usize {
85    dot.push_str(&format!("  {} [label=\"{}\"];\n", id, node.value()));
86
87    let mut next_id = id + 1;
88    if let Some(children) = node.children() {
89        for child in children.iter() {
90            let child_id = next_id;
91            next_id = dot_recursive(child, child_id, dot);
92            dot.push_str(&format!("  {} -> {};\n", id, child_id));
93        }
94    }
95    next_id
96}