o2_tools/

helpers.rs

use crate::{
    CallOption,
    call_handler_ext::CallHandlerExt,
    market_data::{
        OrderData,
        order_book::Balances,
    },
    order_book::{
        CreateOrderParams,
        OrderBookManager,
        OrderType,
    },
    order_book_deploy::{
        OrderBookDeploy,
        OrderCancelledEvent,
        OrderCreatedEvent,
        OrderMatchedEvent,
    },
    trade_account::{
        CallContractArgs,
        TradeAccountManager,
    },
    trade_account_deploy::{
        DeployConfig,
        TradeAccountDeploy,
    },
};
use fuels::{
    prelude::*,
    programs::responses::CallResponse,
    types::{
        Address,
        Bytes32,
        ContractId,
        Identity,
        tx_status::TxStatus,
    },
};
use std::collections::{
    HashMap,
    HashSet,
};

pub async fn setup_order_book<W: Account + Clone>(
    deployer_wallet: &W,
    base_asset: AssetId,
    quote_asset: AssetId,
    deploy_config: &crate::order_book_deploy::OrderBookDeployConfig,
) -> anyhow::Result<OrderBookManager<W>, anyhow::Error> {
    let order_book_deploy =
        OrderBookDeploy::deploy(deployer_wallet, base_asset, quote_asset, deploy_config)
            .await?;
    let order_book = OrderBookManager::new(deployer_wallet, 9, 9, &order_book_deploy);
    Ok(order_book)
}

pub async fn setup_trade_accounts(
    deployer_wallet: &Wallet,
    contract_ids: &[ContractId],
    wallets: &mut Vec<Wallet>,
) -> anyhow::Result<Vec<TradeAccountManager<Wallet>>, anyhow::Error> {
    let config = crate::trade_account_deploy::TradeAccountDeployConfig::default();
    let deploy_config = DeployConfig::Latest(config);
    let trade_account_deploy: TradeAccountDeploy<Wallet> =
        TradeAccountDeploy::deploy(deployer_wallet, &deploy_config).await?;
    let mut trade_accounts = Vec::with_capacity(wallets.len());

    // Create remaining trade accounts
    for user_wallet in wallets {
        let deployment = trade_account_deploy
            .deploy_with_account(
                &user_wallet.address().into(),
                &deploy_config,
                &CallOption::AwaitBlock,
            )
            .await?;
        let trade_account = TradeAccountManager::create_with_session(
            &user_wallet.clone(),
            &user_wallet.clone(),
            contract_ids,
            &deployment,
            CallOption::AwaitBlock,
        )
        .await?;
        trade_accounts.push(trade_account);
    }

    Ok(trade_accounts)
}

pub async fn fund_trade_accounts<W: Account + Clone>(
    trade_accounts: &[TradeAccountManager<W>],
    base_asset: AssetId,
    base_asset_amount: u64,
    quote_asset: AssetId,
    quote_asset_amount: u64,
) -> anyhow::Result<(), anyhow::Error> {
    for trade_account in trade_accounts.iter() {
        let _ = trade_account
            .owner
            .force_transfer_to_contract(
                trade_account.contract.contract_id(),
                base_asset_amount,
                base_asset,
                TxPolicies::default(),
            )
            .await?;
        let _ = trade_account
            .owner
            .force_transfer_to_contract(
                trade_account.contract.contract_id(),
                quote_asset_amount,
                quote_asset,
                TxPolicies::default(),
            )
            .await?;
    }
    Ok(())
}

pub async fn create_order_call(
    order_book: &OrderBookManager<Wallet>,
    order_data: &OrderData,
    order_type: OrderType,
    trade_account: &mut TradeAccountManager<Wallet>,
    gas_per_method: Option<u64>,
) -> anyhow::Result<CallContractArgs> {
    let create_order_params = CreateOrderParams {
        price: order_data.price,
        quantity: order_data.quantity,
        side: order_data.side,
        asset_id: order_book.get_order_side_asset(&order_data.side),
        order_type,
    };
    // Create orders args with signatures
    let contract_call_args = trade_account
        .create_order(order_book, &create_order_params, gas_per_method)
        .await?;

    Ok(contract_call_args)
}

pub async fn create_order_handlers<'a, I>(
    order_book: &OrderBookManager<Wallet>,
    orders: &[OrderData],
    trade_accounts: I,
    gas_per_method: Option<u64>,
) -> anyhow::Result<
    Vec<CallHandler<Wallet, fuels::programs::calls::ContractCall, ()>>,
    anyhow::Error,
