use serde::de::DeserializeOwned;
use super::policy::ParallelPolicy;
use super::stage::{ExecutableStage, Stage};
pub type StageResolver<S> = Box<dyn Fn(&S) -> Vec<Box<dyn ExecutableStage<S>>> + Send + Sync>;
pub enum WorkflowStep<S: Send + Sync> {
Stage(Box<dyn ExecutableStage<S>>),
Parallel {
stages: Vec<Box<dyn ExecutableStage<S>>>,
policy: ParallelPolicy,
},
DynamicParallel {
planner: Box<dyn ExecutableStage<S>>,
resolver: StageResolver<S>,
policy: ParallelPolicy,
},
Branch {
condition: Box<dyn Fn(&S) -> bool + Send + Sync>,
then_flow: Workflow<S>,
else_flow: Option<Workflow<S>>,
},
EarlyExitIf {
condition: Box<dyn Fn(&S) -> bool + Send + Sync>,
reason: &'static str,
},
}
pub struct Workflow<S: Send + Sync> {
pub name: &'static str,
pub steps: Vec<WorkflowStep<S>>,
}
impl<S: Send + Sync + 'static> Workflow<S> {
pub fn builder(name: &'static str) -> WorkflowBuilder<S> {
WorkflowBuilder::new(name)
}
}
pub struct WorkflowBuilder<S: Send + Sync> {
name: &'static str,
steps: Vec<WorkflowStep<S>>,
}
impl<S: Send + Sync + 'static> WorkflowBuilder<S> {
pub fn new(name: &'static str) -> Self {
Self {
name,
steps: Vec::new(),
}
}
pub fn stage<T: DeserializeOwned + Send + 'static>(mut self, stage: Stage<S, T>) -> Self
where
S: Send + Sync,
{
self.steps.push(WorkflowStep::Stage(Box::new(stage)));
self
}
pub fn executable_stage(mut self, stage: Box<dyn ExecutableStage<S>>) -> Self {
self.steps.push(WorkflowStep::Stage(stage));
self
}
pub fn parallel(
mut self,
stages: Vec<Box<dyn ExecutableStage<S>>>,
policy: ParallelPolicy,
) -> Self {
self.steps.push(WorkflowStep::Parallel { stages, policy });
self
}
pub fn dynamic_parallel<P, R>(
mut self,
planner: Stage<S, P>,
resolver: R,
policy: ParallelPolicy,
) -> Self
where
S: Send + Sync,
P: DeserializeOwned + Send + 'static,
R: Fn(&S) -> Vec<Box<dyn ExecutableStage<S>>> + Send + Sync + 'static,
{
self.steps.push(WorkflowStep::DynamicParallel {
planner: Box::new(planner),
resolver: Box::new(resolver),
policy,
});
self
}
pub fn early_exit_if<F>(mut self, condition: F, reason: &'static str) -> Self
where
F: Fn(&S) -> bool + Send + Sync + 'static,
{
self.steps.push(WorkflowStep::EarlyExitIf {
condition: Box::new(condition),
reason,
});
self
}
pub fn branch<C, T, E>(mut self, condition: C, then_branch: T, else_branch: Option<E>) -> Self
where
C: Fn(&S) -> bool + Send + Sync + 'static,
T: FnOnce(WorkflowBuilder<S>) -> WorkflowBuilder<S>,
E: FnOnce(WorkflowBuilder<S>) -> WorkflowBuilder<S>,
{
let then_builder = then_branch(WorkflowBuilder::new("then_branch"));
let else_flow = else_branch.map(|eb| eb(WorkflowBuilder::new("else_branch")).build());
self.steps.push(WorkflowStep::Branch {
condition: Box::new(condition),
then_flow: then_builder.build(),
else_flow,
});
self
}
pub fn build(self) -> Workflow<S> {
Workflow {
name: self.name,
steps: self.steps,
}
}
}