use std::collections::HashMap;
use std::sync::Arc;
use serde::{Deserialize, Serialize};
use thiserror::Error;
#[derive(Error, Debug)]
pub enum FlowError {
#[error("Execution error: {0}")]
Execution(String),
#[error("Node not found")]
NodeNotFound,
#[error("Parameter error: {0}")]
ParameterError(String),
}
pub type Result<T> = std::result::Result<T, FlowError>;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Params {
inner: HashMap<String, serde_json::Value>,
}
impl Params {
pub fn new() -> Self {
Self {
inner: HashMap::new(),
}
}
pub fn insert<S: Into<String>, V: Into<serde_json::Value>>(&mut self, key: S, value: V) {
self.inner.insert(key.into(), value.into());
}
pub fn get<S: AsRef<str>>(&self, key: S) -> Option<&serde_json::Value> {
self.inner.get(key.as_ref())
}
pub fn merge(&mut self, other: &Params) {
self.inner.extend(other.inner.clone());
}
pub fn remove<S: AsRef<str>>(&mut self, key: S) -> Option<serde_json::Value> {
self.inner.remove(key.as_ref())
}
}
impl Default for Params {
fn default() -> Self {
Self::new()
}
}
pub trait SharedData: Clone + Send + Sync {}
type NodeFunc = Arc<dyn Fn(&mut (dyn std::any::Any + Send), &Params) -> Result<Option<String>> + Send + Sync>;
#[derive(Clone)]
pub struct Node {
name: String,
params: Params,
successors: HashMap<String, Node>,
func: NodeFunc,
}
impl Node {
pub fn new<F>(name: impl Into<String>, func: F) -> Self
where
F: Fn(&mut (dyn std::any::Any + Send), &Params) -> Result<Option<String>> + Send + Sync + 'static,
{
Self {
name: name.into(),
params: Params::new(),
successors: HashMap::new(),
func: Arc::new(func),
}
}
pub fn add_successor(&mut self, action: impl Into<String>, node: Node) -> &mut Self {
self.successors.insert(action.into(), node);
self
}
pub fn next(&mut self, node: Node) -> &mut Self {
self.add_successor("default", node)
}
pub fn set_params(&mut self, params: Params) {
self.params = params;
}
pub fn get_params(&self) -> &Params {
&self.params
}
pub fn get_successor(&self, action: &str) -> Option<&Node> {
self.successors.get(action)
}
pub fn has_successors(&self) -> bool {
!self.successors.is_empty()
}
pub fn run(&self, shared: &mut (dyn std::any::Any + Send)) -> Result<()> {
if self.has_successors() {
eprintln!("Warning: Node won't run successors. Use Flow.");
}
(self.func)(shared, &self.params)?;
Ok(())
}
pub fn run_recursive(&self, shared: &mut (dyn std::any::Any + Send)) -> Result<()> {
let action = (self.func)(shared, &self.params)?;
if let Some(next_node) = action
.as_ref()
.and_then(|a| self.successors.get(a))
.or_else(|| self.successors.get("default")) {
next_node.run_recursive(shared)?;
}
Ok(())
}
}
pub struct Flow {
start_node: Option<Node>,
params: Params,
}
impl Flow {
pub fn new() -> Self {
Self {
start_node: None,
params: Params::new(),
}
}
pub fn start(mut self, node: Node) -> Self {
self.start_node = Some(node);
self
}
pub fn set_params(&mut self, params: Params) {
self.params = params;
}
pub fn run(&self, shared: &mut (dyn std::any::Any + Send)) -> Result<()> {
if let Some(ref node) = self.start_node {
let mut node = node.clone();
node.set_params(self.params.clone());
node.run_recursive(shared)?;
}
Ok(())
}
pub fn run_with_params(&self, shared: &mut (dyn std::any::Any + Send), params: Params) -> Result<()> {
if let Some(ref node) = self.start_node {
let mut node = node.clone();
let mut merged_params = self.params.clone();
merged_params.merge(¶ms);
node.set_params(merged_params);
node.run_recursive(shared)?;
}
Ok(())
}
}
impl Default for Flow {
fn default() -> Self {
Self::new()
}
}
pub struct BatchFlow {
start_node: Option<Node>,
params: Params,
}
impl BatchFlow {
pub fn new() -> Self {
Self {
start_node: None,
params: Params::new(),
}
}
pub fn start(mut self, node: Node) -> Self {
self.start_node = Some(node);
self
}
pub fn set_params(&mut self, params: Params) {
self.params = params;
}
pub fn run_batch(&self, shared: &mut (dyn std::any::Any + Send), batch_params: Vec<Params>) -> Result<Vec<()>> {
let mut results = Vec::with_capacity(batch_params.len());
for params in batch_params {
let flow = Flow {
start_node: self.start_node.clone(),
params: self.params.clone(),
};
flow.run_with_params(shared, params)?;
results.push(());
}
Ok(results)
}
}
impl Default for BatchFlow {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::{Arc, Mutex};
#[derive(Default, Clone)]
struct TestShared {
pub counter: Arc<Mutex<i32>>,
}
#[test]
fn test_basic_node() {
let mut shared = TestShared::default();
let node = Node::new("test", |shared, _params| {
if let Some(shared) = shared.downcast_mut::<TestShared>() {
let mut counter = shared.counter.lock().unwrap();
*counter += 1;
}
Ok(None)
});
node.run(&mut shared).unwrap();
let counter = shared.counter.lock().unwrap();
assert_eq!(*counter, 1);
}
#[test]
fn test_flow_execution() {
let mut shared = TestShared::default();
let node = Node::new("test", |shared, _params| {
if let Some(shared) = shared.downcast_mut::<TestShared>() {
let mut counter = shared.counter.lock().unwrap();
*counter += 1;
}
Ok(None)
});
let flow = Flow::new().start(node);
flow.run(&mut shared).unwrap();
let counter = shared.counter.lock().unwrap();
assert_eq!(*counter, 1);
}
#[test]
fn test_chained_flow() {
let mut shared = TestShared::default();
let mut node1 = Node::new("node1", |shared, _params| {
if let Some(shared) = shared.downcast_mut::<TestShared>() {
let mut counter = shared.counter.lock().unwrap();
*counter += 1;
}
Ok(None)
});
let node2 = Node::new("node2", |shared, _params| {
if let Some(shared) = shared.downcast_mut::<TestShared>() {
let mut counter = shared.counter.lock().unwrap();
*counter += 10;
}
Ok(None)
});
node1.next(node2);
let flow = Flow::new().start(node1);
flow.run(&mut shared).unwrap();
let counter = shared.counter.lock().unwrap();
assert_eq!(*counter, 11);
}
#[test]
fn test_batch_flow() {
let mut shared = TestShared::default();
let node = Node::new("batch", |shared, params| {
if let Some(shared) = shared.downcast_mut::<TestShared>() {
let mut counter = shared.counter.lock().unwrap();
if let Some(value) = params.get("value") {
if let Some(num) = value.as_i64() {
*counter += num as i32;
}
}
}
Ok(None)
});
let flow = BatchFlow::new().start(node);
let batch_params = vec![
{
let mut params = Params::new();
params.insert("value", 1);
params
},
{
let mut params = Params::new();
params.insert("value", 2);
params
},
];
flow.run_batch(&mut shared, batch_params
).unwrap();
let counter = shared.counter.lock().unwrap();
assert_eq!(*counter, 3);
}
#[test]
fn test_conditional_flow() {
let mut shared = TestShared::default();
let mut node1 = Node::new("node1", |_shared, params| {
let should_continue = params.get("continue").and_then(|v| v.as_bool()).unwrap_or(false);
if should_continue {
Ok(Some("continue".to_string()))
} else {
Ok(None)
}
});
let node2 = Node::new("node2", |shared, _params| {
if let Some(shared) = shared.downcast_mut::<TestShared>() {
let mut counter = shared.counter.lock().unwrap();
*counter += 100;
}
Ok(None)
});
node1.add_successor("continue", node2);
let mut params = Params::new();
params.insert("continue", true);
let flow = Flow::new().start(node1);
flow.run_with_params(&mut shared, params
).unwrap();
let counter = shared.counter.lock().unwrap();
assert_eq!(*counter, 100);
}
}
pub mod prelude {
pub use super::{Node, Flow, BatchFlow, Params, Result, FlowError};
}
#[macro_export]
macro_rules! node {
($name:expr, $func:expr) => {
Node::new($name, $func)
};
}
#[macro_export]
macro_rules! flow {
($($node:expr),+ $(,)?) => {{
let mut current = None;
$(current = Some($node);)*
Flow::new().start(current.unwrap())
}};
}
#[macro_export]
macro_rules! chain {
($first:expr $(, $rest:expr)* $(,)?) => {{
let mut current = $first;
$(current = current.next($rest);)*
current
}};
}