o2_tools/

utxo_manager.rs

use fuel_core_types::{
    fuel_tx::UtxoId,
    fuel_types::{
        Address,
        AssetId,
    },
};
use fuels::types::{
    coin::Coin,
    coin_type::CoinType,
    input::Input,
};
use std::{
    collections::{
        BTreeSet,
        HashMap,
        HashSet,
        hash_map::Entry,
    },
    sync::Arc,
};

pub struct CoinsResult {
    pub known_coins: Vec<FuelTxCoin>,
    pub unknown_coins: HashSet<UtxoId>,
}

#[derive(Copy, Clone, Debug, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
pub struct FuelTxCoin {
    pub amount: u64,
    pub asset_id: AssetId,
    pub utxo_id: UtxoId,
    pub owner: Address,
}

impl From<Coin> for FuelTxCoin {
    fn from(value: Coin) -> Self {
        Self {
            amount: value.amount,
            asset_id: value.asset_id,
            utxo_id: value.utxo_id,
            owner: value.owner,
        }
    }
}

impl From<FuelTxCoin> for Coin {
    fn from(value: FuelTxCoin) -> Self {
        Self {
            amount: value.amount,
            asset_id: value.asset_id,
            utxo_id: value.utxo_id,
            owner: value.owner,
        }
    }
}

impl TryFrom<&fuel_core_types::fuel_tx::Input> for FuelTxCoin {
    type Error = anyhow::Error;

    fn try_from(input: &fuel_core_types::fuel_tx::Input) -> Result<Self, Self::Error> {
        if let fuel_core_types::fuel_tx::Input::CoinSigned(coin) = input {
            return Ok(FuelTxCoin {
                utxo_id: coin.utxo_id,
                owner: coin.owner,
                amount: coin.amount,
                asset_id: coin.asset_id,
            })
        }
        anyhow::bail!("Invalid input type")
    }
}

impl From<FuelTxCoin> for Input {
    fn from(value: FuelTxCoin) -> Self {
        Input::resource_signed(CoinType::Coin(value.into()))
    }
}

impl Ord for FuelTxCoin {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        self.amount
            .cmp(&other.amount)
            .then_with(|| self.asset_id.cmp(&other.asset_id))
            .then_with(|| self.utxo_id.cmp(&other.utxo_id))
            .then_with(|| self.owner.cmp(&other.owner))
    }
}

impl PartialOrd for FuelTxCoin {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(self.cmp(other))
    }
}

/// Trait for types that can provide UTXO management capabilities.
/// This allows decoupling the concrete UtxoManager implementation from the
/// o2-tools crate, enabling consumers to provide their own implementation.
pub trait UtxoProvider: Send + 'static {
    fn balance_of(&self, owner: Address, asset_id: AssetId) -> u128;
    fn len_per_address(&self, asset_id: AssetId, address: Address) -> usize;
    fn guaranteed_extract_coins(
        &mut self,
        owner: Address,
        asset_id: AssetId,
        amount: u128,
    ) -> anyhow::Result<Vec<FuelTxCoin>>;
    fn load_from_coins_vec(&mut self, coins: Vec<FuelTxCoin>);
    fn number_of_coins_with_amount_greater_or_equal(
        &self,
        owner: Address,
        asset_id: AssetId,
        amount: u128,
    ) -> (u128, usize);
    fn extract_largest_coins(
        &mut self,
        owner: Address,
        asset_id: AssetId,
        max_value: u128,
    ) -> Vec<FuelTxCoin>;
    fn coin_count(&self) -> usize;
    fn utxo_ids(&self) -> Vec<UtxoId>;
    fn remove_coin(&mut self, utxo_id: &UtxoId) -> bool;
    fn total_balance(&self, asset_id: &AssetId) -> u128;
}

pub type SharedUtxoManager = Arc<tokio::sync::Mutex<dyn UtxoProvider>>;

/// A struct to manage UTXOs of internal accounts used to pay fee, match transactions, and
/// users predicate based accounts.
pub struct UtxoManager {
    account_utxos: HashMap<(Address, AssetId), BTreeSet<FuelTxCoin>>,
    coins: HashMap<UtxoId, FuelTxCoin>,
}

impl Default for UtxoManager {
    fn default() -> Self {
        Self::new()
    }
}

