python_to_mermaid/convert/
flowchart_to_mermaid.rs

1use itertools::{chain, Itertools as _};
2
3use crate::{
4    flowchart::{Flowchart, FlowchartItem},
5    mermaid::{MermaidFlowchart, MermaidGraph, NodeId, NodeShape},
6};
7
8#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
9pub struct PreviousNode {
10    pub nid: NodeId,
11    pub label: Option<String>,
12    pub is_terminal: bool,
13}
14
15pub fn convert(fc: &Flowchart) -> MermaidFlowchart {
16    let mut mfc = MermaidFlowchart::new();
17    let mg = mfc.graph_mut();
18    add_flowchart_to(mg, vec![], fc);
19    mfc
20}
21
22pub fn add_flowchart_to(
23    mg: &mut MermaidGraph,
24    prev_nodes: Vec<PreviousNode>,
25    fc: &Flowchart,
26) -> Vec<PreviousNode> {
27    let (begin_shape, end_shape) = if fc.is_root {
28        (NodeShape::Rounded, NodeShape::Rounded)
29    } else {
30        (NodeShape::Trapezoid, NodeShape::InvertedTrapezoid)
31    };
32
33    let nid_begin = mg.add_node(&fc.begin, begin_shape);
34    let nid_end = mg.add_node(&fc.end, end_shape);
35
36    add_edges_to(mg, &prev_nodes, &nid_begin);
37
38    let mut prev_nodes = vec![PreviousNode {
39        nid: nid_begin.clone(),
40        label: None,
41        is_terminal: false,
42    }];
43    for item in &fc.items {
44        prev_nodes = add_item_to(mg, nid_begin.clone(), nid_end.clone(), prev_nodes, item);
45    }
46
47    add_edges_to_end(mg, &prev_nodes, &nid_end);
48
49    vec![PreviousNode {
50        nid: nid_end,
51        label: None,
52        is_terminal: false,
53    }]
54}
55
56pub fn add_item_to(
57    mg: &mut MermaidGraph,
58    nid_begin: NodeId,
59    nid_end: NodeId,
60    prev_nodes: Vec<PreviousNode>,
61    item: &FlowchartItem,
62) -> Vec<PreviousNode> {
63    match item {
64        FlowchartItem::Step(step) => {
65            let nid = mg.add_node(&step.label, NodeShape::Rectangle);
66            add_edges_to(mg, &prev_nodes, &nid);
67
68            vec![PreviousNode {
69                nid,
70                label: None,
71                is_terminal: false,
72            }]
73        }
74        FlowchartItem::Condition(cond) => {
75            let nid_cond = mg.add_node(&cond.condition, NodeShape::Diamond);
76            add_edges_to(mg, &prev_nodes, &nid_cond);
77
78            let mut then_nodes = vec![PreviousNode {
79                nid: nid_cond.clone(),
80                label: Some("T".into()),
81                is_terminal: false,
82            }];
83            for then_item in &cond.then_items {
84                then_nodes = add_item_to(
85                    mg,
86                    nid_begin.clone(),
87                    nid_end.clone(),
88                    then_nodes,
89                    then_item,
90                );
91            }
92
93            let mut else_nodes = vec![PreviousNode {
94                nid: nid_cond.clone(),
95                label: Some("F".into()),
96                is_terminal: false,
97            }];
98            for else_item in &cond.else_items {
99                else_nodes = add_item_to(
100                    mg,
101                    nid_begin.clone(),
102                    nid_end.clone(),
103                    else_nodes,
104                    else_item,
105                );
106            }
107
108            chain!(then_nodes.clone(), else_nodes.clone()).collect_vec()
109        }
110        FlowchartItem::Continue(continue_) => {
111            let nid = mg.add_node(&continue_.label, NodeShape::Rectangle);
112            add_edges_to(mg, &prev_nodes, &nid);
113            add_edges_to(
114                mg,
115                &[PreviousNode {
116                    nid: nid.clone(),
117                    label: None,
118                    is_terminal: false,
119                }],
120                &nid_begin,
121            );
122
123            vec![PreviousNode {
124                nid,
125                label: None,
126                is_terminal: true,
127            }]
128        }
129        FlowchartItem::Break(break_) => {
130            let nid = mg.add_node(&break_.label, NodeShape::Rectangle);
131            add_edges_to(mg, &prev_nodes, &nid);
132            add_edges_to(
133                mg,
134                &[PreviousNode {
135                    nid,
136                    label: None,
137                    is_terminal: false,
138                }],
139                &nid_end,
140            );
141
142            vec![]
143        }
144        FlowchartItem::Terminal(terminal) => {
145            let nid = mg.add_node(&terminal.label, NodeShape::Rounded);
146            add_edges_to(mg, &prev_nodes, &nid);
147
148            vec![PreviousNode {
149                nid,
150                label: None,
151                is_terminal: true,
152            }]
153        }
154        FlowchartItem::SubFlowchart(sub_fc) => {
155            let mg = mg.add_subgraph(sub_fc.begin.clone());
156            add_flowchart_to(mg, prev_nodes.clone(), sub_fc)
157        }
158    }
159}
160
161pub fn add_edges_to(mg: &mut MermaidGraph, prev_nodes: &[PreviousNode], next_node: &NodeId) {
162    add_edges_impl(mg, prev_nodes, next_node, false)
163}
164
165pub fn add_edges_to_end(mg: &mut MermaidGraph, prev_nodes: &[PreviousNode], end_node: &NodeId) {
166    add_edges_impl(mg, prev_nodes, end_node, true)
167}
168
169pub fn add_edges_impl(
170    mg: &mut MermaidGraph,
171    prev_nodes: &[PreviousNode],
172    next_node: &NodeId,
173    is_end: bool,
174) {
175    for prev_node in prev_nodes {
176        if prev_node.is_terminal && !is_end {
177            continue;
178        }
179
180        mg.add_edge(&prev_node.nid, next_node, prev_node.label.as_deref());
181    }
182}
python_to_mermaid/convert/flowchart_to_mermaid.rs

python_to_mermaid/convert/
flowchart_to_mermaid.rs
