Skip to content

Commit

Permalink
refactor wallet utxo token_id uses
Browse files Browse the repository at this point in the history
  • Loading branch information
@OBorce
OBorce committed Aug 24, 2024
1 parent d9eee6f commit bb4e3a7
Show file tree
Hide file tree
Showing 10 changed files with 125 additions and 137 deletions.
20 changes: 0 additions & 20 deletions common/src/chain/tokens/tokens_utils.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -87,23 +87,3 @@ pub fn is_token_or_nft_issuance(output: &TxOutput) -> bool {
TxOutput::IssueFungibleToken(_) | TxOutput::IssueNft(_, _, _) => true,
}
}

pub fn get_token_id_for_tx_output(output: &TxOutput, inputs: &[TxInput]) -> Option<TokenId> {
match output {
TxOutput::Transfer(v, _) | TxOutput::LockThenTransfer(v, _, _) | TxOutput::Burn(v) => {
match v {
OutputValue::Coin(_) | OutputValue::TokenV0(_) => None,
OutputValue::TokenV1(token_id, _) => Some(*token_id),
}
}
TxOutput::CreateStakePool(_, _)
| TxOutput::ProduceBlockFromStake(_, _)
| TxOutput::CreateDelegationId(_, _)
| TxOutput::DelegateStaking(_, _)
| TxOutput::DataDeposit(_)
| TxOutput::Htlc(_, _)
| TxOutput::IssueFungibleToken(_)
| TxOutput::AnyoneCanTake(_) => None,
TxOutput::IssueNft(_, _, _) => make_token_id(inputs),
}
}
2 changes: 1 addition & 1 deletion wallet/src/account/currency_grouper/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ pub enum Currency {
Token(TokenId),
}

