use std::collections::HashMap;
use petgraph::algo::toposort;
use petgraph::dot::{Config, Dot};
use petgraph::graph::{EdgeIndex, NodeIndex};
use petgraph::stable_graph::StableGraph;
use petgraph::visit::{Control, depth_first_search, DfsEvent, Visitable, VisitMap};
use crate::Task;
#[derive(Clone)]
pub struct DependencyGraph<S> {
inner: StableGraph<String, ()>,
nodes: HashMap<String, NodeIndex>,
tasks: HashMap<String, Task<S>>,
}
impl<S> DependencyGraph<S> where S: Clone {
pub fn new() -> Self {
Self { inner: Default::default(), nodes: Default::default(), tasks: Default::default() }
}
pub fn add_task(&mut self, name: &str, task: Task<S>) {
self.find_or_add_node(name.to_string());
self.tasks.insert(name.to_string(), task);
}
pub fn add_dependency(&mut self, name: &str, dep: &str) {
self.find_or_add_edge(name.to_string(), dep.to_string());
}
pub fn pop_next(&mut self) -> Option<(String, Task<S>)> {
if self.is_empty() { return None; }
let weight = self.toposort().first().unwrap().clone();
let node = self.find_or_add_node(weight.clone());
let name = self.inner.remove_node(node)?;
Some((name.clone(), (*self.tasks.get(name.as_str())?).clone()))
}
pub fn is_empty(&self) -> bool {
self.inner.node_count() == 0
}
pub fn retain_dependencies(&mut self, weights: Vec<String>) {
let mut vm = self.inner.visit_map();
for w in weights {
let node = self.find_or_add_node(w);
depth_first_search(&self.inner, Some(node), |event| {
if let DfsEvent::Discover(n, _time) = event {
if vm.is_visited(&n) { return Control::<NodeIndex>::Prune; }
vm.visit(n);
}
Control::<NodeIndex>::Continue
});
}
self.inner.retain_nodes(|_fg, n| vm.is_visited(&n));
}
pub fn toposort(&self) -> Vec<String> {
toposort(&self.inner, None).unwrap().iter().map(|n| {
self.inner.node_weight(*n).unwrap().clone()
}).rev().collect()
}
fn find_or_add_node(&mut self, weight: String) -> NodeIndex {
let entry = self.nodes.entry(weight.clone());
*entry.or_insert_with(|| self.inner.add_node(weight))
}
fn find_or_add_edge(&mut self, weight0: String, weight1: String) -> EdgeIndex {
let node0 = self.find_or_add_node(weight0);
let node1 = self.find_or_add_node(weight1);
self.inner.find_edge(node0, node1).unwrap_or_else(|| self.inner.add_edge(node0, node1, ()))
}
}
impl<S> ToString for DependencyGraph<S> {
fn to_string(&self) -> String {
format!("{:?}", Dot::with_config(&self.inner, &[Config::EdgeNoLabel]))
}
}