>
where
    I: Iterator<Item = &'a mut TradeAccountManager<Wallet>>,
{
    let mut create_orders_handlers = Vec::with_capacity(orders.len());
    let mut trade_accounts = trade_accounts.into_iter().collect::<Vec<_>>();
    for order_data in orders.iter() {
        let trade_account = trade_accounts
            .iter_mut()
            .find(|ta| ta.identity() == order_data.trader_id)
            .unwrap();

        // Create orders args with signatures
        let contract_call_args = create_order_call(
            order_book,
            order_data,
            OrderType::Spot,
            trade_account,
            gas_per_method,
        )
        .await?;

        // Create function handler
        create_orders_handlers.push(
            trade_account
                .session_call_contract(&contract_call_args)
                .with_contract_ids(&[order_book.contract.contract_id()]),
        );
    }
    Ok(create_orders_handlers)
}

pub async fn cancel_order_call(
    order_book: &OrderBookManager<Wallet>,
    order_id: &Bytes32,
    trade_account: &mut TradeAccountManager<Wallet>,
    gas_per_method: Option<u64>,
) -> anyhow::Result<CallContractArgs> {
    trade_account
        .cancel_order(order_book, *order_id, gas_per_method)
        .await
}

pub async fn cancel_order_handlers(
    order_book: &OrderBookManager<Wallet>,
    orders: &[Bytes32],
    trade_account: &mut TradeAccountManager<Wallet>,
    gas_per_method: Option<u64>,
) -> anyhow::Result<
    Vec<CallHandler<Wallet, fuels::programs::calls::ContractCall, ()>>,
    anyhow::Error,
> {
    let mut cancel_orders_handlers = Vec::with_capacity(orders.len());
    for order_id in orders.iter() {
        let contract_call_args =
            cancel_order_call(order_book, order_id, trade_account, gas_per_method)
                .await?;

        // Create function handler
        cancel_orders_handlers
            .push(trade_account.session_call_contract(&contract_call_args));
    }
    Ok(cancel_orders_handlers)
}

pub async fn send_transactions(
    create_orders_handlers: Vec<
        CallHandler<Wallet, fuels::programs::calls::ContractCall, ()>,
    >,
    gaspayer_wallet: Wallet,
    call_option: CallOption,
) -> Vec<Bytes32> {
    let mut order_results = Vec::with_capacity(create_orders_handlers.len());
    for mut call_handler in create_orders_handlers {
        call_handler.account = gaspayer_wallet.clone();

        let tx_id = match call_option.clone() {
            CallOption::AwaitBlock => call_handler.submit().await.unwrap().tx_id(),
            CallOption::AwaitPreconfirmation(ops) => {
                call_handler
                    .almost_sync_call(
                        &ops.data_builder,
                        &ops.utxo_manager,
                        &ops.tx_config,
                    )
                    .await
                    .unwrap()
                    .tx_id
            }
        };
        order_results.push(tx_id);
    }
    order_results
}

#[derive(Debug, Clone, Default)]
pub struct OrderBookEvents {
    pub matches: Vec<OrderMatchedEvent>,
    pub orders: Vec<OrderCreatedEvent>,
    pub cancels: Vec<OrderCancelledEvent>,
}

pub fn get_order_book_events<W: Account + Clone>(
    order_book: &OrderBookManager<W>,
    order_book_events: &mut OrderBookEvents,
    tx_result: &TxStatus,
) -> anyhow::Result<(), anyhow::Error> {
    if let TxStatus::Success(success) = tx_result {
        let mut order_created_events =
            order_book
                .contract
                .log_decoder()
                .decode_logs_with_type::<OrderCreatedEvent>(&success.receipts)?;
        let mut order_match_events =
            order_book
                .contract
                .log_decoder()
                .decode_logs_with_type::<OrderMatchedEvent>(&success.receipts)?;
        let mut order_cancel_events =
            order_book
                .contract
                .log_decoder()
                .decode_logs_with_type::<OrderCancelledEvent>(&success.receipts)?;

        order_book_events.orders.append(&mut order_created_events);
        order_book_events.matches.append(&mut order_match_events);
        order_book_events.cancels.append(&mut order_cancel_events);
    }
    Ok(())
}