pub(crate) fn group_outputs<T, Grouped: Clone>(
pub fn group_outputs<T, Grouped: Clone>(
outputs: impl Iterator<Item = T>,
get_tx_output: impl Fn(&T) -> &TxOutput,
mut combiner: impl FnMut(&mut Grouped, &T, Amount) -> WalletResult<()>,
Expand Down
46 changes: 17 additions & 29 deletions wallet/src/account/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -246,10 +246,7 @@ impl<K: AccountKeyChains> Account<K> {
)
}
SelectedInputs::Inputs(ref inputs) => (
inputs
.iter()
.map(|(outpoint, utxo)| (outpoint.clone(), (utxo, None)))
.collect(),
inputs.iter().map(|(outpoint, utxo)| (outpoint.clone(), utxo)).collect(),
CoinSelectionAlgo::UsePreselected,
),
}
Expand Down Expand Up @@ -415,7 +412,7 @@ impl<K: AccountKeyChains> Account<K> {
&self,
fee_rates: CurrentFeeRate,
pay_fee_with_currency: &currency_grouper::Currency,
utxos: Vec<(UtxoOutPoint, (&TxOutput, Option<TokenId>))>,
utxos: Vec<(UtxoOutPoint, &TxOutput)>,
) -> Result<BTreeMap<currency_grouper::Currency, Vec<OutputGroup>>, WalletError> {
let utxo_to_output_group =
|(outpoint, txo): (UtxoOutPoint, TxOutput)| -> WalletResult<OutputGroup> {
Expand Down Expand Up @@ -443,9 +440,9 @@ impl<K: AccountKeyChains> Account<K> {

currency_grouper::group_utxos_for_input(
utxos.into_iter(),
|(_, (tx_output, _))| tx_output,
|(_, tx_output)| tx_output,
|grouped: &mut Vec<(UtxoOutPoint, TxOutput)>, element, _| -> WalletResult<()> {
grouped.push((element.0.clone(), element.1 .0.clone()));
grouped.push((element.0.clone(), element.1.clone()));
Ok(())
},
vec![],
Expand Down Expand Up @@ -1125,8 +1122,7 @@ impl<K: AccountKeyChains> Account<K> {
outpoint: &UtxoOutPoint,
current_block_info: BlockInfo,
) -> WalletResult<(TxOutput, Destination)> {
let (txo, _) =
self.output_cache.find_unspent_unlocked_utxo(outpoint, current_block_info)?;
let txo = self.output_cache.find_unspent_unlocked_utxo(outpoint, current_block_info)?;

Ok((
txo.clone(),
Expand Down Expand Up @@ -1241,14 +1237,14 @@ impl<K: AccountKeyChains> Account<K> {
};
let amounts_by_currency = currency_grouper::group_utxos_for_input(
self.output_cache
.utxos_with_token_ids(
.utxos(
current_block_info,
UtxoState::Confirmed.into(),
WithLocked::Unlocked,
|txo| get_utxo_type(txo).is_some() && self.is_watched_by(txo, &address),
)
.into_iter(),
|(_, (tx_output, _))| tx_output,
|(_, tx_output)| tx_output,
|total: &mut Amount, _, amount| -> WalletResult<()> {
*total = (*total + amount).ok_or(WalletError::OutputAmountOverflow)?;
Ok(())
Expand Down Expand Up @@ -1433,7 +1429,7 @@ impl<K: AccountKeyChains> Account<K> {
with_locked,
)
.into_iter(),
|(_, (tx_output, _))| tx_output,
|(_, tx_output)| tx_output,
|total: &mut Amount, _, amount| -> WalletResult<()> {
*total = (*total + amount).ok_or(WalletError::OutputAmountOverflow)?;
Ok(())
Expand All @@ -1449,20 +1445,15 @@ impl<K: AccountKeyChains> Account<K> {
median_time: BlockTimestamp,
utxo_states: UtxoStates,
with_locked: WithLocked,
) -> Vec<(UtxoOutPoint, (&TxOutput, Option<TokenId>))> {
) -> Vec<(UtxoOutPoint, &TxOutput)> {
let current_block_info = BlockInfo {
height: self.account_info.best_block_height(),
timestamp: median_time,
};
self.output_cache.utxos_with_token_ids(
current_block_info,
utxo_states,
with_locked,
|txo| {
self.is_watched_multisig_output(txo)
&& get_utxo_type(txo).is_some_and(|v| utxo_types.contains(v))
},
)
self.output_cache.utxos(current_block_info, utxo_states, with_locked, |txo| {
self.is_watched_multisig_output(txo)
&& get_utxo_type(txo).is_some_and(|v| utxo_types.contains(v))
})
}

pub fn get_utxos(
Expand All @@ -1471,17 +1462,14 @@ impl<K: AccountKeyChains> Account<K> {
median_time: BlockTimestamp,
utxo_states: UtxoStates,
with_locked: WithLocked,
) -> Vec<(UtxoOutPoint, (&TxOutput, Option<TokenId>))> {
) -> Vec<(UtxoOutPoint, &TxOutput)> {
let current_block_info = BlockInfo {
height: self.account_info.best_block_height(),
timestamp: median_time,
};
self.output_cache.utxos_with_token_ids(
current_block_info,
utxo_states,
with_locked,
|txo| self.is_mine(txo) && get_utxo_type(txo).is_some_and(|v| utxo_types.contains(v)),
)
self.output_cache.utxos(current_block_info, utxo_states, with_locked, |txo| {
self.is_mine(txo) && get_utxo_type(txo).is_some_and(|v| utxo_types.contains(v))
})
}

pub fn get_transaction_list(&self, skip: usize, count: usize) -> WalletResult<TransactionList> {
Expand Down
54 changes: 20 additions & 34 deletions wallet/src/account/output_cache/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,9 +25,9 @@ use common::{
output_value::OutputValue,
stakelock::StakePoolData,
tokens::{
get_token_id_for_tx_output, make_token_id, IsTokenFreezable, IsTokenUnfreezable,
RPCFungibleTokenInfo, RPCIsTokenFrozen, RPCNonFungibleTokenInfo, RPCTokenTotalSupply,
TokenId, TokenIssuance, TokenTotalSupply,
make_token_id, IsTokenFreezable, IsTokenUnfreezable, RPCFungibleTokenInfo,
RPCIsTokenFrozen, RPCNonFungibleTokenInfo, RPCTokenTotalSupply, TokenId, TokenIssuance,
TokenTotalSupply,
},
AccountCommand, AccountNonce, AccountSpending, DelegationId, Destination, GenBlock,
OutPointSourceId, PoolId, Transaction, TxInput, TxOutput, UtxoOutPoint,
Expand All @@ -47,9 +47,11 @@ use wallet_types::{
AccountWalletTxId, BlockInfo, WalletTx,
};

use crate::{get_tx_output_destination, WalletError, WalletResult};
use crate::{
get_tx_output_destination, send_request::get_referenced_token_ids, WalletError, WalletResult,
};

pub type UtxoWithTxOutput<'a> = (UtxoOutPoint, (&'a TxOutput, Option<TokenId>));
pub type UtxoWithTxOutput<'a> = (UtxoOutPoint, &'a TxOutput);

#[derive(Debug, Eq, PartialEq, Clone, serde::Serialize, serde::Deserialize, HasValueHint)]
pub struct TxInfo {
Expand Down Expand Up @@ -1128,7 +1130,7 @@ impl OutputCache {
&self,
utxo: &UtxoOutPoint,
current_block_info: BlockInfo,
) -> WalletResult<(&TxOutput, Option<TokenId>)> {
) -> WalletResult<&TxOutput> {
let tx = self
.txs
.get(&utxo.source_id())
Expand Down Expand Up @@ -1163,29 +1165,22 @@ impl OutputCache {
WalletError::TokenV0Utxo(utxo.clone())
);

let token_id = match tx {
WalletTx::Tx(tx_data) => {
get_token_id_for_tx_output(output, tx_data.get_transaction().inputs())
}
WalletTx::Block(_) => None,
};

Ok((output, token_id))
Ok(output)
}

pub fn find_used_tokens(
&self,
current_block_info: BlockInfo,
inputs: &[UtxoOutPoint],
) -> WalletResult<BTreeSet<TokenId>> {
inputs
.iter()
.filter_map(|utxo| {
self.find_unspent_unlocked_utxo(utxo, current_block_info)
.map(|(_, token_id)| token_id)
.transpose()
})
.collect()
inputs.iter().try_fold(BTreeSet::new(), |mut token_ids, utxo| {
let new_ids = self
.find_unspent_unlocked_utxo(utxo, current_block_info)
.map(get_referenced_token_ids)?;
token_ids.extend(new_ids);

Ok(token_ids)
})
}

pub fn find_utxos(
Expand All @@ -1202,13 +1197,13 @@ impl OutputCache {
.collect()
}

pub fn utxos_with_token_ids<F: Fn(&TxOutput) -> bool>(
pub fn utxos<F: Fn(&TxOutput) -> bool>(
&self,
current_block_info: BlockInfo,
utxo_states: UtxoStates,
locked_state: WithLocked,
output_filter: F,
) -> Vec<(UtxoOutPoint, (&TxOutput, Option<TokenId>))> {
) -> Vec<(UtxoOutPoint, &TxOutput)> {
let output_filter = &output_filter;
self.txs
.values()
Expand All @@ -1232,16 +1227,7 @@ impl OutputCache {
&& !is_v0_token_output(output)
&& output_filter(output)
})
.map(move |(output, outpoint)| {
let token_id = match tx {
WalletTx::Tx(tx_data) => get_token_id_for_tx_output(
output,
tx_data.get_transaction().inputs(),
),
WalletTx::Block(_) => None,
};
(outpoint, (output, token_id))
})
.map(|(output, outpoint)| (outpoint, output))
})
.collect()
}
Expand Down
32 changes: 31 additions & 1 deletion wallet/src/send_request/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
// See the License for the specific language governing permissions and
// limitations under the License.

use std::collections::BTreeMap;
use std::collections::{BTreeMap, BTreeSet};
use std::mem::take;

use common::address::Address;
Expand Down Expand Up @@ -397,6 +397,36 @@ where
}
}

/// Get any referenced token by this output
/// ignore tokens V0
pub fn get_referenced_token_ids(output: &TxOutput) -> BTreeSet<TokenId> {
match output {
TxOutput::Transfer(v, _)
| TxOutput::LockThenTransfer(v, _, _)
| TxOutput::Burn(v)
| TxOutput::Htlc(v, _) => referenced_token_id(v),
| TxOutput::AnyoneCanTake(data) => {
let mut tokens = referenced_token_id(data.ask());
tokens.extend(referenced_token_id(data.give()));
tokens
}
TxOutput::CreateStakePool(_, _)
| TxOutput::ProduceBlockFromStake(_, _)
| TxOutput::CreateDelegationId(_, _)
| TxOutput::DelegateStaking(_, _)
| TxOutput::DataDeposit(_)
| TxOutput::IssueFungibleToken(_) => BTreeSet::new(),
TxOutput::IssueNft(token_id, _, _) => BTreeSet::from_iter([*token_id]),
}
}

fn referenced_token_id(v: &OutputValue) -> BTreeSet<TokenId> {
match v {
OutputValue::Coin(_) | OutputValue::TokenV0(_) => BTreeSet::new(),
OutputValue::TokenV1(token_id, _) => BTreeSet::from_iter([*token_id]),
}
}

pub enum SelectedInputs {
Utxos(Vec<UtxoOutPoint>),
Inputs(Vec<(UtxoOutPoint, TxOutput)>),
Expand Down
20 changes: 6 additions & 14 deletions wallet/src/wallet/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1118,18 +1118,15 @@ where
utxo_types: UtxoTypes,
utxo_states: UtxoStates,
with_locked: WithLocked,
) -> WalletResult<Vec<(UtxoOutPoint, TxOutput, Option<TokenId>)>> {
) -> WalletResult<Vec<(UtxoOutPoint, TxOutput)>> {
let account = self.get_account(account_index)?;
let utxos = account.get_multisig_utxos(
utxo_types,
self.latest_median_time,
utxo_states,
with_locked,
);
let utxos = utxos
.into_iter()
.map(|(outpoint, (txo, token_id))| (outpoint, txo.clone(), token_id))
.collect();
let utxos = utxos.into_iter().map(|(outpoint, txo)| (outpoint, txo.clone())).collect();
Ok(utxos)
}

Expand All @@ -1139,18 +1136,15 @@ where
utxo_types: UtxoTypes,
utxo_states: UtxoStates,
with_locked: WithLocked,
) -> WalletResult<Vec<(UtxoOutPoint, TxOutput, Option<TokenId>)>> {
) -> WalletResult<Vec<(UtxoOutPoint, TxOutput)>> {
let account = self.get_account(account_index)?;
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(), token_id))
.collect();
let utxos = utxos.into_iter().map(|(outpoint, txo)| (outpoint, txo.clone())).collect();
Ok(utxos)
}

Expand Down Expand Up @@ -1452,14 +1446,12 @@ where
&mut self,
account_index: U31,
destination: Destination,
inputs: Vec<(UtxoOutPoint, TxOutput, Option<TokenId>)>,
inputs: Vec<(UtxoOutPoint, TxOutput)>,
current_fee_rate: FeeRate,
additional_utxo_infos: &BTreeMap<PoolOrTokenId, UtxoAdditionalInfo>,
) -> WalletResult<SignedTransaction> {
let request = SendRequest::new().with_inputs(
inputs
.into_iter()
.map(|(outpoint, output, _)| (TxInput::Utxo(outpoint), output)),
inputs.into_iter().map(|(outpoint, output)| (TxInput::Utxo(outpoint), output)),
&|_| None,
)?;

Expand Down
6 changes: 3 additions & 3 deletions wallet/src/wallet/tests.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1501,7 +1501,7 @@ fn spend_from_user_specified_utxos(#[case] seed: Seed) {

let selected_utxos = utxos
.iter()
.map(|(outpoint, _, _)| outpoint)
.map(|(outpoint, _)| outpoint)
.take(rng.gen_range(1..utxos.len()))
.cloned()
.collect_vec();
Expand Down Expand Up @@ -1635,7 +1635,7 @@ fn create_stake_pool_and_list_pool_ids(#[case] seed: Seed) {
)
.unwrap();
assert_eq!(create_stake_pool_utxos.len(), 1);
let (_, output, _) = create_stake_pool_utxos.pop().unwrap();
let (_, output) = create_stake_pool_utxos.pop().unwrap();
match output {
TxOutput::CreateStakePool(id, data) => {
assert_eq!(id, *pool_id);
Expand Down Expand Up @@ -4555,7 +4555,7 @@ fn sign_send_request_cold_wallet(#[case] seed: Seed) {
.unwrap();

assert_eq!(utxos.len(), 1);
let (_, output, _) = utxos.pop().unwrap();
let (_, output) = utxos.pop().unwrap();

matches!(output, TxOutput::Transfer(OutputValue::Coin(value), dest)
if value == balance && dest == cold_wallet_address.into_object());
Expand Down
2 changes: 1 addition & 1 deletion wallet/wallet-controller/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -762,7 +762,7 @@ where
),
}
.map_err(ControllerError::WalletError)?;
let pool_ids = stake_pool_utxos.into_iter().filter_map(|(_, utxo, _)| match utxo {
let pool_ids = stake_pool_utxos.into_iter().filter_map(|(_, utxo)| match utxo {
TxOutput::ProduceBlockFromStake(_, pool_id) | TxOutput::CreateStakePool(pool_id, _) => {
Some(pool_id)
}
Expand Down
Loading

0 comments on commit bb4e3a7

Please sign in to comment.