1use crate::RollupOrders;
2use alloy::primitives::{Address, U256};
3use std::collections::HashMap;
4
5#[derive(Debug, Default, Clone, Eq, PartialEq)]
7pub struct AggregateOrders {
8 pub outputs: HashMap<(u64, Address), HashMap<Address, U256>>,
11 pub inputs: HashMap<Address, U256>,
14}
15
16impl AggregateOrders {
17 pub fn new() -> Self {
19 Default::default()
20 }
21
22 pub fn with_capacity(capacity: usize) -> Self {
25 Self { outputs: HashMap::with_capacity(capacity), inputs: HashMap::with_capacity(capacity) }
26 }
27
28 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 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 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 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}