satellite-shard 0.31.6

use arch_program::rune::{RuneAmount, RuneId};
use bitcoin::Amount;
use satellite_bitcoin::generic::fixed_set::FixedCapacitySet;
use satellite_bitcoin::utxo_info::UtxoInfoTrait;

/// A container holding the UTXOs owned by a single program **account**.
///
/// The concrete type parameter `U` is a *fixed-size array* auto-generated by
/// the `satellite_bitcoin::declare_fixed_array!` macro – for example
/// `FixedArrayUtxoInfo25` which can hold up to 25 `UtxoInfo` entries without
/// allocating.  Keeping the collections fixed-size and `#[repr(C)]` ensures
/// the struct has a predictable, byte-compatible layout so it can be stored
/// directly inside on-chain state.
///
/// The `RU` type parameter is a *fixed-option* type for holding rune UTXOs.
#[repr(C, align(8))]
#[derive(Clone, Debug)]
pub struct AccountUtxos<U, RU> {
    /// Bitcoin UTXOs that **do not** contain runes.  The size of the collection
    /// is determined by the generic parameter `U`.
    pub btc_utxos: U,
    /// Optional UTXO that *does* contain runes for the account.  Using
    /// a fixed-option type means the option itself is stored without any
    /// dynamic allocation.
    pub rune_utxo: RU,
}

/// Abstraction over a *shard* of state held by an on-chain program.
///
/// A shard is the minimal unit that can be processed in parallel without
/// causing conflicts.  Implementors usually wrap `AccountUtxos` but may provide
/// additional bookkeeping fields.
pub trait StateShard<
    U: UtxoInfoTrait<RuneSet>,
    RuneSet: FixedCapacitySet<Item = RuneAmount> + Default,
>
{
    /// Immutable view over the shard's plain-BTC UTXOs.
    fn btc_utxos(&self) -> &[U];
    /// Mutable view over the shard's plain-BTC UTXOs.
    fn btc_utxos_mut(&mut self) -> &mut [U];

    /// Retain only those BTC-UTXOs that satisfy the provided predicate `f`.
    ///
    /// Implementations should delegate to the underlying fixed-array's
    /// `retain` logic so the backing storage and length tracking stay in sync.
    fn btc_utxos_retain(&mut self, f: &mut dyn FnMut(&U) -> bool);

    /// Inserts a new BTC-UTXO into the collection.
    ///
    /// Returns the index at which the UTXO was placed or `None` when the
    /// collection is already full.
    fn add_btc_utxo(&mut self, utxo: U) -> Option<usize>;

    /// Current number of BTC-UTXOs stored in the shard.
    fn btc_utxos_len(&self) -> usize;

    /// Max amount of btc utxos
    fn btc_utxos_max_len(&self) -> usize;

    /// Returns the rune-carrying UTXO if one exists.
    fn rune_utxo(&self) -> Option<&U>;
    /// Mutable variant of [`Self::rune_utxo`].
    fn rune_utxo_mut(&mut self) -> Option<&mut U>;

    /// Removes the rune-carrying UTXO, leaving the slot empty.
    fn clear_rune_utxo(&mut self);

    /// Overwrites the current rune-UTXO with `utxo`.
    fn set_rune_utxo(&mut self, utxo: U);

    /// Calculate the total amount of BTC held in this shard's BTC UTXOs
    fn total_btc(&self) -> Amount {
        let sat = self
            .btc_utxos()
            .iter()
            .map(|utxo_info| utxo_info.value())
            .sum();

        Amount::from_sat(sat)
    }

    /// Calculate the total amount of Runes held in this shard's Rune UTXOs
    fn total_rune(&self, rune_id: RuneId) -> u128 {
        #[cfg(feature = "runes")]
        {
            if let Some(utxo_info) = self.rune_utxo() {
                utxo_info
                    .runes()
                    .iter()
                    .find(|rune_amount| rune_amount.id == rune_id)
                    .map_or(0, |rune_amount| rune_amount.amount)
            } else {
                0
            }
        }
        #[cfg(not(feature = "runes"))]
        {
            0
        }
    }
}