use crate::types::NodeKind;
use petgraph::graph::{DiGraph, NodeIndex};
use rustc_hash::FxHashMap;
pub type WeaveDiGraph = DiGraph<NodeKind, ()>;
pub type NodeIndexMap = FxHashMap<NodeKind, NodeIndex>;
#[derive(Debug, Clone)]
pub struct PetgraphConversion {
pub graph: WeaveDiGraph,
pub index_map: NodeIndexMap,
}
impl PetgraphConversion {
#[must_use]
pub fn index_of(&self, node: &NodeKind) -> Option<NodeIndex> {
self.index_map.get(node).copied()
}
#[must_use]
pub fn node_at(&self, index: NodeIndex) -> Option<&NodeKind> {
self.graph.node_weight(index)
}
}
pub(super) fn to_petgraph(edges: &FxHashMap<NodeKind, Vec<NodeKind>>) -> PetgraphConversion {
let mut graph = DiGraph::new();
let mut index_map: NodeIndexMap = FxHashMap::default();
let mut all_nodes: Vec<NodeKind> = Vec::new();
for (from, tos) in edges {
if !index_map.contains_key(from) {
all_nodes.push(from.clone());
index_map.insert(from.clone(), NodeIndex::new(0)); }
for to in tos {
if !index_map.contains_key(to) {
all_nodes.push(to.clone());
index_map.insert(to.clone(), NodeIndex::new(0)); }
}
}
all_nodes.sort_by(|a, b| match (a, b) {
(NodeKind::Start, _) => std::cmp::Ordering::Less,
(_, NodeKind::Start) => std::cmp::Ordering::Greater,
(NodeKind::End, _) => std::cmp::Ordering::Greater,
(_, NodeKind::End) => std::cmp::Ordering::Less,
(NodeKind::Custom(a_name), NodeKind::Custom(b_name)) => a_name.cmp(b_name),
});
for node in &all_nodes {
let idx = graph.add_node(node.clone());
index_map.insert(node.clone(), idx);
}
for (from, tos) in edges {
let from_idx = index_map[from];
for to in tos {
let to_idx = index_map[to];
graph.add_edge(from_idx, to_idx, ());
}
}
PetgraphConversion { graph, index_map }
}
pub(super) fn to_dot(edges: &FxHashMap<NodeKind, Vec<NodeKind>>) -> String {
use std::fmt::Write;
let conversion = to_petgraph(edges);
let mut output = String::new();
writeln!(output, "digraph {{").unwrap();
writeln!(output, " rankdir=TB;").unwrap();
writeln!(output, " node [shape=box, style=rounded];").unwrap();
for idx in conversion.graph.node_indices() {
let node = conversion.graph.node_weight(idx).unwrap();
let (label, style) = match node {
NodeKind::Start => ("Start", " style=\"filled\" fillcolor=\"lightgreen\""),
NodeKind::End => ("End", " style=\"filled\" fillcolor=\"lightcoral\""),
NodeKind::Custom(name) => (name.as_str(), ""),
};
writeln!(
output,
" {} [ label=\"{}\"{} ];",
idx.index(),
label,
style
)
.unwrap();
}
writeln!(output).unwrap();
for edge in conversion.graph.edge_indices() {
let (from, to) = conversion.graph.edge_endpoints(edge).unwrap();
writeln!(output, " {} -> {};", from.index(), to.index()).unwrap();
}
writeln!(output, "}}").unwrap();
output
}
#[must_use]
pub fn is_cyclic(edges: &FxHashMap<NodeKind, Vec<NodeKind>>) -> bool {
let conversion = to_petgraph(edges);
petgraph::algo::is_cyclic_directed(&conversion.graph)
}
#[cfg(test)]
mod tests {
use super::*;
fn make_linear_graph() -> FxHashMap<NodeKind, Vec<NodeKind>> {
let mut edges = FxHashMap::default();
edges.insert(NodeKind::Start, vec![NodeKind::Custom("A".into())]);
edges.insert(NodeKind::Custom("A".into()), vec![NodeKind::End]);
edges
}
fn make_cyclic_graph() -> FxHashMap<NodeKind, Vec<NodeKind>> {
let mut edges = FxHashMap::default();
edges.insert(NodeKind::Start, vec![NodeKind::Custom("A".into())]);
edges.insert(
NodeKind::Custom("A".into()),
vec![NodeKind::Custom("B".into())],
);
edges.insert(
NodeKind::Custom("B".into()),
vec![NodeKind::Custom("A".into())], );
edges
}
#[test]
fn test_to_petgraph_linear() {
let edges = make_linear_graph();
let conversion = to_petgraph(&edges);
assert_eq!(conversion.graph.node_count(), 3);
assert_eq!(conversion.graph.edge_count(), 2);
assert!(conversion.index_of(&NodeKind::Start).is_some());
assert!(conversion.index_of(&NodeKind::Custom("A".into())).is_some());
assert!(conversion.index_of(&NodeKind::End).is_some());
}
#[test]
fn test_is_cyclic_linear() {
let edges = make_linear_graph();
assert!(!is_cyclic(&edges));
}
#[test]
fn test_is_cyclic_with_cycle() {
let edges = make_cyclic_graph();
assert!(is_cyclic(&edges));
}
#[test]
fn test_to_dot_output() {
let edges = make_linear_graph();
let dot = to_dot(&edges);
assert!(dot.contains("digraph {"));
assert!(dot.contains("Start"));
assert!(dot.contains("End"));
assert!(dot.contains("->"));
}
#[test]
fn test_deterministic_indices() {
let edges = make_linear_graph();
let conv1 = to_petgraph(&edges);
let conv2 = to_petgraph(&edges);
assert_eq!(
conv1.index_of(&NodeKind::Start),
conv2.index_of(&NodeKind::Start)
);
assert_eq!(
conv1.index_of(&NodeKind::Custom("A".into())),
conv2.index_of(&NodeKind::Custom("A".into()))
);
}
}