zenith_types/orders/
agg.rs

1use crate::RollupOrders;
2use alloy::primitives::{Address, U256};
3use std::collections::HashMap;
4
5/// Aggregated orders for a transaction or set of transactions.
6#[derive(Debug, Default, Clone, Eq, PartialEq)]
7pub struct AggregateOrders {
8    /// Outputs to be transferred to the user. These may be on the rollup or
9    /// the host or potentially elsewhere in the future.
10    pub outputs: HashMap<(u64, Address), HashMap<Address, U256>>,
11    /// Inputs to be transferred to the filler. These are always on the
12    /// rollup.
13    pub inputs: HashMap<Address, U256>,
14}
15
16impl AggregateOrders {
17    /// Instantiate a new [`AggregateOrders`].
18    pub fn new() -> Self {
19        Default::default()
20    }
21
22    /// Instantiate a new [`AggregateOrders`] with a custom capacity. The
23    /// capcity is for the number of assets in inputs or outputs.
24    pub fn with_capacity(capacity: usize) -> Self {
25        Self { outputs: HashMap::with_capacity(capacity), inputs: HashMap::with_capacity(capacity) }
26    }
27
28    /// Ingest an output into the aggregate orders.
29    fn ingest_output(&mut self, output: &RollupOrders::Output) {
30        let entry = self
31            .outputs
32            .entry((output.chain_id() as u64, output.token))
33            .or_default()
34            .entry(output.recipient)
35            .or_default();
36        *entry = entry.saturating_add(output.amount);
37    }
38
39    /// Ingest an input into the aggregate orders.
40    fn ingest_input(&mut self, input: &RollupOrders::Input) {
41        let entry = self.inputs.entry(input.token).or_default();
42        *entry = entry.saturating_add(input.amount);
43    }
44
45    /// Ingest a new order into the aggregate orders.
46    pub fn ingest(&mut self, order: &RollupOrders::Order) {
47        order.outputs.iter().for_each(|o| self.ingest_output(o));
48        order.inputs.iter().for_each(|i| self.ingest_input(i));
49    }
50
51    /// Extend the orders with a new set of orders.
52    pub fn extend<'a>(&mut self, orders: impl IntoIterator<Item = &'a RollupOrders::Order>) {
53        for order in orders {
54            self.ingest(order);
55        }
56    }
57}
58
59impl<'a> FromIterator<&'a RollupOrders::Order> for AggregateOrders {
60    fn from_iter<T: IntoIterator<Item = &'a RollupOrders::Order>>(iter: T) -> Self {
61        let mut orders = AggregateOrders::new();
62        orders.extend(iter);
63        orders
64    }
65}
66
67#[cfg(test)]
68mod test {
69    use super::*;
70    use alloy::primitives::{Address, U256};
71
72    const ASSET_A: Address = Address::repeat_byte(1);
73    const ASSET_B: Address = Address::repeat_byte(2);
74    const ASSET_C: Address = Address::repeat_byte(3);
75
76    const USER_A: Address = Address::repeat_byte(4);
77    const USER_B: Address = Address::repeat_byte(5);
78    const USER_C: Address = Address::repeat_byte(6);
79
80    fn input(asset: Address, amount: u64) -> RollupOrders::Input {
81        RollupOrders::Input { token: asset, amount: U256::from(amount) }
82    }
83
84    fn output(asset: Address, recipient: Address, amount: u64) -> RollupOrders::Output {
85        RollupOrders::Output { chainId: 1, token: asset, recipient, amount: U256::from(amount) }
86    }
87
88    #[test]
89    fn test_single_order() {
90        let order = RollupOrders::Order {
91            inputs: vec![input(ASSET_A, 100), input(ASSET_B, 200)],
92            outputs: vec![
93                output(ASSET_A, USER_A, 50),
94                output(ASSET_A, USER_B, 50),
95                output(ASSET_B, USER_B, 100),
96                output(ASSET_C, USER_C, 200),
97                output(ASSET_C, USER_C, 200),
98            ],
99            deadline: U256::ZERO,
100        };
101
102        let agg: AggregateOrders = [&order].into_iter().collect();
103        assert_eq!(agg.inputs.get(&ASSET_A), Some(&U256::from(100)), "ASSET_A input");
104        assert_eq!(agg.inputs.get(&ASSET_B), Some(&U256::from(200)), "ASSET_B input");
105
106        assert_eq!(
107            agg.outputs.get(&(1, ASSET_A)).map(|m| m.get(&USER_A)),
108            Some(Some(&U256::from(50))),
109            "ASSET_A USER_A output"
110        );
111        assert_eq!(
112            agg.outputs.get(&(1, ASSET_A)).map(|m| m.get(&USER_B)),
113            Some(Some(&U256::from(50))),
114            "ASSET_A USER_B output"
115        );
116        assert_eq!(
117            agg.outputs.get(&(1, ASSET_B)).map(|m| m.get(&USER_B)),
118            Some(Some(&U256::from(100))),
119            "ASSET_B USER_B output"
120        );
121        assert_eq!(
122            agg.outputs.get(&(1, ASSET_C)).map(|m| m.get(&USER_C)),
123            Some(Some(&U256::from(400))),
124            "ASSET_C USER_C output"
125        );
126    }
127
128    #[test]
129    fn test_two_orders() {
130        let order_1 = RollupOrders::Order {
131            inputs: vec![input(ASSET_A, 100), input(ASSET_B, 200)],
132            outputs: vec![
133                output(ASSET_A, USER_A, 50),
134                output(ASSET_A, USER_B, 50),
135                output(ASSET_B, USER_B, 100),
136                output(ASSET_C, USER_C, 200),
137                output(ASSET_C, USER_C, 200),
138            ],
139            deadline: U256::ZERO,
140        };
141        let order_2 = RollupOrders::Order {
142            inputs: vec![input(ASSET_A, 50), input(ASSET_C, 100)],
143            outputs: vec![
144                output(ASSET_A, USER_A, 50),
145                output(ASSET_B, USER_B, 100),
146                output(ASSET_C, USER_C, 100),
147            ],
148            deadline: U256::ZERO,
149        };
150
151        let agg: AggregateOrders = [&order_1, &order_2].into_iter().collect();
152
153        assert_eq!(agg.inputs.get(&ASSET_A), Some(&U256::from(150)), "ASSET_A input");
154        assert_eq!(agg.inputs.get(&ASSET_B), Some(&U256::from(200)), "ASSET_B input");
155        assert_eq!(agg.inputs.get(&ASSET_C), Some(&U256::from(100)), "ASSET_C input");
156
157        assert_eq!(
158            agg.outputs.get(&(1, ASSET_A)).map(|m| m.get(&USER_A)),
159            Some(Some(&U256::from(100))),
160            "ASSET_A USER_A output"
161        );
162        assert_eq!(
163            agg.outputs.get(&(1, ASSET_A)).map(|m| m.get(&USER_B)),
164            Some(Some(&U256::from(50))),
165            "ASSET_A USER_B output"
166        );
167        assert_eq!(
168            agg.outputs.get(&(1, ASSET_B)).map(|m| m.get(&USER_B)),
169            Some(Some(&U256::from(200))),
170            "ASSET_B USER_B output"
171        );
172        assert_eq!(
173            agg.outputs.get(&(1, ASSET_C)).map(|m| m.get(&USER_C)),
174            Some(Some(&U256::from(500))),
175            "ASSET_C USER_C output"
176        );
177    }
178}