Behavior Tree Pattern

// Node execution results
#[derive(Debug, Clone, PartialEq)]
pub enum NodeResult {
    Success,
    Failure,
    Running,
}
 
// Context for behavior execution
pub struct BehaviorContext {
    pub agent_id: String,
    pub variables: std::collections::HashMap<String, f64>,
}
 
// Base behavior node trait
pub trait BehaviorNode {
    fn execute(&self, context: &mut BehaviorContext) -> NodeResult;
    fn reset(&self);
}
 
// Action node - performs actual work
pub struct ActionNode {
    name: String,
    action_fn: fn(&mut BehaviorContext) -> NodeResult,
}
 
impl ActionNode {
    pub fn new(name: &str, action_fn: fn(&mut BehaviorContext) -> NodeResult) -> Self {
        Self {
            name: name.to_string(),
            action_fn,
        }
    }
}
 
impl BehaviorNode for ActionNode {
    fn execute(&self, context: &mut BehaviorContext) -> NodeResult {
        (self.action_fn)(context)
    }
 
    fn reset(&self) {
        // Actions typically don't need reset
    }
}
 
// Condition node - checks state
pub struct ConditionNode {
    name: String,
    condition_fn: fn(&BehaviorContext) -> bool,
}
 
impl ConditionNode {
    pub fn new(name: &str, condition_fn: fn(&BehaviorContext) -> bool) -> Self {
        Self {
            name: name.to_string(),
            condition_fn,
        }
    }
}
 
impl BehaviorNode for ConditionNode {
    fn execute(&self, context: &mut BehaviorContext) -> NodeResult {
        if (self.condition_fn)(context) {
            NodeResult::Success
        } else {
            NodeResult::Failure
        }
    }
 
    fn reset(&self) {
        // Conditions typically don't need reset
    }
}
 
// Sequence node - executes children in order, fails if any fails
pub struct SequenceNode {
    children: Vec<Box<dyn BehaviorNode>>,
    current_child: usize,
}
 
impl SequenceNode {
    pub fn new(children: Vec<Box<dyn BehaviorNode>>) -> Self {
        Self {
            children,
            current_child: 0,
        }
    }
}
 
impl BehaviorNode for SequenceNode {
    fn execute(&self, context: &mut BehaviorContext) -> NodeResult {
        for child in &self.children {
            match child.execute(context) {
                NodeResult::Success => continue,
                NodeResult::Failure => return NodeResult::Failure,
                NodeResult::Running => return NodeResult::Running,
            }
        }
        NodeResult::Success
    }
 
    fn reset(&self) {
        for child in &self.children {
            child.reset();
        }
    }
}
 
// Selector node - tries children until one succeeds
pub struct SelectorNode {
    children: Vec<Box<dyn BehaviorNode>>,
}
 
impl SelectorNode {
    pub fn new(children: Vec<Box<dyn BehaviorNode>>) -> Self {
        Self { children }
    }
}
 
impl BehaviorNode for SelectorNode {
    fn execute(&self, context: &mut BehaviorContext) -> NodeResult {
        for child in &self.children {
            match child.execute(context) {
                NodeResult::Success => return NodeResult::Success,
                NodeResult::Failure => continue,
                NodeResult::Running => return NodeResult::Running,
            }
        }
        NodeResult::Failure
    }
 
    fn reset(&self) {
        for child in &self.children {
            child.reset();
        }
    }
}
 
// Behavior tree agent
pub struct BehaviorTreeAgent {
    id: String,
    root_node: Box<dyn BehaviorNode>,
    context: BehaviorContext,
}
 
impl BehaviorTreeAgent {
    pub fn new(id: &str, root_node: Box<dyn BehaviorNode>) -> Self {
        Self {
            id: id.to_string(),
            root_node,
            context: BehaviorContext {
                agent_id: id.to_string(),
                variables: std::collections::HashMap::new(),
            },
        }
    }
 
    pub fn tick(&mut self) -> NodeResult {
        self.root_node.execute(&mut self.context)
    }
 
    pub fn set_variable(&mut self, key: &str, value: f64) {
        self.context.variables.insert(key.to_string(), value);
    }
 
    pub fn get_variable(&self, key: &str) -> Option<&f64> {
        self.context.variables.get(key)
    }
}
 
// Example behavior functions
fn check_energy_level(context: &BehaviorContext) -> bool {
    context.variables.get("energy").unwrap_or(&0.0) > &50.0
}
 
fn recharge_action(context: &mut BehaviorContext) -> NodeResult {
    println!("Agent {} is recharging", context.agent_id);
    context.variables.insert("energy".to_string(), 100.0);
    NodeResult::Success
}
 
fn work_action(context: &mut BehaviorContext) -> NodeResult {
    let current_energy = context.variables.get("energy").unwrap_or(&0.0);
    if *current_energy > 10.0 {
        context.variables.insert("energy".to_string(), current_energy - 10.0);
        println!("Agent {} is working (energy: {})", context.agent_id, current_energy - 10.0);
        NodeResult::Success
    } else {
        NodeResult::Failure
    }
}

// Build behavior tree: (Check Energy -> Work) OR Recharge
let work_sequence = SequenceNode::new(vec![
    Box::new(ConditionNode::new("check_energy", check_energy_level)),
    Box::new(ActionNode::new("work", work_action)),
]);
 
let root_selector = SelectorNode::new(vec![
    Box::new(work_sequence),
    Box::new(ActionNode::new("recharge", recharge_action)),
]);
 
let mut agent = BehaviorTreeAgent::new("WorkerBot", Box::new(root_selector));
agent.set_variable("energy", 30.0);
 
// Execute behavior tree
for _tick in 0..5 {
    let result = agent.tick();
    println!("Tick result: {:?}", result);
}