use crate::actions::{Action, ActionType};
use std::collections::{HashMap, HashSet};
pub fn build_dag(actions: &[Action]) -> Vec<Vec<usize>> {
let n = actions.len();
if n == 0 {
return vec![];
}
let mut writers: HashMap<&str, usize> = HashMap::new();
for (i, action) in actions.iter().enumerate() {
if action.action_type == ActionType::StateWrite {
if let Some(key) = action.parameters.get("key").and_then(|v| v.as_str()) {
writers.insert(key, i);
}
}
for key in action.expected_effects.keys() {
writers.insert(key.as_str(), i);
}
}
let mut deps: Vec<HashSet<usize>> = vec![HashSet::new(); n];
for (i, action) in actions.iter().enumerate() {
for dep_key in &action.state_dependencies {
if let Some(&writer_idx) = writers.get(dep_key.as_str()) {
if writer_idx < i {
deps[i].insert(writer_idx);
}
}
}
if action.action_type == ActionType::StateWrite {
if let Some(key) = action.parameters.get("key").and_then(|v| v.as_str()) {
for j in 0..i {
if actions[j].action_type == ActionType::StateWrite {
if let Some(prev_key) =
actions[j].parameters.get("key").and_then(|v| v.as_str())
{
if prev_key == key {
deps[i].insert(j);
}
}
}
}
}
}
}
let mut completed: HashSet<usize> = HashSet::new();
let mut levels: Vec<Vec<usize>> = Vec::new();
let mut remaining: HashSet<usize> = (0..n).collect();
while !remaining.is_empty() {
let mut ready: Vec<usize> = remaining
.iter()
.filter(|&&i| deps[i].is_subset(&completed))
.copied()
.collect();
ready.sort();
if ready.is_empty() {
ready.push(*remaining.iter().min().unwrap());
}
for &i in &ready {
completed.insert(i);
remaining.remove(&i);
}
levels.push(ready);
}
levels
}
#[cfg(test)]
mod tests {
use super::*;
use crate::actions::{Action, FailureBehavior};
use serde_json::Value;
use std::collections::HashMap;
fn tool_call(id: &str, tool: &str) -> Action {
Action {
id: id.to_string(),
action_type: ActionType::ToolCall,
tool: Some(tool.to_string()),
parameters: HashMap::new(),
preconditions: vec![],
expected_effects: HashMap::new(),
state_dependencies: vec![],
idempotent: false,
max_retries: 3,
failure_behavior: FailureBehavior::Abort,
timeout_ms: None,
metadata: HashMap::new(),
}
}
fn state_write(id: &str, key: &str, value: i64) -> Action {
Action {
id: id.to_string(),
action_type: ActionType::StateWrite,
tool: None,
parameters: [
("key".to_string(), Value::from(key)),
("value".to_string(), Value::from(value)),
]
.into(),
preconditions: vec![],
expected_effects: HashMap::new(),
state_dependencies: vec![],
idempotent: false,
max_retries: 3,
failure_behavior: FailureBehavior::Abort,
timeout_ms: None,
metadata: HashMap::new(),
}
}
#[test]
fn empty_actions() {
assert_eq!(build_dag(&[]), Vec::<Vec<usize>>::new());
}
#[test]
fn no_deps_single_level() {
let actions = vec![tool_call("a", "t1"), tool_call("b", "t2"), tool_call("c", "t3")];
let levels = build_dag(&actions);
assert_eq!(levels.len(), 1);
assert_eq!(levels[0], vec![0, 1, 2]);
}
#[test]
fn diamond_deps() {
let mut a1 = tool_call("a1", "t1");
a1.state_dependencies = vec!["x".to_string()];
let mut a2 = tool_call("a2", "t2");
a2.state_dependencies = vec!["x".to_string()];
let actions = vec![state_write("w", "x", 1), a1, a2];
let levels = build_dag(&actions);
assert_eq!(levels[0], vec![0]); assert_eq!(levels[1], vec![1, 2]); }
#[test]
fn linear_state_writes() {
let actions = vec![state_write("a", "x", 1), state_write("b", "x", 2)];
let levels = build_dag(&actions);
assert_eq!(levels.len(), 2);
assert_eq!(levels[0], vec![0]);
assert_eq!(levels[1], vec![1]);
}
}