lyra-engine/lyra-ecs/src/system/graph.rs

use std::{collections::{HashMap, VecDeque, HashSet}, ptr::NonNull};

use super::System;

use crate::{World, CommandQueue, Commands};

#[derive(thiserror::Error, Debug)]
pub enum GraphExecutorError {
    #[error("could not find a system's dependency named `{0}`")]
    MissingSystem(String),
    #[error("system `{0}` returned with an error: `{1}`")]
    SystemError(String, anyhow::Error),
    #[error("a command returned with an error: `{0}`")]
    Command(anyhow::Error)
}

/// A single system in the graph.
/// 
/// This is a helper struct for the [`GraphExecutor`] that stores the name
/// and dependencies of a system.
pub struct GraphSystem {
    /// The name of this system
    name: String,
    /// The actual system
    system: Box<dyn System>,
    /// The dependencies of this system
    depends: Vec<String>,
}

/// A system executor that represents the systems in a graph to handle dependencies.
/// 
/// The graph uses an adjacency list.
#[derive(Default)]
pub struct GraphExecutor {
    systems: HashMap<String, GraphSystem>,
}

impl GraphExecutor {
    pub fn new() -> Self {
        Self::default()
    }

    /// Inserts a system into the Graph.
    /// 
    /// This does not validate that the dependencies of this system are existing. This is done on
    /// execution. It makes it easier when adding systems so you don't have to make sure you
    /// insert everything in the correct order.
    pub fn insert_system<S>(&mut self, name: &str, system: S, depends: &[&str])
    where
        S: System + 'static
    {
        let system = GraphSystem {
            name: name.to_string(),
            system: Box::new(system),
            depends: depends.iter().map(|d| d.to_string()).collect(),
        };
        self.systems.insert(name.to_string(), system);
    }

    /// Executes the systems in the graph
    pub fn execute(&mut self, mut world_ptr: NonNull<World>, stop_on_error: bool)
            -> Result<Vec<GraphExecutorError>, GraphExecutorError> {
        let mut stack = VecDeque::new();
        let mut visited = HashSet::new();

        for (_name, node) in self.systems.iter() {
            self.topological_sort(&mut stack, &mut visited, node)?;
        }

        let mut possible_errors = Vec::new();

        while let Some(node) = stack.pop_front() {
            let system = self.systems.get_mut(node.as_str()).unwrap();

            if let Err(e) = system.system.execute(world_ptr)
                    .map_err(|e| GraphExecutorError::SystemError(node, e)) {
                if stop_on_error {
                    return Err(e);
                }

                possible_errors.push(e);
                unimplemented!("Cannot resume topological execution from error"); // TODO: resume topological execution from error
            }

            if let Err(e) = system.system.execute_deferred(world_ptr)
                    .map_err(|e| GraphExecutorError::Command(e)) {

                if stop_on_error {
                    return Err(e);
                }

                possible_errors.push(e);
                unimplemented!("Cannot resume topological execution from error"); // TODO: resume topological execution from error
            }

            let world = unsafe { world_ptr.as_mut() };
            if let Some(mut queue) = world.try_get_resource_mut::<CommandQueue>() {
                // Safety: Commands only borrows world.entities when adding commands
                let world = unsafe { world_ptr.as_mut() };
                let mut commands = Commands::new(&mut queue, world);

                let world = unsafe { world_ptr.as_mut() };
                commands.execute(world);
            }
        }

        Ok(possible_errors)
    }

    fn topological_sort<'a>(&'a self, stack: &mut VecDeque<String>,
            visited: &mut HashSet<&'a str>, node: &'a GraphSystem) -> Result<(), GraphExecutorError> {
        
        if !visited.contains(node.name.as_str()) {
            visited.insert(&node.name);

            for depend in node.depends.iter() {
                let node = self.systems.get(depend)
                    .ok_or_else(|| GraphExecutorError::MissingSystem(depend.clone()))?;
    
                if !visited.contains(depend.as_str()) {
                    self.topological_sort(stack, visited, node)?;
                }
            }
    
            stack.push_back(node.name.clone());
        }

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use std::ptr::NonNull;

    use crate::{query::{ResMut, View}, system::IntoSystem, World};

    use super::GraphExecutor;

    #[test]
    pub fn execution() {
        let mut world = World::new();
        let order_list = Vec::<String>::new();
        world.add_resource(order_list);

        let mut exec = GraphExecutor::new();

        let a_system = |view: View<ResMut<Vec<String>>, ()>| -> anyhow::Result<()> {
            println!("System 'a' ran!");
            
            let mut order = view.into_iter().next().unwrap();
            order.push("a".to_string());

            Ok(())
        };

        let b_system = |view: View<ResMut<Vec<String>>, ()>| -> anyhow::Result<()> {
            println!("System 'b' ran!");

            let mut order = view.into_iter().next().unwrap();
            order.push("b".to_string());

            Ok(())
        };

        let c_system = |view: View<ResMut<Vec<String>>, ()>| -> anyhow::Result<()> {
            println!("System 'c' ran!");

            let mut order = view.into_iter().next().unwrap();
            order.push("c".to_string());

            Ok(())
        };

        // inserted into the graph out of order...
        exec.insert_system("c", c_system.into_system(), &["b"]);
        exec.insert_system("a", a_system.into_system(), &[]);
        exec.insert_system("b", b_system.into_system(), &["a"]);

        exec.execute(NonNull::from(&world), true).unwrap();
        println!("Executed systems");

        let order = world.get_resource::<Vec<String>>();
        let mut order_iter = order.iter();

        // ... but still executed in order
        assert_eq!(order_iter.next().unwrap().clone(), "a".to_string());
        assert_eq!(order_iter.next().unwrap().clone(), "b".to_string());
        assert_eq!(order_iter.next().unwrap().clone(), "c".to_string());
        println!("Systems executed in order");
    }
}