Skip to main content

o2_tools/
helpers.rs

1use crate::{
2    CallOption,
3    call_handler_ext::CallHandlerExt,
4    market_data::{
5        OrderData,
6        order_book::Balances,
7    },
8    order_book::{
9        CreateOrderParams,
10        OrderBookManager,
11        OrderType,
12    },
13    order_book_deploy::{
14        OrderBookDeploy,
15        OrderCancelledEvent,
16        OrderCreatedEvent,
17        OrderMatchedEvent,
18        SharedExecutionEvent,
19        SharedOrderMatchedEvent,
20        TriggerOrderCreatedEvent,
21        TriggerOrderEjectedEvent,
22    },
23    trade_account::{
24        CallContractArgs,
25        TradeAccountManager,
26    },
27    trade_account_deploy::{
28        DeployConfig,
29        TradeAccountDeploy,
30    },
31};
32use fuels::{
33    prelude::*,
34    programs::responses::CallResponse,
35    types::{
36        Address,
37        Bytes32,
38        ContractId,
39        Identity,
40        tx_status::TxStatus,
41    },
42};
43use std::collections::{
44    HashMap,
45    HashSet,
46};
47
48pub async fn setup_order_book<W: Account + Clone>(
49    deployer_wallet: &W,
50    base_asset: AssetId,
51    quote_asset: AssetId,
52    deploy_config: &crate::order_book_deploy::OrderBookDeployConfig,
53) -> anyhow::Result<OrderBookManager<W>, anyhow::Error> {
54    let order_book_deploy =
55        OrderBookDeploy::deploy(deployer_wallet, base_asset, quote_asset, deploy_config)
56            .await?;
57    let order_book = OrderBookManager::new(deployer_wallet, 9, 9, &order_book_deploy);
58    Ok(order_book)
59}
60
61pub async fn setup_trade_accounts(
62    deployer_wallet: &Wallet,
63    contract_ids: &[ContractId],
64    wallets: &mut Vec<Wallet>,
65) -> anyhow::Result<Vec<TradeAccountManager<Wallet>>, anyhow::Error> {
66    let config = crate::trade_account_deploy::TradeAccountDeployConfig::default();
67    let deploy_config = DeployConfig::Latest(config);
68    let trade_account_deploy: TradeAccountDeploy<Wallet> =
69        TradeAccountDeploy::deploy(deployer_wallet, &deploy_config).await?;
70    let mut trade_accounts = Vec::with_capacity(wallets.len());
71
72    // Create remaining trade accounts
73    for user_wallet in wallets {
74        let deployment = trade_account_deploy
75            .deploy_with_account(
76                &user_wallet.address().into(),
77                &deploy_config,
78                &CallOption::AwaitBlock,
79                None,
80                &[],
81            )
82            .await?;
83        let trade_account = TradeAccountManager::create_with_session(
84            &user_wallet.clone(),
85            &user_wallet.clone(),
86            contract_ids,
87            &deployment,
88            CallOption::AwaitBlock,
89        )
90        .await?;
91        trade_accounts.push(trade_account);
92    }
93
94    Ok(trade_accounts)
95}
96
97pub async fn fund_trade_accounts<W: Account + Clone>(
98    trade_accounts: &[TradeAccountManager<W>],
99    base_asset: AssetId,
100    base_asset_amount: u64,
101    quote_asset: AssetId,
102    quote_asset_amount: u64,
103) -> anyhow::Result<(), anyhow::Error> {
104    for trade_account in trade_accounts.iter() {
105        let _ = trade_account
106            .owner
107            .force_transfer_to_contract(
108                trade_account.contract.contract_id(),
109                base_asset_amount,
110                base_asset,
111                TxPolicies::default(),
112            )
113            .await?;
114        let _ = trade_account
115            .owner
116            .force_transfer_to_contract(
117                trade_account.contract.contract_id(),
118                quote_asset_amount,
119                quote_asset,
120                TxPolicies::default(),
121            )
122            .await?;
123    }
124    Ok(())
125}
126
127pub async fn create_order_call(
128    order_book: &OrderBookManager<Wallet>,
129    order_data: &OrderData,
130    order_type: OrderType,
131    trade_account: &mut TradeAccountManager<Wallet>,
132    gas_per_method: Option<u64>,
133) -> anyhow::Result<CallContractArgs> {
134    let create_order_params = CreateOrderParams {
135        price: order_data.price,
136        quantity: order_data.quantity,
137        side: order_data.side,
138        asset_id: order_book.get_order_side_asset(&order_data.side),
139        order_type,
140    };
141    // Create orders args with signatures
142    let contract_call_args = trade_account
143        .create_order(order_book, &create_order_params, gas_per_method)
144        .await?;
145
146    Ok(contract_call_args)
147}
148
149pub async fn create_order_handlers<'a, I>(
150    order_book: &OrderBookManager<Wallet>,
151    orders: &[OrderData],
152    trade_accounts: I,
153    gas_per_method: Option<u64>,
154) -> anyhow::Result<
155    Vec<CallHandler<Wallet, fuels::programs::calls::ContractCall, ()>>,
156    anyhow::Error,
157>
158where
159    I: Iterator<Item = &'a mut TradeAccountManager<Wallet>>,
160{
161    let mut create_orders_handlers = Vec::with_capacity(orders.len());
162    let mut trade_accounts = trade_accounts.into_iter().collect::<Vec<_>>();
163    for order_data in orders.iter() {
164        let trade_account = trade_accounts
165            .iter_mut()
166            .find(|ta| ta.identity() == order_data.trader_id)
167            .unwrap();
168
169        // Create orders args with signatures
170        let contract_call_args = create_order_call(
171            order_book,
172            order_data,
173            OrderType::Spot,
174            trade_account,
175            gas_per_method,
176        )
177        .await?;
178
179        // Create function handler
180        create_orders_handlers.push(
181            trade_account
182                .session_call_contract(&contract_call_args)
183                .with_contract_ids(&[order_book.contract.contract_id()]),
184        );
185    }
186    Ok(create_orders_handlers)
187}
188
189pub async fn cancel_order_call(
190    order_book: &OrderBookManager<Wallet>,
191    order_id: &Bytes32,
192    trade_account: &mut TradeAccountManager<Wallet>,
193    gas_per_method: Option<u64>,
194) -> anyhow::Result<CallContractArgs> {
195    trade_account
196        .cancel_order(order_book, *order_id, gas_per_method)
197        .await
198}
199
200pub async fn cancel_order_handlers(
201    order_book: &OrderBookManager<Wallet>,
202    orders: &[Bytes32],
203    trade_account: &mut TradeAccountManager<Wallet>,
204    gas_per_method: Option<u64>,
205) -> anyhow::Result<
206    Vec<CallHandler<Wallet, fuels::programs::calls::ContractCall, ()>>,
207    anyhow::Error,
208> {
209    let mut cancel_orders_handlers = Vec::with_capacity(orders.len());
210    for order_id in orders.iter() {
211        let contract_call_args =
212            cancel_order_call(order_book, order_id, trade_account, gas_per_method)
213                .await?;
214
215        // Create function handler
216        cancel_orders_handlers
217            .push(trade_account.session_call_contract(&contract_call_args));
218    }
219    Ok(cancel_orders_handlers)
220}
221
222pub async fn send_transactions(
223    create_orders_handlers: Vec<
224        CallHandler<Wallet, fuels::programs::calls::ContractCall, ()>,
225    >,
226    gaspayer_wallet: Wallet,
227    call_option: CallOption,
228) -> Vec<Bytes32> {
229    let mut order_results = Vec::with_capacity(create_orders_handlers.len());
230    for mut call_handler in create_orders_handlers {
231        call_handler.account = gaspayer_wallet.clone();
232
233        let tx_id = match call_option.clone() {
234            CallOption::AwaitBlock => call_handler.submit().await.unwrap().tx_id(),
235            CallOption::AwaitPreconfirmation(ops) => {
236                call_handler
237                    .almost_sync_call(
238                        &ops.data_builder,
239                        &ops.utxo_manager,
240                        &ops.tx_config,
241                        None,
242                        &[],
243                    )
244                    .await
245                    .unwrap()
246                    .tx_id
247            }
248        };
249        order_results.push(tx_id);
250    }
251    order_results
252}
253
254#[derive(Debug, Clone, Default)]
255pub struct OrderBookEvents {
256    pub matches: Vec<OrderMatchedEvent>,
257    /// Shared fills emitted by the requested canonical book.
258    pub shared_matches: Vec<SharedOrderMatchedEvent>,
259    /// Self-contained executions emitted by the requested canonical book.
260    pub shared_executions: Vec<SharedExecutionEvent>,
261    pub orders: Vec<OrderCreatedEvent>,
262    pub cancels: Vec<OrderCancelledEvent>,
263    pub trigger_orders: Vec<TriggerOrderCreatedEvent>,
264    pub trigger_ejections: Vec<TriggerOrderEjectedEvent>,
265}
266
267/// Decode only events emitted by the manager's contract.
268/// House fills are emitted by canonical; decode them with a canonical manager
269/// or explicitly use `HouseOrderBookDeploy::decode_shared_matches`.
270pub fn get_order_book_events<W: Account + Clone>(
271    order_book: &OrderBookManager<W>,
272    order_book_events: &mut OrderBookEvents,
273    tx_result: &TxStatus,
274) -> anyhow::Result<(), anyhow::Error> {
275    if let TxStatus::Success(success) = tx_result {
276        let mut order_created_events =
277            order_book
278                .contract
279                .log_decoder()
280                .decode_logs_with_type::<OrderCreatedEvent>(&success.receipts)?;
281        let mut order_match_events =
282            order_book
283                .contract
284                .log_decoder()
285                .decode_logs_with_type::<OrderMatchedEvent>(&success.receipts)?;
286        let mut shared_matches = crate::house_order_book::decode_legacy_shared_matches(
287            order_book.contract.contract_id(),
288            &success.receipts,
289        )?;
290        let mut shared_executions =
291            order_book
292                .contract
293                .log_decoder()
294                .decode_logs_with_type::<SharedExecutionEvent>(&success.receipts)?;
295        let mut order_cancel_events =
296            order_book
297                .contract
298                .log_decoder()
299                .decode_logs_with_type::<OrderCancelledEvent>(&success.receipts)?;
300
301        let mut trigger_created_events =
302            order_book
303                .contract
304                .log_decoder()
305                .decode_logs_with_type::<TriggerOrderCreatedEvent>(&success.receipts)?;
306        let mut trigger_ejected_events =
307            order_book
308                .contract
309                .log_decoder()
310                .decode_logs_with_type::<TriggerOrderEjectedEvent>(&success.receipts)?;
311
312        order_book_events.orders.append(&mut order_created_events);
313        order_book_events.matches.append(&mut order_match_events);
314        order_book_events.shared_matches.append(&mut shared_matches);
315        order_book_events
316            .shared_executions
317            .append(&mut shared_executions);
318        order_book_events.cancels.append(&mut order_cancel_events);
319        order_book_events
320            .trigger_orders
321            .append(&mut trigger_created_events);
322        order_book_events
323            .trigger_ejections
324            .append(&mut trigger_ejected_events);
325    }
326    Ok(())
327}
328
329pub async fn get_wallets_balances<W: Account + Clone>(
330    wallets: &[W],
331    base_asset: AssetId,
332    quote_asset: AssetId,
333) -> anyhow::Result<HashMap<Address, (u128, u128)>, anyhow::Error> {
334    let mut balances = HashMap::new();
335    for wallet in wallets {
336        let balance = wallet.get_balances().await?;
337        let base_asset_balance =
338            balance.get(&base_asset.to_string()).cloned().unwrap_or(0);
339        let quote_asset_balance =
340            balance.get(&quote_asset.to_string()).cloned().unwrap_or(0);
341        balances.insert(wallet.address(), (base_asset_balance, quote_asset_balance));
342    }
343    Ok(balances)
344}
345
346pub async fn get_contracts_balances(
347    provider: &Provider,
348    contracts: &[ContractId],
349    base_asset: &AssetId,
350    quote_asset: &AssetId,
351) -> anyhow::Result<Balances, anyhow::Error> {
352    let mut balances = Balances::new();
353    for contract_id in contracts {
354        let balance = provider.get_contract_balances(contract_id).await?;
355        let base_asset_balance = balance.get(&base_asset.clone()).cloned().unwrap_or(0);
356        let quote_asset_balance = balance.get(&quote_asset.clone()).cloned().unwrap_or(0);
357        balances.insert(
358            Identity::ContractId(*contract_id),
359            (base_asset_balance, quote_asset_balance),
360        );
361    }
362    Ok(balances)
363}
364
365pub async fn settle_trade_accounts_balances<'a, I>(
366    fee_payer: &Wallet,
367    order_book: &OrderBookManager<Wallet>,
368    trade_accounts: I,
369    call_option: CallOption,
370) -> anyhow::Result<CallResponse<()>, anyhow::Error>
371where
372    I: Iterator<Item = &'a TradeAccountManager<Wallet>>,
373{
374    let mut accounts = vec![];
375    let mut contracts = vec![];
376    for account in trade_accounts {
377        accounts.push(Identity::from(account.contract.contract_id()));
378        contracts.push(account.contract.contract_id());
379    }
380
381    let mut call_handler = order_book
382        .contract
383        .methods()
384        .settle_balances(accounts)
385        .with_contract_ids(&contracts);
386    call_handler.account = fee_payer.clone();
387
388    let result = match call_option {
389        CallOption::AwaitBlock => call_handler.call().await?,
390        CallOption::AwaitPreconfirmation(ops) => {
391            call_handler
392                .almost_sync_call(
393                    &ops.data_builder,
394                    &ops.utxo_manager,
395                    &ops.tx_config,
396                    None,
397                    &[],
398                )
399                .await?
400                .tx_status?
401        }
402    };
403
404    Ok(result)
405}
406
407pub async fn get_trade_accounts_balances<W: Account + Clone>(
408    provider: &Provider,
409    trade_accounts: &[TradeAccountManager<W>],
410    base_asset: &AssetId,
411    quote_asset: &AssetId,
412) -> anyhow::Result<Balances, anyhow::Error> {
413    let contracts = trade_accounts
414        .iter()
415        .map(|trade_account| trade_account.contract.contract_id())
416        .collect::<Vec<_>>();
417    get_contracts_balances(provider, &contracts, base_asset, quote_asset).await
418}
419
420/// Wait for transactions and collect only events emitted by the manager's contract.
421pub async fn wait_for_book_events<W: Account + Clone>(
422    tx_ids: &[Bytes32],
423    trade_account: &TradeAccountManager<W>,
424    order_book: &OrderBookManager<W>,
425    gaspayer_wallet: W,
426) -> anyhow::Result<OrderBookEvents, anyhow::Error> {
427    let mut order_book_events = OrderBookEvents::default();
428    let mut tx_completed = HashSet::new();
429
430    while tx_completed.len() != tx_ids.len() {
431        let provider = gaspayer_wallet.try_provider()?;
432        for order_result in tx_ids.iter() {
433            let result = provider.get_transaction_by_id(order_result).await?;
434
435            if let Some(result) = result {
436                get_order_book_events(
437                    order_book,
438                    &mut order_book_events,
439                    &result.status,
440                )?;
441                match result.status {
442                    TxStatus::Success(_) => {
443                        tx_completed.insert(*order_result);
444                    }
445                    TxStatus::Failure(failure) => {
446                        let logs = order_book
447                            .contract
448                            .log_decoder()
449                            .decode_logs(&failure.receipts);
450                        let logs_trade = trade_account
451                            .contract
452                            .log_decoder()
453                            .decode_logs(&failure.receipts);
454                        println!("{logs:#?}");
455                        println!("{logs_trade:#?}");
456                        panic!("{:#}", failure.reason);
457                    }
458                    _ => {
459                        continue;
460                    }
461                }
462            }
463        }
464    }
465
466    Ok(order_book_events)
467}
468
469pub fn get_asset_balance(balances: &HashMap<String, u128>, asset_id: &AssetId) -> u128 {
470    *balances.get(&asset_id.to_string()).unwrap_or(&0)
471}
472
473pub fn get_total_amount(quantity: u64, price: u64, decimals: u64) -> u64 {
474    (quantity * price) / 10u64.pow(decimals as u32)
475}