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

use brontes_types::{
    normalized_actions::{
        Action, MultiCallFrameClassification, MultiFrameAction, MultiFrameRequest,
    },
    Protocol, TreeSearchBuilder,
};

use crate::multi_frame_classification::MultiCallFrameClassifier;

pub struct ZeroXAgg;

impl MultiCallFrameClassifier for ZeroXAgg {
    const KEY: [u8; 2] = [Protocol::ZeroX as u8, MultiFrameAction::Aggregator as u8];

    fn create_classifier(
        request: MultiFrameRequest,
    ) -> Option<MultiCallFrameClassification<Action>> {
        Some(MultiCallFrameClassification {
            trace_index:         request.trace_idx,
            tree_search_builder: TreeSearchBuilder::new().with_actions([
                Action::is_swap,
                Action::is_transfer,
                Action::is_eth_transfer,
            ]),
            parse_fn:            Box::new(|this_action, child_nodes| {
                let this = this_action.try_aggregator_mut().unwrap();
                let mut prune_nodes = Vec::new();

                for (trace_index, action) in child_nodes {
                    match action {
                        Action::Swap(_)
                        | Action::SwapWithFee(_)
                        | Action::Transfer(_)
                        | Action::EthTransfer(_) => {
                            this.child_actions.push(action.clone());
                            prune_nodes.push(trace_index);
                        }
                        _ => {}
                    }
                }
                prune_nodes
            }),
        })
    }
}