1use crate::SignedOrder;
2use alloy::primitives::{Address, U256};
3use serde::{Deserialize, Serialize};
4use signet_zenith::RollupOrders;
5use std::borrow::Cow;
6use std::collections::{HashMap, HashSet};
7
8#[derive(Debug, Default, Clone, Eq, PartialEq, Deserialize, Serialize)]
10pub struct AggregateOrders {
11 pub outputs: HashMap<(u64, Address), HashMap<Address, U256>>,
15 pub inputs: HashMap<Address, U256>,
18}
19
20impl AggregateOrders {
21 pub fn new() -> Self {
23 Default::default()
24 }
25
26 pub fn with_capacity(capacity: usize) -> Self {
29 Self { outputs: HashMap::with_capacity(capacity), inputs: HashMap::with_capacity(capacity) }
30 }
31
32 pub(crate) fn ingest_output(&mut self, output: &RollupOrders::Output) {
34 self.ingest_raw_output(
35 output.chainId as u64,
36 output.token,
37 output.recipient,
38 output.amount,
39 );
40 }
41
42 fn ingest_raw_output(
44 &mut self,
45 chain_id: u64,
46 token: Address,
47 recipient: Address,
48 amount: U256,
49 ) {
50 let entry =
51 self.outputs.entry((chain_id, token)).or_default().entry(recipient).or_default();
52 *entry = entry.saturating_add(amount);
53 }
54
55 pub(crate) fn ingest_input(&mut self, input: &RollupOrders::Input) {
57 self.ingest_raw_input(input.token, input.amount);
58 }
59
60 fn ingest_raw_input(&mut self, token: Address, amount: U256) {
62 let entry = self.inputs.entry(token).or_default();
63 *entry = entry.saturating_add(amount);
64 }
65
66 pub fn ingest(&mut self, order: &RollupOrders::Order) {
68 order.outputs.iter().for_each(|o| self.ingest_output(o));
69 order.inputs.iter().for_each(|i| self.ingest_input(i));
70 }
71
72 pub fn ingest_signed(&mut self, order: &SignedOrder) {
74 order
75 .outputs
76 .iter()
77 .for_each(|o| self.ingest_raw_output(o.chainId as u64, o.token, o.recipient, o.amount));
78 order
79 .permit
80 .permit
81 .permitted
82 .iter()
83 .for_each(|tp| self.ingest_raw_input(tp.token, tp.amount));
84 }
85
86 pub fn extend<'a>(&mut self, orders: impl IntoIterator<Item = &'a RollupOrders::Order>) {
88 for order in orders {
89 self.ingest(order);
90 }
91 }
92
93 pub fn extend_signed<'a>(&mut self, orders: impl IntoIterator<Item = &'a SignedOrder>) {
95 for order in orders {
96 self.ingest_signed(order);
97 }
98 }
99
100 pub fn destination_chain_ids(&self) -> Vec<u64> {
102 HashSet::<u64>::from_iter(self.outputs.keys().map(|(chain_id, _)| *chain_id))
103 .into_iter()
104 .collect()
105 }
106
107 pub fn outputs_for(&self, target_chain_id: u64) -> Vec<RollupOrders::Output> {
109 let mut o = Vec::new();
110 for ((chain_id, token), recipient_map) in &self.outputs {
111 if *chain_id == target_chain_id {
112 for (recipient, amount) in recipient_map {
113 o.push(RollupOrders::Output {
114 token: *token,
115 amount: U256::from(*amount),
116 recipient: *recipient,
117 chainId: *chain_id as u32,
118 });
119 }
120 }
121 }
122 o
123 }
124}
125
126impl<'a> FromIterator<&'a RollupOrders::Order> for AggregateOrders {
127 fn from_iter<T: IntoIterator<Item = &'a RollupOrders::Order>>(iter: T) -> Self {
128 let mut orders = AggregateOrders::new();
129 orders.extend(iter);
130 orders
131 }
132}
133
134impl<'a> FromIterator<&'a SignedOrder> for AggregateOrders {
135 fn from_iter<T: IntoIterator<Item = &'a SignedOrder>>(iter: T) -> Self {
136 let mut orders = AggregateOrders::new();
137 orders.extend_signed(iter);
138 orders
139 }
140}
141
142impl<'a> From<&'a RollupOrders::Order> for AggregateOrders {
143 fn from(order: &'a RollupOrders::Order) -> Self {
144 let mut orders = AggregateOrders::new();
145 orders.ingest(order);
146 orders
147 }
148}
149
150impl<'a> From<&'a SignedOrder> for AggregateOrders {
151 fn from(order: &'a SignedOrder) -> Self {
152 let mut orders = AggregateOrders::new();
153 orders.ingest_signed(order);
154 orders
155 }
156}
157
158impl<'a> From<&'a AggregateOrders> for Cow<'a, AggregateOrders> {
159 fn from(orders: &'a AggregateOrders) -> Self {
160 Cow::Borrowed(orders)
161 }
162}
163
164#[cfg(test)]
165mod test {
166 use super::*;
167 use alloy::primitives::{Address, U256};
168
169 const ASSET_A: Address = Address::repeat_byte(1);
170 const ASSET_B: Address = Address::repeat_byte(2);
171 const ASSET_C: Address = Address::repeat_byte(3);
172
173 const USER_A: Address = Address::repeat_byte(4);
174 const USER_B: Address = Address::repeat_byte(5);
175 const USER_C: Address = Address::repeat_byte(6);
176
177 fn input(asset: Address, amount: u64) -> RollupOrders::Input {
178 RollupOrders::Input { token: asset, amount: U256::from(amount) }
179 }
180
181 fn output(asset: Address, recipient: Address, amount: u64) -> RollupOrders::Output {
182 RollupOrders::Output { chainId: 1, token: asset, recipient, amount: U256::from(amount) }
183 }
184
185 #[test]
186 fn test_single_order() {
187 let order = RollupOrders::Order {
188 inputs: vec![input(ASSET_A, 100), input(ASSET_B, 200)],
189 outputs: vec![
190 output(ASSET_A, USER_A, 50),
191 output(ASSET_A, USER_B, 50),
192 output(ASSET_B, USER_B, 100),
193 output(ASSET_C, USER_C, 200),
194 output(ASSET_C, USER_C, 200),
195 ],
196 deadline: U256::ZERO,
197 };
198
199 let agg: AggregateOrders = [&order].into_iter().collect();
200 assert_eq!(agg.inputs.get(&ASSET_A), Some(&U256::from(100)), "ASSET_A input");
201 assert_eq!(agg.inputs.get(&ASSET_B), Some(&U256::from(200)), "ASSET_B input");
202
203 assert_eq!(
204 agg.outputs.get(&(1, ASSET_A)).map(|m| m.get(&USER_A)),
205 Some(Some(&U256::from(50))),
206 "ASSET_A USER_A output"
207 );
208 assert_eq!(
209 agg.outputs.get(&(1, ASSET_A)).map(|m| m.get(&USER_B)),
210 Some(Some(&U256::from(50))),
211 "ASSET_A USER_B output"
212 );
213 assert_eq!(
214 agg.outputs.get(&(1, ASSET_B)).map(|m| m.get(&USER_B)),
215 Some(Some(&U256::from(100))),
216 "ASSET_B USER_B output"
217 );
218 assert_eq!(
219 agg.outputs.get(&(1, ASSET_C)).map(|m| m.get(&USER_C)),
220 Some(Some(&U256::from(400))),
221 "ASSET_C USER_C output"
222 );
223 }
224
225 #[test]
226 fn test_two_orders() {
227 let order_1 = RollupOrders::Order {
228 inputs: vec![input(ASSET_A, 100), input(ASSET_B, 200)],
229 outputs: vec![
230 output(ASSET_A, USER_A, 50),
231 output(ASSET_A, USER_B, 50),
232 output(ASSET_B, USER_B, 100),
233 output(ASSET_C, USER_C, 200),
234 output(ASSET_C, USER_C, 200),
235 ],
236 deadline: U256::ZERO,
237 };
238 let order_2 = RollupOrders::Order {
239 inputs: vec![input(ASSET_A, 50), input(ASSET_C, 100)],
240 outputs: vec![
241 output(ASSET_A, USER_A, 50),
242 output(ASSET_B, USER_B, 100),
243 output(ASSET_C, USER_C, 100),
244 ],
245 deadline: U256::ZERO,
246 };
247
248 let agg: AggregateOrders = [&order_1, &order_2].into_iter().collect();
249
250 assert_eq!(agg.inputs.get(&ASSET_A), Some(&U256::from(150)), "ASSET_A input");
251 assert_eq!(agg.inputs.get(&ASSET_B), Some(&U256::from(200)), "ASSET_B input");
252 assert_eq!(agg.inputs.get(&ASSET_C), Some(&U256::from(100)), "ASSET_C input");
253
254 assert_eq!(
255 agg.outputs.get(&(1, ASSET_A)).map(|m| m.get(&USER_A)),
256 Some(Some(&U256::from(100))),
257 "ASSET_A USER_A output"
258 );
259 assert_eq!(
260 agg.outputs.get(&(1, ASSET_A)).map(|m| m.get(&USER_B)),
261 Some(Some(&U256::from(50))),
262 "ASSET_A USER_B output"
263 );
264 assert_eq!(
265 agg.outputs.get(&(1, ASSET_B)).map(|m| m.get(&USER_B)),
266 Some(Some(&U256::from(200))),
267 "ASSET_B USER_B output"
268 );
269 assert_eq!(
270 agg.outputs.get(&(1, ASSET_C)).map(|m| m.get(&USER_C)),
271 Some(Some(&U256::from(500))),
272 "ASSET_C USER_C output"
273 );
274 }
275}