taskflowrs 0.1.1

A Rust implementation of TaskFlow — task-parallel programming with heterogeneous GPU support
Documentation
use std::sync::{Arc, Mutex};
use std::collections::HashSet;

#[cfg(feature = "async")]
use std::future::Future;
#[cfg(feature = "async")]
use std::pin::Pin;

pub type TaskId = usize;

/// Task work - either a static closure, subflow, condition, or async task
pub enum TaskWork {
    Static(Box<dyn FnOnce() + Send + 'static>),
    Subflow(Box<dyn FnOnce(&mut crate::Subflow) + Send + 'static>),
    Condition(Box<dyn FnOnce() -> usize + Send + 'static>),
    #[cfg(feature = "async")]
    Async(Box<dyn FnOnce() -> Pin<Box<dyn Future<Output = ()> + Send>> + Send + 'static>),
}

/// Internal task node structure
pub struct TaskNode {
    pub id: TaskId,
    pub name: String,
    pub work: Option<TaskWork>,
    pub successors: HashSet<TaskId>,
    pub dependents: HashSet<TaskId>,
    pub num_dependents: usize,
}

impl TaskNode {
    pub fn new(id: TaskId, work: TaskWork) -> Self {
        Self {
            id,
            name: format!("task_{}", id),
            work: Some(work),
            successors: HashSet::new(),
            dependents: HashSet::new(),
            num_dependents: 0,
        }
    }
}

/// Task handle for building task graphs
#[derive(Clone)]
pub struct TaskHandle {
    pub(crate) id: TaskId,
    graph: Arc<Mutex<Vec<TaskNode>>>,
}

impl TaskHandle {
    pub(crate) fn new(id: TaskId, graph: Arc<Mutex<Vec<TaskNode>>>) -> Self {
        Self { id, graph }
    }

    /// Set the task name
    pub fn name(self, name: &str) -> Self {
        {
            let mut graph = self.graph.lock().unwrap();
            if let Some(node) = graph.iter_mut().find(|n| n.id == self.id) {
                node.name = name.to_string();
            }
        } // Lock is dropped here
        self
    }

    /// Make this task precede another task (this -> other)
    pub fn precede(&self, other: &TaskHandle) {
        let mut graph = self.graph.lock().unwrap();
        
        // Add edge from self to other
        if let Some(node) = graph.iter_mut().find(|n| n.id == self.id) {
            node.successors.insert(other.id);
        }
        
        // Update other's dependents and increment num_dependents
        if let Some(node) = graph.iter_mut().find(|n| n.id == other.id) {
            if node.dependents.insert(self.id) {
                // Only increment if this is a new dependency (insert returns true if newly inserted)
                node.num_dependents += 1;
            }
        }
    }

    /// Make this task succeed another task (other -> this)
    pub fn succeed(&self, other: &TaskHandle) {
        other.precede(self);
    }
    
    /// Get the task ID
    pub fn id(&self) -> TaskId {
        self.id
    }

    #[allow(dead_code)]
    pub(crate) fn get_id(&self) -> TaskId {
        self.id
    }
}

/// Task builder trait for creating tasks with different types of work
pub trait Task {
    fn id(&self) -> TaskId;
}

impl Task for TaskHandle {
    fn id(&self) -> TaskId {
        self.id
    }
}