pub async fn get_wallets_balances<W: Account + Clone>(
    wallets: &[W],
    base_asset: AssetId,
    quote_asset: AssetId,
) -> anyhow::Result<HashMap<Address, (u128, u128)>, anyhow::Error> {
    let mut balances = HashMap::new();
    for wallet in wallets {
        let balance = wallet.get_balances().await?;
        let base_asset_balance =
            balance.get(&base_asset.to_string()).cloned().unwrap_or(0);
        let quote_asset_balance =
            balance.get(&quote_asset.to_string()).cloned().unwrap_or(0);
        balances.insert(wallet.address(), (base_asset_balance, quote_asset_balance));
    }
    Ok(balances)
}

pub async fn get_contracts_balances(
    provider: &Provider,
    contracts: &[ContractId],
    base_asset: &AssetId,
    quote_asset: &AssetId,
) -> anyhow::Result<Balances, anyhow::Error> {
    let mut balances = Balances::new();
    for contract_id in contracts {
        let balance = provider.get_contract_balances(contract_id).await?;
        let base_asset_balance = balance.get(&base_asset.clone()).cloned().unwrap_or(0);
        let quote_asset_balance = balance.get(&quote_asset.clone()).cloned().unwrap_or(0);
        balances.insert(
            Identity::ContractId(*contract_id),
            (base_asset_balance, quote_asset_balance),
        );
    }
    Ok(balances)
}

pub async fn settle_trade_accounts_balances<'a, I>(
    fee_payer: &Wallet,
    order_book: &OrderBookManager<Wallet>,
    trade_accounts: I,
    call_option: CallOption,
) -> anyhow::Result<CallResponse<()>, anyhow::Error>
where
    I: Iterator<Item = &'a TradeAccountManager<Wallet>>,
{
    let mut accounts = vec![];
    let mut contracts = vec![];
    for account in trade_accounts {
        accounts.push(Identity::from(account.contract.contract_id()));
        contracts.push(account.contract.contract_id());
    }

    let mut call_handler = order_book
        .contract
        .methods()
        .settle_balances(accounts)
        .with_contract_ids(&contracts);
    call_handler.account = fee_payer.clone();

    let result = match call_option {
        CallOption::AwaitBlock => call_handler.call().await?,
        CallOption::AwaitPreconfirmation(ops) => {
            call_handler
                .almost_sync_call(&ops.data_builder, &ops.utxo_manager, &ops.tx_config)
                .await?
                .tx_status?
        }
    };

    Ok(result)
}

pub async fn get_trade_accounts_balances<W: Account + Clone>(
    provider: &Provider,
    trade_accounts: &[TradeAccountManager<W>],
    base_asset: &AssetId,
    quote_asset: &AssetId,
) -> anyhow::Result<Balances, anyhow::Error> {
    let contracts = trade_accounts
        .iter()
        .map(|trade_account| trade_account.contract.contract_id())
        .collect::<Vec<_>>();
    get_contracts_balances(provider, &contracts, base_asset, quote_asset).await
}

pub async fn wait_for_book_events<W: Account + Clone>(
    tx_ids: &[Bytes32],
    trade_account: &TradeAccountManager<W>,
    order_book: &OrderBookManager<W>,
    gaspayer_wallet: W,
) -> anyhow::Result<OrderBookEvents, anyhow::Error> {
    let mut order_book_events = OrderBookEvents::default();
    let mut tx_completed = HashSet::new();

    while tx_completed.len() != tx_ids.len() {
        let provider = gaspayer_wallet.try_provider()?;
        for order_result in tx_ids.iter() {
            let result = provider.get_transaction_by_id(order_result).await?;

            if let Some(result) = result {
                get_order_book_events(
                    order_book,
                    &mut order_book_events,
                    &result.status,
                )?;
                match result.status {
                    TxStatus::Success(_) => {
                        tx_completed.insert(*order_result);
                    }
                    TxStatus::Failure(failure) => {
                        let logs = order_book
                            .contract
                            .log_decoder()
                            .decode_logs(&failure.receipts);
                        let logs_trade = trade_account
                            .contract
                            .log_decoder()
                            .decode_logs(&failure.receipts);
                        println!("{logs:#?}");
                        println!("{logs_trade:#?}");
                        panic!("{:#}", failure.reason);
                    }
                    _ => {
                        continue;
                    }
                }
            }
        }
    }

    Ok(order_book_events)
}

pub fn get_asset_balance(balances: &HashMap<String, u128>, asset_id: &AssetId) -> u128 {
    *balances.get(&asset_id.to_string()).unwrap_or(&0)
}

pub fn get_total_amount(quantity: u64, price: u64, decimals: u64) -> u64 {
    (quantity * price) / 10u64.pow(decimals as u32)
}