impl UtxoManager {
    pub fn new() -> Self {
        Self {
            account_utxos: HashMap::new(),
            coins: HashMap::new(),
        }
    }

    pub fn len_per_address(&self, asset_id: AssetId, address: Address) -> usize {
        self.account_utxos
            .get(&(address, asset_id))
            .map_or(0, |utxos| utxos.len())
    }

    pub fn new_from_coins<I>(coins: I) -> Self
    where
        I: Iterator<Item = FuelTxCoin>,
    {
        let mut _self = Self::new();

        _self.load_from_coins(coins);

        _self
    }

    pub fn load_from_coins<I>(&mut self, coins: I)
    where
        I: Iterator<Item = FuelTxCoin>,
    {
        for coin in coins {
            if coin.amount == 0 {
                continue;
            }

            let key = (coin.owner, coin.asset_id);
            self.account_utxos.entry(key).or_default().insert(coin);
            self.coins.insert(coin.utxo_id, coin);
        }
    }

    fn extract_utxos(&mut self, utxos: &[UtxoId]) -> anyhow::Result<Vec<FuelTxCoin>> {
        let mut coins = vec![];

        for utxo_id in utxos {
            let coin = self.coins.remove(utxo_id).ok_or_else(|| {
                anyhow::anyhow!("UTXO {utxo_id} not found in the UTXO manager")
            })?;

            let key = (coin.owner, coin.asset_id);
            let account = self.account_utxos.entry(key);

            match account {
                Entry::Occupied(mut occupied) => {
                    occupied.get_mut().remove(&coin);

                    if occupied.get().is_empty() {
                        occupied.remove();
                    }

                    coins.push(coin);
                }
                Entry::Vacant(_) => {}
            }
        }

        Ok(coins)
    }

    fn utxos_for(
        &self,
        owner: Address,
        asset_id: AssetId,
        amount: u128,
        fail_if_not_enough: bool,
    ) -> anyhow::Result<Vec<UtxoId>> {
        let Some(coins) = self.account_utxos.get(&(owner, asset_id)) else {
            return Err(anyhow::anyhow!(
                "No UTXOs found for the given {owner} and {asset_id}"
            ));
        };

        let mut total_amount = 0;
        let mut utxos_to_remove = vec![];

        // Iterate through entries (already sorted by amount ascending)
        for coin in coins.iter() {
            if total_amount >= amount {
                break;
            }

            utxos_to_remove.push(coin.utxo_id);
            total_amount += coin.amount as u128;
        }

        if fail_if_not_enough && total_amount < amount {
            return Err(anyhow::anyhow!(
                "Not enough UTXOs({total_amount}) found \
                for the given {owner} and {asset_id} to cover {amount}."
            ));
        }

        Ok(utxos_to_remove)
    }

    pub fn guaranteed_extract_coins(
        &mut self,
        owner: Address,
        asset_id: AssetId,
        amount: u128,
    ) -> anyhow::Result<Vec<FuelTxCoin>> {
        let utxos = self.utxos_for(owner, asset_id, amount, true)?;
        self.extract_utxos(&utxos)
    }

    pub fn number_of_coins_with_amount_greater_or_equal(
        &self,
        owner: Address,
        asset_id: AssetId,
        amount: u128,
    ) -> (u128, usize) {
        self.account_utxos
            .get(&(owner, asset_id))
            .map_or((0, 0), |coins| {
                let mut count = 0;
                let mut total_balance = 0;

                for coin in coins.iter() {
                    if coin.amount as u128 >= amount {
                        count += 1;
                        total_balance += coin.amount as u128;
                    }
                }

                (total_balance, count)
            })
    }

    pub fn balance_of(&self, owner: Address, asset_id: AssetId) -> u128 {
        self.account_utxos
            .get(&(owner, asset_id))
            .map_or(0, |coins| {
                coins.iter().map(|coin| coin.amount as u128).sum()
            })
    }

    /// Returns a reference to all tracked coins.
    pub fn coins(&self) -> &HashMap<UtxoId, FuelTxCoin> {
        &self.coins
    }

    /// Returns the number of coins tracked.
    pub fn coin_count(&self) -> usize {
        self.coins.len()
    }

    /// Returns the UtxoIds of all tracked coins.
    pub fn utxo_ids(&self) -> Vec<UtxoId> {
        self.coins.keys().copied().collect()
    }

