From 02dee20cc4d911ca846232a962c5522ecc2cef84 Mon Sep 17 00:00:00 2001 From: Boris Oncev Date: Tue, 1 Aug 2023 21:51:56 +0200 Subject: [PATCH] Add CLI argument to show locked balance and utxos --- node-gui/src/backend/backend_impl.rs | 7 +- wallet/src/account/mod.rs | 10 +- wallet/src/account/output_cache/mod.rs | 136 +++++++++++------- wallet/src/wallet/mod.rs | 11 +- wallet/src/wallet/tests.rs | 85 +++++++++++ wallet/types/src/lib.rs | 1 + wallet/types/src/with_locked.rs | 22 +++ .../src/commands/helper_types.rs | 18 +++ wallet/wallet-cli-lib/src/commands/mod.rs | 16 ++- wallet/wallet-controller/src/lib.rs | 8 +- 10 files changed, 251 insertions(+), 63 deletions(-) create mode 100644 wallet/types/src/with_locked.rs diff --git a/node-gui/src/backend/backend_impl.rs b/node-gui/src/backend/backend_impl.rs index 9d52b7c344..8bbcaac67c 100644 --- a/node-gui/src/backend/backend_impl.rs +++ b/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 { controller - .get_balance(account_index, UtxoState::Confirmed.into()) + .get_balance( + account_index, + UtxoState::Confirmed.into(), + WithLocked::Unlocked, + ) .expect("get_balance should not fail normally") } diff --git a/wallet/src/account/mod.rs b/wallet/src/account/mod.rs index 486ead432c..d9d7219118 100644 --- a/wallet/src/account/mod.rs +++ b/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> { 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)> { 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)) }); diff --git a/wallet/src/account/output_cache/mod.rs b/wallet/src/account/output_cache/mod.rs index da7d73a96f..5b2277d047 100644 --- a/wallet/src/account/output_cache/mod.rs +++ b/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, - 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)> { - 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, - ¤t_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> { @@ -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, + outpoint: &UtxoOutPoint, +) -> bool { + match locked_state { + WithLocked::Any => true, + WithLocked::Locked => { + !valid_timelock(output, ¤t_block_info, &tx_block_info, outpoint) + } + WithLocked::Unlocked => { + valid_timelock(output, ¤t_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 { + 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())) +} diff --git a/wallet/src/wallet/mod.rs b/wallet/src/wallet/mod.rs index 046f5a8b47..9b870eded7 100644 --- a/wallet/src/wallet/mod.rs +++ b/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 Wallet { account_index: U31, utxo_types: UtxoTypes, utxo_states: UtxoStates, + with_locked: WithLocked, ) -> WalletResult> { self.get_account(account_index)?.get_balance( utxo_types, utxo_states, self.latest_median_time, + with_locked, ) } @@ -369,9 +372,15 @@ impl Wallet { account_index: U31, utxo_types: UtxoTypes, utxo_states: UtxoStates, + with_locked: WithLocked, ) -> WalletResult> { 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())) diff --git a/wallet/src/wallet/tests.rs b/wallet/src/wallet/tests.rs index abcf6ef55b..65c3a2a2ba 100644 --- a/wallet/src/wallet/tests.rs +++ b/wallet/src/wallet/tests.rs @@ -165,6 +165,7 @@ fn verify_wallet_balance( DEFAULT_ACCOUNT_INDEX, UtxoType::Transfer | UtxoType::LockThenTransfer, UtxoState::Confirmed.into(), + WithLocked::Unlocked, ) .unwrap() .get(&Currency::Coin) @@ -180,6 +181,7 @@ fn verify_wallet_balance( DEFAULT_ACCOUNT_INDEX, UtxoType::Transfer | UtxoType::LockThenTransfer, UtxoState::Confirmed.into(), + WithLocked::Unlocked, ) .unwrap() .get(&Currency::Coin) @@ -306,6 +308,7 @@ fn locked_wallet_balance_works(#[case] seed: Seed) { DEFAULT_ACCOUNT_INDEX, UtxoType::Transfer | UtxoType::LockThenTransfer, UtxoState::Confirmed.into(), + WithLocked::Unlocked, ) .unwrap() .get(&Currency::Coin) @@ -322,6 +325,7 @@ fn locked_wallet_balance_works(#[case] seed: Seed) { DEFAULT_ACCOUNT_INDEX, UtxoType::Transfer | UtxoType::LockThenTransfer, UtxoState::Confirmed.into(), + WithLocked::Unlocked, ) .unwrap() .get(&Currency::Coin) @@ -342,6 +346,7 @@ fn wallet_balance_block_reward() { DEFAULT_ACCOUNT_INDEX, UtxoType::Transfer | UtxoType::LockThenTransfer, UtxoState::Confirmed.into(), + WithLocked::Unlocked, ) .unwrap() .get(&Currency::Coin) @@ -655,6 +660,7 @@ fn locked_wallet_cant_sign_transaction(#[case] seed: Seed) { DEFAULT_ACCOUNT_INDEX, UtxoType::Transfer | UtxoType::LockThenTransfer, UtxoState::Confirmed.into(), + WithLocked::Unlocked, ) .unwrap() .get(&Currency::Coin) @@ -696,6 +702,7 @@ fn locked_wallet_cant_sign_transaction(#[case] seed: Seed) { DEFAULT_ACCOUNT_INDEX, UtxoType::Transfer | UtxoType::LockThenTransfer, UtxoState::Confirmed.into(), + WithLocked::Unlocked, ) .unwrap() .get(&Currency::Coin) @@ -778,6 +785,7 @@ fn wallet_transaction_with_fees(#[case] seed: Seed) { DEFAULT_ACCOUNT_INDEX, UtxoType::Transfer | UtxoType::LockThenTransfer, UtxoState::Confirmed.into(), + WithLocked::Unlocked, ) .unwrap() .get(&Currency::Coin) @@ -815,6 +823,7 @@ fn wallet_transaction_with_fees(#[case] seed: Seed) { DEFAULT_ACCOUNT_INDEX, UtxoType::Transfer | UtxoType::LockThenTransfer, UtxoState::Confirmed.into(), + WithLocked::Unlocked, ) .unwrap() .get(&Currency::Coin) @@ -845,6 +854,7 @@ fn wallet_transaction_with_fees(#[case] seed: Seed) { DEFAULT_ACCOUNT_INDEX, UtxoType::Transfer | UtxoType::LockThenTransfer, UtxoState::Confirmed | UtxoState::Inactive, + WithLocked::Unlocked, ) .unwrap() .get(&Currency::Coin) @@ -883,6 +893,7 @@ fn create_stake_pool_and_list_pool_ids(#[case] seed: Seed) { DEFAULT_ACCOUNT_INDEX, UtxoType::Transfer | UtxoType::LockThenTransfer, UtxoState::Confirmed.into(), + WithLocked::Unlocked, ) .unwrap() .get(&Currency::Coin) @@ -926,6 +937,7 @@ fn create_stake_pool_and_list_pool_ids(#[case] seed: Seed) { DEFAULT_ACCOUNT_INDEX, UtxoType::Transfer | UtxoType::LockThenTransfer, UtxoState::Confirmed.into(), + WithLocked::Unlocked, ) .unwrap() .get(&Currency::Coin) @@ -964,6 +976,7 @@ fn create_stake_pool_and_list_pool_ids(#[case] seed: Seed) { DEFAULT_ACCOUNT_INDEX, UtxoType::Transfer | UtxoType::LockThenTransfer | UtxoType::CreateStakePool, UtxoState::Confirmed.into(), + WithLocked::Unlocked, ) .unwrap(); assert_eq!( @@ -973,6 +986,40 @@ fn create_stake_pool_and_list_pool_ids(#[case] seed: Seed) { let pool_ids = wallet.get_pool_ids(DEFAULT_ACCOUNT_INDEX).unwrap(); assert_eq!(pool_ids.len(), 1); + + let pool_id = pool_ids.first().unwrap().0; + let decommission_tx = wallet + .decommission_stake_pool( + &mut WalletEventsNoOp, + DEFAULT_ACCOUNT_INDEX, + pool_id, + pool_amount, + FeeRate::new(Amount::from_atoms(0)), + ) + .unwrap(); + + let block3 = Block::new( + vec![decommission_tx], + block1_id.into(), + block1_timestamp, + ConsensusData::None, + BlockReward::new(vec![]), + ) + .unwrap(); + scan_wallet(&mut wallet, BlockHeight::new(2), vec![block3]); + + let currency_balances = wallet + .get_balance( + DEFAULT_ACCOUNT_INDEX, + UtxoType::Transfer | UtxoType::LockThenTransfer | UtxoType::CreateStakePool, + UtxoState::Confirmed.into(), + WithLocked::Unlocked, + ) + .unwrap(); + assert_eq!( + currency_balances.get(&Currency::Coin).copied().unwrap_or(Amount::ZERO), + pool_amount, + ); } #[rstest] @@ -990,6 +1037,7 @@ fn issue_and_transfer_tokens(#[case] seed: Seed) { DEFAULT_ACCOUNT_INDEX, UtxoType::Transfer | UtxoType::LockThenTransfer, UtxoState::Confirmed.into(), + WithLocked::Unlocked, ) .unwrap() .get(&Currency::Coin) @@ -1032,6 +1080,7 @@ fn issue_and_transfer_tokens(#[case] seed: Seed) { DEFAULT_ACCOUNT_INDEX, UtxoType::Transfer | UtxoType::LockThenTransfer, UtxoState::Confirmed.into(), + WithLocked::Unlocked, ) .unwrap() .get(&Currency::Coin) @@ -1102,6 +1151,7 @@ fn issue_and_transfer_tokens(#[case] seed: Seed) { DEFAULT_ACCOUNT_INDEX, UtxoType::Transfer | UtxoType::LockThenTransfer, UtxoState::Confirmed.into(), + WithLocked::Unlocked, ) .unwrap(); assert_eq!( @@ -1164,6 +1214,7 @@ fn issue_and_transfer_tokens(#[case] seed: Seed) { DEFAULT_ACCOUNT_INDEX, UtxoType::Transfer | UtxoType::LockThenTransfer, UtxoState::Confirmed.into(), + WithLocked::Unlocked, ) .unwrap(); assert_eq!( @@ -1234,6 +1285,7 @@ fn lock_then_transfer(#[case] seed: Seed) { DEFAULT_ACCOUNT_INDEX, UtxoType::Transfer | UtxoType::LockThenTransfer, UtxoState::Confirmed.into(), + WithLocked::Unlocked, ) .unwrap() .get(&Currency::Coin) @@ -1279,6 +1331,7 @@ fn lock_then_transfer(#[case] seed: Seed) { DEFAULT_ACCOUNT_INDEX, UtxoType::Transfer | UtxoType::LockThenTransfer, UtxoState::Confirmed.into(), + WithLocked::Unlocked, ) .unwrap() .get(&Currency::Coin) @@ -1346,6 +1399,7 @@ fn lock_then_transfer(#[case] seed: Seed) { DEFAULT_ACCOUNT_INDEX, UtxoType::Transfer | UtxoType::LockThenTransfer, UtxoState::Confirmed.into(), + WithLocked::Unlocked, ) .unwrap(); assert_eq!( @@ -1360,13 +1414,29 @@ fn lock_then_transfer(#[case] seed: Seed) { DEFAULT_ACCOUNT_INDEX, UtxoType::Transfer | UtxoType::LockThenTransfer, UtxoState::Confirmed.into(), + WithLocked::Unlocked, ) .unwrap(); + // check that the amount is still not unlocked assert_eq!( currency_balances.get(&Currency::Coin).copied().unwrap_or(Amount::ZERO), balance_without_locked_transer ); + let currency_balances = wallet + .get_balance( + DEFAULT_ACCOUNT_INDEX, + UtxoType::Transfer | UtxoType::LockThenTransfer, + UtxoState::Confirmed.into(), + WithLocked::Locked, + ) + .unwrap(); + // check that the amount is still locked + assert_eq!( + currency_balances.get(&Currency::Coin).copied().unwrap_or(Amount::ZERO), + amount_to_lock_then_transfer + ); + let new_block = Block::new( vec![], chain_config.genesis_block_id(), @@ -1387,6 +1457,7 @@ fn lock_then_transfer(#[case] seed: Seed) { DEFAULT_ACCOUNT_INDEX, UtxoType::Transfer | UtxoType::LockThenTransfer, UtxoState::Confirmed.into(), + WithLocked::Unlocked, ) .unwrap(); assert_eq!( @@ -1456,6 +1527,7 @@ fn wallet_sync_new_account(#[case] seed: Seed) { new_account_index, UtxoType::Transfer.into(), UtxoState::Confirmed.into(), + WithLocked::Unlocked, ) .unwrap(); @@ -1475,6 +1547,7 @@ fn wallet_sync_new_account(#[case] seed: Seed) { new_account_index, UtxoType::Transfer | UtxoType::LockThenTransfer, UtxoState::Confirmed.into(), + WithLocked::Unlocked, ) .unwrap() .get(&Currency::Coin) @@ -1529,6 +1602,7 @@ fn wallet_multiple_transactions_in_single_block(#[case] seed: Seed) { DEFAULT_ACCOUNT_INDEX, UtxoType::Transfer | UtxoType::LockThenTransfer, UtxoState::Confirmed.into(), + WithLocked::Unlocked, ) .unwrap() .get(&Currency::Coin) @@ -1598,6 +1672,7 @@ fn wallet_multiple_transactions_in_single_block(#[case] seed: Seed) { DEFAULT_ACCOUNT_INDEX, UtxoType::Transfer | UtxoType::LockThenTransfer, UtxoState::Confirmed.into(), + WithLocked::Unlocked, ) .unwrap() .get(&Currency::Coin) @@ -1621,6 +1696,7 @@ fn wallet_scan_multiple_transactions_from_mempool(#[case] seed: Seed) { DEFAULT_ACCOUNT_INDEX, UtxoType::Transfer | UtxoType::LockThenTransfer, UtxoState::Confirmed.into(), + WithLocked::Unlocked, ) .unwrap() .get(&Currency::Coin) @@ -1666,6 +1742,7 @@ fn wallet_scan_multiple_transactions_from_mempool(#[case] seed: Seed) { DEFAULT_ACCOUNT_INDEX, UtxoType::Transfer | UtxoType::LockThenTransfer, UtxoState::Confirmed.into(), + WithLocked::Unlocked, ) .unwrap() .get(&Currency::Coin) @@ -1816,6 +1893,7 @@ fn wallet_scan_multiple_transactions_from_mempool(#[case] seed: Seed) { DEFAULT_ACCOUNT_INDEX, UtxoType::Transfer | UtxoType::LockThenTransfer, UtxoState::Confirmed | UtxoState::InMempool | UtxoState::Inactive, + WithLocked::Unlocked, ) .unwrap() .get(&Currency::Coin) @@ -1830,6 +1908,7 @@ fn wallet_scan_multiple_transactions_from_mempool(#[case] seed: Seed) { DEFAULT_ACCOUNT_INDEX, UtxoType::Transfer | UtxoType::LockThenTransfer, UtxoState::Confirmed | UtxoState::InMempool | UtxoState::Inactive, + WithLocked::Unlocked, ) .unwrap() .get(&Currency::Coin) @@ -1844,6 +1923,7 @@ fn wallet_scan_multiple_transactions_from_mempool(#[case] seed: Seed) { DEFAULT_ACCOUNT_INDEX, UtxoType::Transfer | UtxoType::LockThenTransfer, UtxoState::Confirmed | UtxoState::InMempool | UtxoState::Inactive, + WithLocked::Unlocked, ) .unwrap() .get(&Currency::Coin) @@ -1867,6 +1947,7 @@ fn wallet_abandone_transactions(#[case] seed: Seed) { DEFAULT_ACCOUNT_INDEX, UtxoType::Transfer | UtxoType::LockThenTransfer, UtxoState::Confirmed.into(), + WithLocked::Unlocked, ) .unwrap() .get(&Currency::Coin) @@ -1910,6 +1991,7 @@ fn wallet_abandone_transactions(#[case] seed: Seed) { DEFAULT_ACCOUNT_INDEX, UtxoType::Transfer | UtxoType::LockThenTransfer, UtxoState::Confirmed.into(), + WithLocked::Unlocked, ) .unwrap() .get(&Currency::Coin) @@ -1972,6 +2054,7 @@ fn wallet_abandone_transactions(#[case] seed: Seed) { DEFAULT_ACCOUNT_INDEX, UtxoType::Transfer | UtxoType::LockThenTransfer, UtxoState::Confirmed | UtxoState::InMempool | UtxoState::Inactive, + WithLocked::Unlocked, ) .unwrap() .get(&Currency::Coin) @@ -1998,6 +2081,7 @@ fn wallet_abandone_transactions(#[case] seed: Seed) { DEFAULT_ACCOUNT_INDEX, UtxoType::Transfer | UtxoType::LockThenTransfer, UtxoState::Confirmed | UtxoState::InMempool | UtxoState::Inactive, + WithLocked::Unlocked, ) .unwrap() .get(&Currency::Coin) @@ -2012,6 +2096,7 @@ fn wallet_abandone_transactions(#[case] seed: Seed) { DEFAULT_ACCOUNT_INDEX, UtxoType::Transfer | UtxoType::LockThenTransfer, UtxoState::Confirmed | UtxoState::InMempool, + WithLocked::Unlocked, ) .unwrap() .get(&Currency::Coin) diff --git a/wallet/types/src/lib.rs b/wallet/types/src/lib.rs index 6eaf1f3c27..21e629e7b9 100644 --- a/wallet/types/src/lib.rs +++ b/wallet/types/src/lib.rs @@ -18,6 +18,7 @@ pub mod account_info; pub mod keys; pub mod utxo_types; pub mod wallet_tx; +pub mod with_locked; pub use account_id::{ AccountDerivationPathId, AccountId, AccountKeyPurposeId, AccountWalletCreatedTxId, diff --git a/wallet/types/src/with_locked.rs b/wallet/types/src/with_locked.rs new file mode 100644 index 0000000000..1f01aa56d9 --- /dev/null +++ b/wallet/types/src/with_locked.rs @@ -0,0 +1,22 @@ +// Copyright (c) 2023 RBB S.r.l +// opensource@mintlayer.org +// SPDX-License-Identifier: MIT +// Licensed under the MIT License; +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://github.com/mintlayer/mintlayer-core/blob/master/LICENSE +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +/// Enum used to specify whether to include locked balance/utxos for wallet commands +#[derive(Debug, Clone, Copy)] +pub enum WithLocked { + Any, + Unlocked, + Locked, +} diff --git a/wallet/wallet-cli-lib/src/commands/helper_types.rs b/wallet/wallet-cli-lib/src/commands/helper_types.rs index 9b39457272..5efb84b342 100644 --- a/wallet/wallet-cli-lib/src/commands/helper_types.rs +++ b/wallet/wallet-cli-lib/src/commands/helper_types.rs @@ -23,6 +23,7 @@ use common::{ chain::{block::timestamp::BlockTimestamp, ChainConfig, DelegationId, PoolId}, primitives::{Amount, BlockHeight}, }; +use wallet_types::with_locked::WithLocked; #[derive(Debug, Clone, Copy, ValueEnum)] pub enum CliUtxoTypes { @@ -124,3 +125,20 @@ pub fn format_delegation_info( balance.into_fixedpoint_str(chain_config.coin_decimals()), ) } + +#[derive(Debug, Clone, Copy, ValueEnum)] +pub enum CliWithLocked { + Any, + Unlocked, + Locked, +} + +impl CliWithLocked { + pub fn to_wallet_type(self) -> WithLocked { + match self { + CliWithLocked::Any => WithLocked::Any, + CliWithLocked::Unlocked => WithLocked::Unlocked, + CliWithLocked::Locked => WithLocked::Locked, + } + } +} diff --git a/wallet/wallet-cli-lib/src/commands/mod.rs b/wallet/wallet-cli-lib/src/commands/mod.rs index 824c18a9fc..e1caca5fdc 100644 --- a/wallet/wallet-cli-lib/src/commands/mod.rs +++ b/wallet/wallet-cli-lib/src/commands/mod.rs @@ -35,7 +35,9 @@ use wallet_controller::{NodeInterface, NodeRpcClient, PeerId, DEFAULT_ACCOUNT_IN use crate::{errors::WalletCliError, CliController}; -use self::helper_types::{format_delegation_info, format_pool_info, CliUtxoState, CliUtxoTypes}; +use self::helper_types::{ + format_delegation_info, format_pool_info, CliUtxoState, CliUtxoTypes, CliWithLocked, +}; #[derive(Debug, Parser)] #[clap(rename_all = "lower")] @@ -181,6 +183,8 @@ pub enum WalletCommand { SyncWallet, GetBalance { + #[arg(value_enum, default_value_t = CliWithLocked::Unlocked)] + with_locked: CliWithLocked, #[arg(default_values_t = vec![CliUtxoState::Confirmed])] utxo_states: Vec, }, @@ -188,6 +192,8 @@ pub enum WalletCommand { ListUtxo { #[arg(value_enum, default_value_t = CliUtxoTypes::All)] utxo_type: CliUtxoTypes, + #[arg(value_enum, default_value_t = CliWithLocked::Unlocked)] + with_locked: CliWithLocked, #[arg(default_values_t = vec![CliUtxoState::Confirmed])] utxo_states: Vec, }, @@ -849,13 +855,17 @@ impl CommandHandler { Ok(ConsoleCommand::Print("Success".to_owned())) } - WalletCommand::GetBalance { utxo_states } => { + WalletCommand::GetBalance { + utxo_states, + with_locked, + } => { let mut balances = controller_opt .as_mut() .ok_or(WalletCliError::NoWallet)? .get_balance( selected_account.ok_or(WalletCliError::NoSelectedAccount)?, CliUtxoState::to_wallet_states(utxo_states), + with_locked.to_wallet_type(), ) .map_err(WalletCliError::Controller)?; let coin_balance = balances.remove(&Currency::Coin).unwrap_or(Amount::ZERO); @@ -884,6 +894,7 @@ impl CommandHandler { WalletCommand::ListUtxo { utxo_type, utxo_states, + with_locked, } => { let utxos = controller_opt .as_mut() @@ -892,6 +903,7 @@ impl CommandHandler { selected_account.ok_or(WalletCliError::NoSelectedAccount)?, utxo_type.to_wallet_types(), CliUtxoState::to_wallet_states(utxo_states), + with_locked.to_wallet_type(), ) .map_err(WalletCliError::Controller)?; Ok(ConsoleCommand::Print(format!("{utxos:#?}"))) diff --git a/wallet/wallet-controller/src/lib.rs b/wallet/wallet-controller/src/lib.rs index 0e9700596b..823edb3a27 100644 --- a/wallet/wallet-controller/src/lib.rs +++ b/wallet/wallet-controller/src/lib.rs @@ -76,11 +76,11 @@ use wallet::{ wallet_events::WalletEvents, DefaultWallet, WalletError, }; -use wallet_types::BlockInfo; pub use wallet_types::{ account_info::DEFAULT_ACCOUNT_INDEX, utxo_types::{UtxoState, UtxoStates, UtxoType, UtxoTypes}, }; +use wallet_types::{with_locked::WithLocked, BlockInfo}; #[derive(thiserror::Error, Debug)] pub enum ControllerError { @@ -225,12 +225,14 @@ impl Controll &self, account_index: U31, utxo_states: UtxoStates, + with_locked: WithLocked, ) -> Result, ControllerError> { self.wallet .get_balance( account_index, UtxoType::Transfer | UtxoType::LockThenTransfer, utxo_states, + with_locked, ) .map_err(ControllerError::WalletError) } @@ -240,9 +242,10 @@ impl Controll account_index: U31, utxo_types: UtxoTypes, utxo_states: UtxoStates, + with_locked: WithLocked, ) -> Result, ControllerError> { self.wallet - .get_utxos(account_index, utxo_types, utxo_states) + .get_utxos(account_index, utxo_types, utxo_states, with_locked) .map_err(ControllerError::WalletError) } @@ -747,6 +750,7 @@ impl Controll account_index, UtxoType::CreateStakePool | UtxoType::ProduceBlockFromStake, UtxoState::Confirmed.into(), + WithLocked::Unlocked, ) .map_err(ControllerError::WalletError)?; let pool_ids = stake_pool_utxos.values().filter_map(|utxo| match utxo {