1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53

use std::collections::HashSet;

use crate::chain::Chain;

use super::SequentialChain;

pub struct SequentialChainBuilder {
    chains: Vec<Box<dyn Chain>>,
}

impl SequentialChainBuilder {
    pub fn new() -> Self {
        Self { chains: Vec::new() }
    }

    pub fn add_chain<C: Chain + 'static>(mut self, chain: C) -> Self {
        self.chains.push(Box::new(chain));
        self
    }

    pub fn build(self) -> SequentialChain {
        let outputs: HashSet<String> = self
            .chains
            .iter()
            .flat_map(|c| c.get_output_keys())
            .collect();

        let input_keys: HashSet<String> = self
            .chains
            .iter()
            .flat_map(|c| c.get_input_keys())
            .collect();

        SequentialChain {
            chains: self.chains,
            input_keys,
            outputs,
        }
    }
}

#[macro_export]
macro_rules! sequential_chain {
    ( $( $chain:expr ),* $(,)? ) => {
        {
            let mut builder = $crate::chain::SequentialChainBuilder::new();
            $(
                builder = builder.add_chain($chain);
            )*
            builder.build()
        }
    };
}