    /// Checks if a specific coin is tracked.
    pub fn contains(&self, utxo_id: &UtxoId) -> bool {
        self.coins.contains_key(utxo_id)
    }

    /// Returns the total balance for a given asset across all owners.
    pub fn total_balance(&self, asset_id: &AssetId) -> u128 {
        self.coins
            .values()
            .filter(|coin| &coin.asset_id == asset_id)
            .map(|coin| coin.amount as u128)
            .sum()
    }

    /// Removes a specific coin by UtxoId. Returns true if it was found and removed.
    pub fn remove_coin(&mut self, utxo_id: &UtxoId) -> bool {
        if let Some(coin) = self.coins.remove(utxo_id) {
            let key = (coin.owner, coin.asset_id);
            if let Entry::Occupied(mut entry) = self.account_utxos.entry(key) {
                entry.get_mut().remove(&coin);
                if entry.get().is_empty() {
                    entry.remove();
                }
            }
            true
        } else {
            false
        }
    }

    /// Extracts the largest coins for the given owner/asset up to
    /// `max_value` total. Iterates coins in descending order of amount.
    /// Returns the extracted coins (removed from the manager).
    pub fn extract_largest_coins(
        &mut self,
        owner: Address,
        asset_id: AssetId,
        max_value: u128,
    ) -> Vec<FuelTxCoin> {
        let utxo_ids: Vec<UtxoId> = {
            let Some(coins) = self.account_utxos.get(&(owner, asset_id)) else {
                return vec![];
            };
            let mut total = 0u128;
            coins
                .iter()
                .rev() // largest first
                .take_while(|coin| {
                    if total >= max_value {
                        return false;
                    }
                    total += coin.amount as u128;
                    true
                })
                .map(|coin| coin.utxo_id)
                .collect()
        };
        self.extract_utxos(&utxo_ids).unwrap_or_default()
    }
}

impl UtxoProvider for UtxoManager {
    fn balance_of(&self, owner: Address, asset_id: AssetId) -> u128 {
        UtxoManager::balance_of(self, owner, asset_id)
    }

    fn len_per_address(&self, asset_id: AssetId, address: Address) -> usize {
        UtxoManager::len_per_address(self, asset_id, address)
    }

    fn guaranteed_extract_coins(
        &mut self,
        owner: Address,
        asset_id: AssetId,
        amount: u128,
    ) -> anyhow::Result<Vec<FuelTxCoin>> {
        UtxoManager::guaranteed_extract_coins(self, owner, asset_id, amount)
    }

    fn load_from_coins_vec(&mut self, coins: Vec<FuelTxCoin>) {
        self.load_from_coins(coins.into_iter());
    }

    fn number_of_coins_with_amount_greater_or_equal(
        &self,
        owner: Address,
        asset_id: AssetId,
        amount: u128,
    ) -> (u128, usize) {
        UtxoManager::number_of_coins_with_amount_greater_or_equal(
            self, owner, asset_id, amount,
        )
    }

    fn extract_largest_coins(
        &mut self,
        owner: Address,
        asset_id: AssetId,
        max_value: u128,
    ) -> Vec<FuelTxCoin> {
        UtxoManager::extract_largest_coins(self, owner, asset_id, max_value)
    }

    fn coin_count(&self) -> usize {
        UtxoManager::coin_count(self)
    }

    fn utxo_ids(&self) -> Vec<UtxoId> {
        UtxoManager::utxo_ids(self)
    }

    fn remove_coin(&mut self, utxo_id: &UtxoId) -> bool {
        UtxoManager::remove_coin(self, utxo_id)
    }

    fn total_balance(&self, asset_id: &AssetId) -> u128 {
        UtxoManager::total_balance(self, asset_id)
    }
}

#[cfg(test)]
#[allow(non_snake_case)]
mod tests {
    use super::*;

