Add CLI argument to show locked balance and utxos

node-gui/src/backend/backend_impl.rs

@@ -31,6 +31,7 @@ use wallet::{
     DefaultWallet,
 };
 use wallet_controller::{HandlesController, UtxoState, WalletHandlesClient};
+use wallet_types::with_locked::WithLocked;
 
 use super::{
     chainstate_event_handler::ChainstateEventHandler,
@@ -411,7 +412,11 @@ impl Backend {
         account_index: U31,
     ) -> BTreeMap<Currency, Amount> {
         controller
-            .get_balance(account_index, UtxoState::Confirmed.into())
+            .get_balance(
+                account_index,
+                UtxoState::Confirmed.into(),
+                WithLocked::Unlocked,
+            )
             .expect("get_balance should not fail normally")
     }
 

wallet/src/account/mod.rs

@@ -26,6 +26,7 @@ use common::Uint256;
 use crypto::key::hdkd::child_number::ChildNumber;
 use mempool::FeeRate;
 pub use utxo_selector::UtxoSelectorError;
+use wallet_types::with_locked::WithLocked;
 
 use crate::account::utxo_selector::{select_coins, OutputGroup};
 use crate::key_chain::{make_path_to_vrf_key, AccountKeyChain, KeyChainError};
@@ -167,6 +168,7 @@ impl Account {
                 UtxoType::Transfer | UtxoType::LockThenTransfer,
                 median_time,
                 UtxoState::Confirmed | UtxoState::InMempool | UtxoState::Inactive,
+                WithLocked::Unlocked,
             )
             .into_iter(),
             |(_, (tx_output, _))| tx_output,
@@ -549,6 +551,7 @@ impl Account {
             UtxoType::CreateStakePool | UtxoType::ProduceBlockFromStake,
             median_time,
             UtxoState::Confirmed.into(),
+            WithLocked::Unlocked,
         );
         // TODO: Select by pool_id if there is more than one UTXO
         let (kernel_input_outpoint, (kernel_input_utxo, _token_id)) =
@@ -748,9 +751,10 @@ impl Account {
         utxo_types: UtxoTypes,
         utxo_states: UtxoStates,
         median_time: BlockTimestamp,
+        with_locked: WithLocked,
     ) -> WalletResult<BTreeMap<Currency, Amount>> {
         let amounts_by_currency = group_utxos_for_input(
-            self.get_utxos(utxo_types, median_time, utxo_states).into_iter(),
+            self.get_utxos(utxo_types, median_time, utxo_states, with_locked).into_iter(),
             |(_, (tx_output, _))| tx_output,
             |total: &mut Amount, _, amount| -> WalletResult<()> {
                 *total = (*total + amount).ok_or(WalletError::OutputAmountOverflow)?;
@@ -767,13 +771,15 @@ impl Account {
         utxo_types: UtxoTypes,
         median_time: BlockTimestamp,
         utxo_states: UtxoStates,
+        with_locked: WithLocked,
     ) -> BTreeMap<UtxoOutPoint, (&TxOutput, Option<TokenId>)> {
         let current_block_info = BlockInfo {
             height: self.account_info.best_block_height(),
             timestamp: median_time,
         };
         let mut all_outputs =
-            self.output_cache.utxos_with_token_ids(current_block_info, utxo_states);
+            self.output_cache
+                .utxos_with_token_ids(current_block_info, utxo_states, with_locked);
         all_outputs.retain(|_outpoint, (txo, _token_id)| {
             self.is_mine_or_watched(txo) && utxo_types.contains(get_utxo_type(txo))
         });

wallet/src/account/output_cache/mod.rs

@@ -30,6 +30,7 @@ use utils::ensure;
 use wallet_types::{
     utxo_types::{get_utxo_state, UtxoStates},
     wallet_tx::TxState,
+    with_locked::WithLocked,
     AccountWalletTxId, BlockInfo, WalletTx,
 };
 
@@ -202,13 +203,8 @@ impl OutputCache {
             }
         }
 
-        let tx_block_info = match tx.state() {
-            TxState::Confirmed(height, timestamp) => Some(BlockInfo { height, timestamp }),
-            TxState::Inactive
-            | TxState::Conflicted(_)
-            | TxState::InMempool
-            | TxState::Abandoned => None,
-        };
+        let tx_block_info = get_block_info(&tx);
+
         if let Some(block_info) = tx_block_info {
             for (idx, output) in tx.outputs().iter().enumerate() {
                 match output {
@@ -298,18 +294,6 @@ impl OutputCache {
         Ok(())
     }
 
-    fn valid_utxo(
-        &self,
-        outpoint: &UtxoOutPoint,
-        output: &TxOutput,
-        transaction_block_info: &Option<BlockInfo>,
-        current_block_info: &BlockInfo,
-        utxo_states: UtxoStates,
-    ) -> bool {
-        !self.is_consumed(utxo_states, outpoint)
-            && valid_timelock(output, current_block_info, transaction_block_info, outpoint)
-    }
-
     fn is_consumed(&self, utxo_states: UtxoStates, outpoint: &UtxoOutPoint) -> bool {
         self.consumed.get(outpoint).map_or(false, |consumed_state| {
             utxo_states.contains(get_utxo_state(consumed_state))
@@ -320,44 +304,47 @@ impl OutputCache {
         &self,
         current_block_info: BlockInfo,
         utxo_states: UtxoStates,
+        locked_state: WithLocked,
     ) -> BTreeMap<UtxoOutPoint, (&TxOutput, Option<TokenId>)> {
-        let mut utxos = BTreeMap::new();
-
-        for tx in self.txs.values() {
-            if !utxo_states.contains(get_utxo_state(&tx.state())) {
-                continue;
-            }
-
-            let tx_block_info = match tx.state() {
-                TxState::Confirmed(height, timestamp) => Some(BlockInfo { height, timestamp }),
-                TxState::InMempool
-                | TxState::Inactive
-                | TxState::Conflicted(_)
-                | TxState::Abandoned => None,
-            };
-            for (index, output) in tx.outputs().iter().enumerate() {
-                let outpoint = UtxoOutPoint::new(tx.id(), index as u32);
-                if self.valid_utxo(
-                    &outpoint,
-                    output,
-                    &tx_block_info,
-                    &current_block_info,
-                    utxo_states,
-                ) {
-                    let token_id = if output.is_token_or_nft_issuance() {
-                        match tx {
-                            WalletTx::Tx(tx_data) => token_id(tx_data.get_transaction()),
-                            WalletTx::Block(_) => None,
-                        }
-                    } else {
-                        None
-                    };
-                    utxos.insert(outpoint, (output, token_id));
-                }
-            }
-        }
+        self.txs
+            .values()
+            .filter(|tx| is_in_state(tx, utxo_states))
+            .flat_map(|tx| {
+                let tx_block_info = get_block_info(tx);
+                let token_id = match tx {
+                    WalletTx::Tx(tx_data) => token_id(tx_data.get_transaction()),
+                    WalletTx::Block(_) => None,
+                };
 
-        utxos
+                tx.outputs()
+                    .iter()
+                    .enumerate()
+                    .map(|(idx, output)| (output, UtxoOutPoint::new(tx.id(), idx as u32)))
+                    .filter(move |(output, outpoint)| {
+                        !self.is_consumed(utxo_states, outpoint)
+                            && is_specific_lock_state(
+                                locked_state,
+                                output,
+                                current_block_info,
+                                tx_block_info,
+                                outpoint,
+                            )
+                    })
+                    .map(move |(output, outpoint)| {
+                        (
+                            outpoint,
+                            (
+                                output,
+                                if output.is_token_or_nft_issuance() {
+                                    token_id
+                                } else {
+                                    None
+                                },
+                            ),
+                        )
+                    })
+            })
+            .collect()
     }
 
     pub fn pending_transactions(&self) -> Vec<&WithId<Transaction>> {
@@ -437,6 +424,40 @@ impl OutputCache {
     }
 }
 
+/// Checks the output against the current block height and compares it with the locked_state parameter.
+/// If they match, the function return true, if they don't, it returns false.
+/// For example, if we would like to check that an output is locked,
+/// we pass locked_state = WithLocked::Locked, and pass the output in question.
+/// If the output is locked, the function returns true. Otherwise, it returns false
+fn is_specific_lock_state(
+    locked_state: WithLocked,
+    output: &TxOutput,
+    current_block_info: BlockInfo,
+    tx_block_info: Option<BlockInfo>,
+    outpoint: &UtxoOutPoint,
+) -> bool {
+    match locked_state {
+        WithLocked::Any => true,
+        WithLocked::Locked => {
+            !valid_timelock(output, &current_block_info, &tx_block_info, outpoint)
+        }
+        WithLocked::Unlocked => {
+            valid_timelock(output, &current_block_info, &tx_block_info, outpoint)
+        }
+    }
+}
+
+/// Get the block info (block height and timestamp) if the Tx is in confirmed state
+fn get_block_info(tx: &WalletTx) -> Option<BlockInfo> {
+    match tx.state() {
+        TxState::Confirmed(height, timestamp) => Some(BlockInfo { height, timestamp }),
+        TxState::InMempool | TxState::Inactive | TxState::Conflicted(_) | TxState::Abandoned => {
+            None
+        }
+    }
+}
+
+/// Check the TxOutput's timelock is unlocked
 fn valid_timelock(
     output: &TxOutput,
     current_block_info: &BlockInfo,
@@ -457,3 +478,8 @@ fn valid_timelock(
         })
     })
 }
+
+/// Check Tx is in the selected state Confirmed/Inactive/Abandoned...
+fn is_in_state(tx: &WalletTx, utxo_states: UtxoStates) -> bool {
+    utxo_states.contains(get_utxo_state(&tx.state()))
+}

wallet/src/wallet/mod.rs

@@ -53,6 +53,7 @@ use wallet_storage::{
 use wallet_storage::{StoreTxRwUnlocked, TransactionRwUnlocked};
 use wallet_types::utxo_types::{UtxoStates, UtxoTypes};
 use wallet_types::wallet_tx::TxState;
+use wallet_types::with_locked::WithLocked;
 use wallet_types::{AccountId, BlockInfo, KeyPurpose};
 
 pub const WALLET_VERSION_UNINITIALIZED: u32 = 0;
@@ -356,11 +357,13 @@ impl<B: storage::Backend> Wallet<B> {
         account_index: U31,
         utxo_types: UtxoTypes,
         utxo_states: UtxoStates,
+        with_locked: WithLocked,
     ) -> WalletResult<BTreeMap<Currency, Amount>> {
         self.get_account(account_index)?.get_balance(
             utxo_types,
             utxo_states,
             self.latest_median_time,
+            with_locked,
         )
     }
 
@@ -369,9 +372,15 @@ impl<B: storage::Backend> Wallet<B> {
         account_index: U31,
         utxo_types: UtxoTypes,
         utxo_states: UtxoStates,
+        with_locked: WithLocked,
     ) -> WalletResult<BTreeMap<UtxoOutPoint, TxOutput>> {
         let account = self.get_account(account_index)?;
-        let utxos = account.get_utxos(utxo_types, self.latest_median_time, utxo_states);
+        let utxos = account.get_utxos(
+            utxo_types,
+            self.latest_median_time,
+            utxo_states,
+            with_locked,
+        );
         let utxos = utxos
             .into_iter()
             .map(|(outpoint, (txo, _token_id))| (outpoint, txo.clone()))