#[derive(Debug, Clone)]
pub struct TaskNode {
pub id: u64,
pub name: String,
pub estimated_ms: u64,
pub resource_requirements: Vec<String>,
}
impl TaskNode {
#[must_use]
pub fn new(id: u64, name: impl Into<String>, estimated_ms: u64) -> Self {
Self {
id,
name: name.into(),
estimated_ms,
resource_requirements: Vec::new(),
}
}
#[must_use]
pub fn with_resource(mut self, resource: impl Into<String>) -> Self {
self.resource_requirements.push(resource.into());
self
}
#[must_use]
pub fn total_resource_count(&self) -> usize {
self.resource_requirements.len()
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum EdgeType {
Sequential,
Parallel,
ConditionalSuccess,
ConditionalFailure,
}
impl EdgeType {
#[must_use]
pub const fn is_conditional(self) -> bool {
matches!(self, Self::ConditionalSuccess | Self::ConditionalFailure)
}
}
#[derive(Debug, Clone)]
pub struct TaskEdge {
pub from_id: u64,
pub to_id: u64,
pub edge_type: EdgeType,
}
impl TaskEdge {
#[must_use]
pub fn new(from_id: u64, to_id: u64, edge_type: EdgeType) -> Self {
Self {
from_id,
to_id,
edge_type,
}
}
}
#[derive(Debug, Clone, Default)]
pub struct TaskGraph {
pub nodes: Vec<TaskNode>,
pub edges: Vec<TaskEdge>,
}
impl TaskGraph {
#[must_use]
pub fn new() -> Self {
Self::default()
}
pub fn add_node(&mut self, node: TaskNode) {
self.nodes.push(node);
}
pub fn add_edge(&mut self, edge: TaskEdge) {
self.edges.push(edge);
}
#[must_use]
pub fn successors(&self, id: u64) -> Vec<u64> {
self.edges
.iter()
.filter(|e| e.from_id == id)
.map(|e| e.to_id)
.collect()
}
#[must_use]
pub fn predecessors(&self, id: u64) -> Vec<u64> {
self.edges
.iter()
.filter(|e| e.to_id == id)
.map(|e| e.from_id)
.collect()
}
#[must_use]
pub fn root_nodes(&self) -> Vec<u64> {
self.nodes
.iter()
.filter(|n| self.predecessors(n.id).is_empty())
.map(|n| n.id)
.collect()
}
#[must_use]
pub fn critical_path_ms(&self) -> u64 {
if self.nodes.is_empty() {
return 0;
}
let mut best: std::collections::HashMap<u64, u64> =
self.nodes.iter().map(|n| (n.id, 0u64)).collect();
let mut in_degree: std::collections::HashMap<u64, usize> =
self.nodes.iter().map(|n| (n.id, 0)).collect();
for e in &self.edges {
*in_degree.entry(e.to_id).or_insert(0) += 1;
}
let mut queue: std::collections::VecDeque<u64> = in_degree
.iter()
.filter(|(_, °)| deg == 0)
.map(|(&id, _)| id)
.collect();
for &id in &queue {
if let Some(node) = self.nodes.iter().find(|n| n.id == id) {
best.insert(id, node.estimated_ms);
}
}
while let Some(node_id) = queue.pop_front() {
let current_best = best[&node_id];
let node_cost = self
.nodes
.iter()
.find(|n| n.id == node_id)
.map_or(0, |n| n.estimated_ms);
for succ_id in self.successors(node_id) {
let candidate = current_best
+ node_cost.saturating_sub(
node_cost, );
let succ_cost = self
.nodes
.iter()
.find(|n| n.id == succ_id)
.map_or(0, |n| n.estimated_ms);
let new_val = current_best + succ_cost;
let entry = best.entry(succ_id).or_insert(0);
if new_val > *entry {
*entry = new_val;
}
let _ = candidate;
if let Some(deg) = in_degree.get_mut(&succ_id) {
if *deg > 0 {
*deg -= 1;
}
if *deg == 0 {
queue.push_back(succ_id);
}
}
}
}
best.values().copied().max().unwrap_or(0)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_task_node_resource_count_empty() {
let n = TaskNode::new(1, "encode", 5000);
assert_eq!(n.total_resource_count(), 0);
}
#[test]
fn test_task_node_resource_count_with_resources() {
let n = TaskNode::new(1, "encode", 5000)
.with_resource("gpu")
.with_resource("network");
assert_eq!(n.total_resource_count(), 2);
}
#[test]
fn test_edge_type_is_conditional_false() {
assert!(!EdgeType::Sequential.is_conditional());
assert!(!EdgeType::Parallel.is_conditional());
}
#[test]
fn test_edge_type_is_conditional_true() {
assert!(EdgeType::ConditionalSuccess.is_conditional());
assert!(EdgeType::ConditionalFailure.is_conditional());
}
#[test]
fn test_add_node_and_edge() {
let mut g = TaskGraph::new();
g.add_node(TaskNode::new(1, "a", 100));
g.add_node(TaskNode::new(2, "b", 200));
g.add_edge(TaskEdge::new(1, 2, EdgeType::Sequential));
assert_eq!(g.nodes.len(), 2);
assert_eq!(g.edges.len(), 1);
}
#[test]
fn test_successors() {
let mut g = TaskGraph::new();
g.add_node(TaskNode::new(1, "a", 100));
g.add_node(TaskNode::new(2, "b", 100));
g.add_node(TaskNode::new(3, "c", 100));
g.add_edge(TaskEdge::new(1, 2, EdgeType::Sequential));
g.add_edge(TaskEdge::new(1, 3, EdgeType::Parallel));
let mut succs = g.successors(1);
succs.sort_unstable();
assert_eq!(succs, vec![2, 3]);
}
#[test]
fn test_predecessors() {
let mut g = TaskGraph::new();
g.add_node(TaskNode::new(1, "a", 100));
g.add_node(TaskNode::new(2, "b", 100));
g.add_node(TaskNode::new(3, "c", 100));
g.add_edge(TaskEdge::new(1, 3, EdgeType::Sequential));
g.add_edge(TaskEdge::new(2, 3, EdgeType::Sequential));
let mut preds = g.predecessors(3);
preds.sort_unstable();
assert_eq!(preds, vec![1, 2]);
}
#[test]
fn test_root_nodes_single() {
let mut g = TaskGraph::new();
g.add_node(TaskNode::new(1, "root", 0));
g.add_node(TaskNode::new(2, "leaf", 0));
g.add_edge(TaskEdge::new(1, 2, EdgeType::Sequential));
assert_eq!(g.root_nodes(), vec![1]);
}
#[test]
fn test_root_nodes_multiple() {
let mut g = TaskGraph::new();
g.add_node(TaskNode::new(1, "a", 0));
g.add_node(TaskNode::new(2, "b", 0));
g.add_node(TaskNode::new(3, "c", 0));
g.add_edge(TaskEdge::new(1, 3, EdgeType::Sequential));
g.add_edge(TaskEdge::new(2, 3, EdgeType::Sequential));
let mut roots = g.root_nodes();
roots.sort_unstable();
assert_eq!(roots, vec![1, 2]);
}
#[test]
fn test_successors_empty() {
let g = TaskGraph::new();
assert!(g.successors(99).is_empty());
}
#[test]
fn test_predecessors_empty() {
let g = TaskGraph::new();
assert!(g.predecessors(99).is_empty());
}
#[test]
fn test_critical_path_empty_graph() {
let g = TaskGraph::new();
assert_eq!(g.critical_path_ms(), 0);
}
#[test]
fn test_critical_path_single_node() {
let mut g = TaskGraph::new();
g.add_node(TaskNode::new(1, "only", 500));
assert_eq!(g.critical_path_ms(), 500);
}
#[test]
fn test_critical_path_linear_chain() {
let mut g = TaskGraph::new();
g.add_node(TaskNode::new(1, "a", 100));
g.add_node(TaskNode::new(2, "b", 200));
g.add_node(TaskNode::new(3, "c", 300));
g.add_edge(TaskEdge::new(1, 2, EdgeType::Sequential));
g.add_edge(TaskEdge::new(2, 3, EdgeType::Sequential));
assert_eq!(g.critical_path_ms(), 600);
}
#[test]
fn test_critical_path_parallel_paths() {
let mut g = TaskGraph::new();
g.add_node(TaskNode::new(1, "a", 100));
g.add_node(TaskNode::new(2, "b", 200));
g.add_node(TaskNode::new(3, "c", 500));
g.add_edge(TaskEdge::new(1, 2, EdgeType::Sequential));
g.add_edge(TaskEdge::new(3, 2, EdgeType::Sequential));
assert_eq!(g.critical_path_ms(), 700);
}
}