Code wars - behaviour tree

code-wars.rs

// Petr Šťastný <[email protected]>
// 2019-07-14

// trait = interface
pub trait CodingBlock {
    fn execute(&self) -> Option<&Action>;
}

// Define enums. Each enum variant can 
// have another type inside it.
#[derive(Debug)] // Allow debug printing this enum
#[derive(PartialEq)] // Allow comparing this enum
pub enum Action {
    Move(Direction),
    Shoot(RobotSelector),
}

#[derive(Debug)]
#[derive(PartialEq)]
pub enum Direction {
    Towards(RobotSelector),
    AwayFrom(RobotSelector),
}

#[derive(Debug)]
#[derive(PartialEq)]
pub enum RobotSelector {
    Closest(RobotRelationship),
    FarthestAway(RobotRelationship),
    LeastHp(RobotRelationship),
    MostHp(RobotRelationship),
    ThisRobot,
}

#[derive(Debug)]
#[derive(PartialEq)]
pub enum RobotRelationship {
    Ally,
    Enemy,
}

// Define query enums
#[derive(Debug)]
#[derive(PartialEq)]
pub enum Query {
    LowerThan(i32, Property),
    HigherThan(i32, Property),
}

#[derive(Debug)]
#[derive(PartialEq)]
pub enum Property {
    Health(RobotSelector),
    Distance(RobotSelector)
}


// Define Action coding block
// This coding block returns an action
// to be done.
pub struct ActionCodingBlock<'a> {
    action: &'a Action
}
impl CodingBlock for ActionCodingBlock<'_> {
    fn execute(&self) -> Option<&Action> {
        Some(self.action)
    }
}
impl ActionCodingBlock<'_> {
    fn new(action: &Action) -> ActionCodingBlock {
        ActionCodingBlock {
            action
        }
    }
}
// Define query coding block.
pub struct QueryCodingBlock<'a> {
    query: &'a Query,
    child: Option<&'a CodingBlock>, // QueryCodingBlock has one child, that gets executed if query is true
}
impl CodingBlock for QueryCodingBlock<'_> {
    fn execute(&self) -> Option<&Action> {
        if let Some(child) = self.child { // if child exists
            // TODO: Test query and MAYBE execute child
            return child.execute();
        }else{
            // We can't return any action
            return None;
        }
    }
}
impl<'a> QueryCodingBlock<'a> {
    fn new(query: &'a Query) -> QueryCodingBlock<'a> {
        QueryCodingBlock {
            query,
            child: None
        }
    }
    fn register_child(&mut self, child: &'a CodingBlock) {
        self.child = Some(child)
    }
}
// Define hub. Hub connects to multiple coding blocks.
// It tries to execute them from left to right,
// and the first coding block to return an action
// is actually executed.
pub struct HubCodingBlock<'a> {
    children: Vec<&'a CodingBlock>
}
impl CodingBlock for HubCodingBlock<'_> {
    fn execute(&self) -> Option<&Action> {
        for child in &self.children {
            if let Some(action) = child.execute() {
                return Some(action);
            }
        }
        return None;
    }
}
impl<'a> HubCodingBlock<'a> {
    fn new() -> HubCodingBlock<'a> {
        HubCodingBlock {
            children: Vec::new()
        }
    }
    fn add_child(&mut self, child: &'a CodingBlock) {
        self.children.push(child);
    }
}


fn main() {
    // Construct the coding blocks
    let move_to_closest_enemy = Action::Move(Direction::Towards(RobotSelector::Closest(RobotRelationship::Enemy)));
    let shoot_weakest_enemy = Action::Shoot(RobotSelector::LeastHp(RobotRelationship::Enemy));
    let if_hp_below_twenty = Query::LowerThan(20, Property::Health(RobotSelector::ThisRobot));
    let run_away_from_closest_enemy = Action::Move(Direction::AwayFrom(RobotSelector::Closest(RobotRelationship::Enemy)));
    let if_closest_enemy_50m_or_closer = Query::LowerThan(50, Property::Distance(RobotSelector::Closest(RobotRelationship::Enemy)));

    // Construct decision tree
    let mut root = HubCodingBlock::new();
    let mut query_hp_below_twenty = QueryCodingBlock::new(&if_hp_below_twenty);
    let mut query_enemy_close_enough = QueryCodingBlock::new(&if_closest_enemy_50m_or_closer);
    let action_run_away_from_enemy = ActionCodingBlock::new(&run_away_from_closest_enemy);
    let action_shoot_enemy = ActionCodingBlock::new(&shoot_weakest_enemy);
    let action_go_towards_enemy = ActionCodingBlock::new(&move_to_closest_enemy);

    // Nodes in decision tree will be executed from left to right (just like going thr. binary tree)
    // First action that is found is executed. Actions below queries don't get executed if query fails.

    // Because of how Rust works with memory, register queries first
    // Root node will contain immutable reference to each node. We need guarantee that
    // nodes to which Root node has references will never change. So we must give queries reference to children
    // before connecting them to root node.
    query_hp_below_twenty.register_child(&action_run_away_from_enemy); // What happens when query is successful?
    query_enemy_close_enough.register_child(&action_shoot_enemy);
    // Now we can start connecting queries (which will never change at this point) to root node
    root.add_child(&query_hp_below_twenty); // hp_below_twenty is most important, so it goes first
    root.add_child(&query_enemy_close_enough); // if we have enough HP, check if there is an enemy close enough
    root.add_child(&action_go_towards_enemy); // If we didn't run away and we didn't shoot, go towards closest enemy.
    
    // Now we can execute the tree. If there was runtime, we could execute the action. Since
    // queries are not evaluated and always return true, we should get &run_away_from_closest_enemy,
    // since it's action with highest priority.

    // Test if this decision tree returns an action (.unwrap will cause the assert to fail if no value is returned)
    // and if the action is run_away_from_closest_enemy.
    assert_eq!(root.execute().unwrap(), &run_away_from_closest_enemy);
}