    #[test]
    fn guaranteed_extract_coins__returns_coins_in_ascending_order_by_amount__when_coins_inserted_in_random_order()
     {
        // given
        let owner = Address::from([1u8; 32]);
        let asset_id = AssetId::from([2u8; 32]);

        let coin1 = FuelTxCoin {
            amount: 100,
            asset_id,
            utxo_id: UtxoId::new(fuel_core_types::fuel_tx::TxId::from([1u8; 32]), 0),
            owner,
        };
        let coin2 = FuelTxCoin {
            amount: 50,
            asset_id,
            utxo_id: UtxoId::new(fuel_core_types::fuel_tx::TxId::from([2u8; 32]), 0),
            owner,
        };
        let coin3 = FuelTxCoin {
            amount: 200,
            asset_id,
            utxo_id: UtxoId::new(fuel_core_types::fuel_tx::TxId::from([3u8; 32]), 0),
            owner,
        };
        let coin4 = FuelTxCoin {
            amount: 75,
            asset_id,
            utxo_id: UtxoId::new(fuel_core_types::fuel_tx::TxId::from([4u8; 32]), 0),
            owner,
        };

        let mut manager = UtxoManager::new();
        manager.load_from_coins(vec![coin1, coin2, coin3, coin4].into_iter());

        // when
        let total_amount = 100 + 50 + 200 + 75;
        let extracted = manager
            .guaranteed_extract_coins(owner, asset_id, total_amount)
            .unwrap();

        // then
        assert_eq!(extracted.len(), 4);
        assert_eq!(extracted[0].amount, 50); // coin2 (smallest)
        assert_eq!(extracted[1].amount, 75); // coin4
        assert_eq!(extracted[2].amount, 100); // coin1
        assert_eq!(extracted[3].amount, 200); // coin3 (largest)
        assert_eq!(manager.balance_of(owner, asset_id), 0);
    }

    #[test]
    fn extract_largest_coins__takes_largest_first_up_to_max_value() {
        let owner = Address::from([1u8; 32]);
        let asset_id = AssetId::from([2u8; 32]);

        let coins: Vec<FuelTxCoin> = (1..=5)
            .map(|i| FuelTxCoin {
                amount: i * 100, // 100, 200, 300, 400, 500
                asset_id,
                utxo_id: UtxoId::new(
                    fuel_core_types::fuel_tx::TxId::from([i as u8; 32]),
                    0,
                ),
                owner,
            })
            .collect();
        let mut manager = UtxoManager::new();
        manager.load_from_coins(coins.into_iter());

        // Extract up to 600 in value (largest first: 500, then 400 -> total 900 > 600)
        // take_while stops after accumulating >= max_value, so we get 500 (total=500 < 600)
        // then 400 (total=900 >= 600, but take_while already entered), so we get 500 + 400.
        // Actually take_while: first 500 → total=500 < 600 → true; next 400 → total=900 → true
        // (check is at start: total >= max_value? 500 >= 600? no, so take it);
        // next iteration: total=900 >= 600? yes → stop.
        let extracted = manager.extract_largest_coins(owner, asset_id, 600);
        let amounts: Vec<u64> = extracted.iter().map(|c| c.amount).collect();
        // Should have extracted the two largest coins (500 and 400)
        assert_eq!(amounts.len(), 2);
        assert!(amounts.contains(&500));
        assert!(amounts.contains(&400));
        // Remaining: 100+200+300 = 600
        assert_eq!(manager.balance_of(owner, asset_id), 600);
    }

    #[test]
    fn extract_largest_coins__returns_empty_when_no_coins() {
        let owner = Address::from([1u8; 32]);
        let asset_id = AssetId::from([2u8; 32]);
        let mut manager = UtxoManager::new();
        let extracted = manager.extract_largest_coins(owner, asset_id, 1000);
        assert!(extracted.is_empty());
    }

    #[test]
    fn extract_largest_coins__extracts_single_coin_when_only_one() {
        let owner = Address::from([1u8; 32]);
        let asset_id = AssetId::from([2u8; 32]);
        let coin = FuelTxCoin {
            amount: 1_000_000,
            asset_id,
            utxo_id: UtxoId::new(fuel_core_types::fuel_tx::TxId::from([1u8; 32]), 0),
            owner,
        };
        let mut manager = UtxoManager::new();
        manager.load_from_coins(vec![coin].into_iter());

        let extracted = manager.extract_largest_coins(owner, asset_id, 500_000);
        assert_eq!(extracted.len(), 1);
        assert_eq!(extracted[0].amount, 1_000_000);
        assert_eq!(manager.balance_of(owner, asset_id), 0);
    }
}