diff --git a/.changelog/unreleased/improvements/3029-token-balance-update-fns.md b/.changelog/unreleased/improvements/3029-token-balance-update-fns.md new file mode 100644 index 0000000000..30cd83dcbd --- /dev/null +++ b/.changelog/unreleased/improvements/3029-token-balance-update-fns.md @@ -0,0 +1,2 @@ +- Refactor and modularize the token balance and supply API. + ([\#3029](https://github.com/anoma/namada/pull/3029)) \ No newline at end of file diff --git a/crates/apps/src/lib/node/ledger/shell/finalize_block.rs b/crates/apps/src/lib/node/ledger/shell/finalize_block.rs index 4c6b6c419a..5dc22ddeff 100644 --- a/crates/apps/src/lib/node/ledger/shell/finalize_block.rs +++ b/crates/apps/src/lib/node/ledger/shell/finalize_block.rs @@ -766,7 +766,10 @@ mod test_finalize_block { use namada::proof_of_stake::{unjail_validator, ADDRESS as pos_address}; use namada::replay_protection; use namada::tendermint::abci::types::{Misbehavior, MisbehaviorKind}; - use namada::token::{Amount, DenominatedAmount, NATIVE_MAX_DECIMAL_PLACES}; + use namada::token::{ + read_balance, update_balance, Amount, DenominatedAmount, + NATIVE_MAX_DECIMAL_PLACES, + }; use namada::tx::data::Fee; use namada::tx::{Authorization, Code, Data}; use namada::vote_ext::ethereum_events; @@ -842,14 +845,14 @@ mod test_finalize_block { let mut processed_txs = vec![]; // Add unshielded balance for fee payment - let balance_key = token::storage_key::balance_key( - &shell.state.in_mem().native_token, + let native_token = shell.state.in_mem().native_token.clone(); + update_balance( + &mut shell.state, + &native_token, &Address::from(&keypair.ref_to()), - ); - shell - .state - .write(&balance_key, Amount::native_whole(1000)) - .unwrap(); + |_| Ok(Amount::native_whole(1000)), + ) + .unwrap(); // create some wrapper txs for i in 0u64..4 { @@ -1143,14 +1146,14 @@ mod test_finalize_block { // add bertha's gas fees the pool { let amt: Amount = 999_999_u64.into(); - let pool_balance_key = token::storage_key::balance_key( - &shell.state.in_mem().native_token, + let native_token = shell.state.in_mem().native_token.clone(); + update_balance( + &mut shell.state, + &native_token, &bridge_pool::BRIDGE_POOL_ADDRESS, - ); - shell - .state - .write(&pool_balance_key, amt) - .expect("Test failed"); + |_| Ok(amt), + ) + .expect("Test failed"); } // write transfer to storage let transfer = { @@ -3117,12 +3120,12 @@ mod test_finalize_block { initial_balance, ) .unwrap(); - let balance_key = token::storage_key::balance_key( + let balance = read_balance( + &shell.state, &native_token, &Address::from(&keypair.to_public()), - ); - let balance: Amount = - shell.state.read(&balance_key).unwrap().unwrap_or_default(); + ) + .unwrap(); assert_eq!(balance, initial_balance); let mut wrapper = @@ -3178,8 +3181,12 @@ mod test_finalize_block { assert_eq!(event.event_type.to_string(), String::from("applied")); let code = event.attributes.get("code").expect("Test failed").as_str(); assert_eq!(code, String::from(ResultCode::InvalidTx).as_str()); - let balance: Amount = - shell.state.read(&balance_key).unwrap().unwrap_or_default(); + let balance = read_balance( + &shell.state, + &native_token, + &Address::from(&keypair.to_public()), + ) + .unwrap(); assert_eq!(balance, 0.into()) } @@ -3701,15 +3708,12 @@ mod test_finalize_block { // Slash pool balance let nam_address = shell.state.in_mem().native_token.clone(); - let slash_balance_key = token::storage_key::balance_key( + let slash_pool_balance_init = read_balance( + &shell.state, &nam_address, &namada_proof_of_stake::SLASH_POOL_ADDRESS, - ); - let slash_pool_balance_init: token::Amount = shell - .state - .read(&slash_balance_key) - .expect("must be able to read") - .unwrap_or_default(); + ) + .unwrap(); debug_assert_eq!(slash_pool_balance_init, token::Amount::zero()); let consensus_set: Vec = @@ -4858,14 +4862,8 @@ mod test_finalize_block { // NOTE: assumed that the only change in pos address balance by // advancing to the next epoch is minted inflation - no change occurs // due to slashing - let pos_balance_pre = shell - .state - .read::(&token::storage_key::balance_key( - &staking_token, - &pos_address, - )) - .unwrap() - .unwrap_or_default(); + let pos_balance_pre = + read_balance(&shell.state, &staking_token, &pos_address).unwrap(); loop { next_block_for_inflation( shell, @@ -4877,14 +4875,8 @@ mod test_finalize_block { break; } } - let pos_balance_post = shell - .state - .read::(&token::storage_key::balance_key( - &staking_token, - &pos_address, - )) - .unwrap() - .unwrap_or_default(); + let pos_balance_post = + read_balance(&shell.state, &staking_token, &pos_address).unwrap(); ( shell.state.in_mem().block.epoch, diff --git a/crates/apps/src/lib/node/ledger/shell/init_chain.rs b/crates/apps/src/lib/node/ledger/shell/init_chain.rs index c75e72615f..3fd22151f8 100644 --- a/crates/apps/src/lib/node/ledger/shell/init_chain.rs +++ b/crates/apps/src/lib/node/ledger/shell/init_chain.rs @@ -511,7 +511,6 @@ where continue; }; - let mut total_token_balance = token::Amount::zero(); for (owner, balance) in balances { if let genesis::GenesisAddress::PublicKey(pk) = owner { namada::account::init_account_storage( @@ -535,15 +534,7 @@ where balance.amount(), ) .expect("Couldn't credit initial balance"); - total_token_balance += balance.amount(); } - // Write the total amount of tokens for the ratio - self.state - .write( - &token::storage_key::minted_balance_key(token_address), - total_token_balance, - ) - .unwrap(); } self.proceed_with(()) } diff --git a/crates/ethereum_bridge/src/protocol/transactions/ethereum_events/events.rs b/crates/ethereum_bridge/src/protocol/transactions/ethereum_events/events.rs index 12055b4e44..ac998bedbe 100644 --- a/crates/ethereum_bridge/src/protocol/transactions/ethereum_events/events.rs +++ b/crates/ethereum_bridge/src/protocol/transactions/ethereum_events/events.rs @@ -23,8 +23,8 @@ use namada_parameters::read_epoch_duration_parameter; use namada_state::{DBIter, StorageHasher, WlState, DB}; use namada_storage::{StorageRead, StorageWrite}; use namada_trans_token::storage_key::{balance_key, minted_balance_key}; +use token::{burn_tokens, decrement_total_supply, increment_total_supply}; -use crate::protocol::transactions::update; use crate::storage::bridge_pool::{ get_nonce_key, is_pending_transfer_key, BRIDGE_POOL_ADDRESS, }; @@ -162,48 +162,9 @@ where let native_werc20_supply_key = minted_balance_key(&erc20_token_address(native_erc20)); - update::amount(state, ð_bridge_native_token_balance_key, |balance| { - tracing::debug!( - %eth_bridge_native_token_balance_key, - ?balance, - "Existing value found", - ); - balance.spend(amount)?; - tracing::debug!( - %eth_bridge_native_token_balance_key, - ?balance, - "New value calculated", - ); - Ok(()) - })?; - update::amount(state, &receiver_native_token_balance_key, |balance| { - tracing::debug!( - %receiver_native_token_balance_key, - ?balance, - "Existing value found", - ); - balance.receive(amount)?; - tracing::debug!( - %receiver_native_token_balance_key, - ?balance, - "New value calculated", - ); - Ok(()) - })?; - update::amount(state, &native_werc20_supply_key, |balance| { - tracing::debug!( - %native_werc20_supply_key, - ?balance, - "Existing value found", - ); - balance.spend(amount)?; - tracing::debug!( - %native_werc20_supply_key, - ?balance, - "New value calculated", - ); - Ok(()) - })?; + let native_token = state.in_mem().native_token.clone(); + token::transfer(state, &native_token, &BRIDGE_ADDRESS, receiver, *amount)?; + decrement_total_supply(state, &erc20_token_address(native_erc20), *amount)?; tracing::info!( amount = %amount.to_string_native(), @@ -255,38 +216,11 @@ where .flatten(); for (token, ref amount) in assets_to_mint { - let balance_key = balance_key(&token, receiver); - update::amount(state, &balance_key, |balance| { - tracing::debug!( - %balance_key, - ?balance, - "Existing value found", - ); - balance.receive(amount)?; - tracing::debug!( - %balance_key, - ?balance, - "New value calculated", - ); - Ok(()) - })?; - _ = changed_keys.insert(balance_key); + token::credit_tokens(state, &token, receiver, *amount)?; + let balance_key = balance_key(&token, receiver); let supply_key = minted_balance_key(&token); - update::amount(state, &supply_key, |supply| { - tracing::debug!( - %supply_key, - ?supply, - "Existing value found", - ); - supply.receive(amount)?; - tracing::debug!( - %supply_key, - ?supply, - "New value calculated", - ); - Ok(()) - })?; + _ = changed_keys.insert(balance_key); _ = changed_keys.insert(supply_key); } @@ -347,14 +281,16 @@ where balance_key(&pending_transfer.gas_fee.token, &BRIDGE_POOL_ADDRESS); let relayer_rewards_key = balance_key(&pending_transfer.gas_fee.token, relayer); - // give the relayer the gas fee for this transfer. - update::amount(state, &relayer_rewards_key, |balance| { - balance.receive(&pending_transfer.gas_fee.amount) - })?; - // the gas fee is removed from escrow. - update::amount(state, &pool_balance_key, |balance| { - balance.spend(&pending_transfer.gas_fee.amount) - })?; + // give the relayer the gas fee for this transfer and remove it from + // escrow. + token::transfer( + state, + &pending_transfer.gas_fee.token, + &BRIDGE_POOL_ADDRESS, + relayer, + pending_transfer.gas_fee.amount, + )?; + state.delete(&key)?; _ = pending_keys.swap_remove(&key); _ = changed_keys.insert(key); @@ -452,12 +388,14 @@ where balance_key(&transfer.gas_fee.token, &transfer.gas_fee.payer); let pool_balance_key = balance_key(&transfer.gas_fee.token, &BRIDGE_POOL_ADDRESS); - update::amount(state, &payer_balance_key, |balance| { - balance.receive(&transfer.gas_fee.amount) - })?; - update::amount(state, &pool_balance_key, |balance| { - balance.spend(&transfer.gas_fee.amount) - })?; + + token::transfer( + state, + &transfer.gas_fee.token, + &BRIDGE_POOL_ADDRESS, + &transfer.gas_fee.payer, + transfer.gas_fee.amount, + )?; tracing::debug!(?transfer, "Refunded Bridge pool transfer fees"); _ = changed_keys.insert(payer_balance_key); @@ -478,26 +416,33 @@ where let native_erc20_addr = state .read(&bridge_storage::native_erc20_key())? .ok_or_else(|| eyre::eyre!("Could not read wNam key from storage"))?; - let (source, target) = if transfer.transfer.asset == native_erc20_addr { + let (source, target, token) = if transfer.transfer.asset + == native_erc20_addr + { let escrow_balance_key = balance_key(&state.in_mem().native_token, &BRIDGE_ADDRESS); let sender_balance_key = balance_key( &state.in_mem().native_token, &transfer.transfer.sender, ); - (escrow_balance_key, sender_balance_key) + ( + escrow_balance_key, + sender_balance_key, + state.in_mem().native_token.clone(), + ) } else { let token = transfer.token_address(); let escrow_balance_key = balance_key(&token, &BRIDGE_POOL_ADDRESS); let sender_balance_key = balance_key(&token, &transfer.transfer.sender); - (escrow_balance_key, sender_balance_key) + (escrow_balance_key, sender_balance_key, token) }; - update::amount(state, &source, |balance| { - balance.spend(&transfer.transfer.amount) - })?; - update::amount(state, &target, |balance| { - balance.receive(&transfer.transfer.amount) - })?; + token::transfer( + state, + &token, + &BRIDGE_POOL_ADDRESS, + &transfer.transfer.sender, + transfer.transfer.amount, + )?; tracing::debug!(?transfer, "Refunded Bridge pool transferred assets"); _ = changed_keys.insert(source); @@ -533,26 +478,23 @@ where unreachable!("Attempted to mint wNAM NUTs!"); } let supply_key = minted_balance_key(&token); - update::amount(state, &supply_key, |supply| { - supply.receive(&transfer.transfer.amount) - })?; + increment_total_supply(state, &token, transfer.transfer.amount)?; _ = changed_keys.insert(supply_key); tracing::debug!(?transfer, "Updated wrapped NAM supply"); return Ok(changed_keys); } // other asset kinds must be burned + burn_tokens( + state, + &token, + &BRIDGE_POOL_ADDRESS, + transfer.transfer.amount, + )?; let escrow_balance_key = balance_key(&token, &BRIDGE_POOL_ADDRESS); - update::amount(state, &escrow_balance_key, |balance| { - balance.spend(&transfer.transfer.amount) - })?; - _ = changed_keys.insert(escrow_balance_key); - let supply_key = minted_balance_key(&token); - update::amount(state, &supply_key, |supply| { - supply.spend(&transfer.transfer.amount) - })?; + _ = changed_keys.insert(escrow_balance_key); _ = changed_keys.insert(supply_key); tracing::debug!(?transfer, "Burned wrapped ERC20 tokens"); @@ -574,6 +516,7 @@ mod tests { use namada_core::{address, eth_bridge_pool}; use namada_parameters::{update_epoch_parameter, EpochDuration}; use namada_state::testing::TestState; + use token::increment_balance; use super::*; use crate::storage::bridge_pool::get_pending_key; @@ -731,12 +674,12 @@ mod tests { let payer_key = balance_key(&transfer.gas_fee.token, &payer); let payer_balance = Amount::from(0); state.write(&payer_key, payer_balance).expect("Test failed"); - let escrow_key = - balance_key(&transfer.gas_fee.token, &BRIDGE_POOL_ADDRESS); - update::amount(state, &escrow_key, |balance| { - let gas_fee = Amount::from_u64(1); - balance.receive(&gas_fee) - }) + increment_balance( + state, + &transfer.gas_fee.token, + &BRIDGE_POOL_ADDRESS, + Amount::from_u64(1), + ) .expect("Test failed"); if transfer.transfer.asset == wnam() { @@ -763,10 +706,8 @@ mod tests { state .write(&escrow_key, escrow_balance) .expect("Test failed"); - update::amount(state, &minted_balance_key(&token), |supply| { - supply.receive(&transfer.transfer.amount) - }) - .expect("Test failed"); + increment_total_supply(state, &token, transfer.transfer.amount) + .expect("Test failed"); }; } } diff --git a/crates/ethereum_bridge/src/protocol/transactions/read.rs b/crates/ethereum_bridge/src/protocol/transactions/read.rs index 7c94dd5939..1841676a6a 100644 --- a/crates/ethereum_bridge/src/protocol/transactions/read.rs +++ b/crates/ethereum_bridge/src/protocol/transactions/read.rs @@ -2,11 +2,13 @@ use borsh::BorshDeserialize; use eyre::{eyre, Result}; use namada_core::storage; +#[cfg(test)] use namada_core::token::Amount; use namada_state::{DBIter, StorageHasher, WlState, DB}; use namada_storage::StorageRead; /// Returns the stored Amount, or 0 if not stored +#[cfg(test)] pub(super) fn amount_or_default( state: &WlState, key: &storage::Key, diff --git a/crates/ethereum_bridge/src/protocol/transactions/update.rs b/crates/ethereum_bridge/src/protocol/transactions/update.rs index 800ede9cad..c6113d82c0 100644 --- a/crates/ethereum_bridge/src/protocol/transactions/update.rs +++ b/crates/ethereum_bridge/src/protocol/transactions/update.rs @@ -3,26 +3,9 @@ use eyre::Result; use namada_core::borsh::{BorshDeserialize, BorshSerialize}; use namada_core::hash::StorageHasher; use namada_core::storage; -use namada_core::token::{Amount, AmountError}; use namada_state::{DBIter, WlState, DB}; use namada_storage::StorageWrite; -/// Reads the `Amount` from key, applies update then writes it back -pub fn amount( - state: &mut WlState, - key: &storage::Key, - update: impl FnOnce(&mut Amount) -> Result<(), AmountError>, -) -> Result -where - D: 'static + DB + for<'iter> DBIter<'iter> + Sync, - H: 'static + StorageHasher + Sync, -{ - let mut amount = super::read::amount_or_default(state, key)?; - update(&mut amount)?; - state.write(key, amount)?; - Ok(amount) -} - #[allow(dead_code)] /// Reads an arbitrary value, applies update then writes it back pub fn value( diff --git a/crates/trans_token/src/storage.rs b/crates/trans_token/src/storage.rs index 4324f0be29..04048b466c 100644 --- a/crates/trans_token/src/storage.rs +++ b/crates/trans_token/src/storage.rs @@ -1,8 +1,9 @@ use namada_core::address::{Address, InternalAddress}; use namada_core::hints; -use namada_core::token::{self, Amount, DenominatedAmount}; +use namada_core::token::{self, Amount, AmountError, DenominatedAmount}; use namada_storage as storage; use namada_storage::{StorageRead, StorageWrite}; +use storage::ResultExt; use crate::storage_key::*; @@ -32,6 +33,59 @@ where Ok(balance) } +/// Update the balance of a given token and owner. +pub fn update_balance( + storage: &mut S, + token: &Address, + owner: &Address, + f: F, +) -> storage::Result<()> +where + S: StorageRead + StorageWrite, + F: FnOnce(token::Amount) -> storage::Result, +{ + let key = balance_key(token, owner); + let balance = storage.read::(&key)?.unwrap_or_default(); + let new_balance = f(balance)?; + storage.write(&key, new_balance) +} + +/// Increment the balance of a given token and owner. +pub fn increment_balance( + storage: &mut S, + token: &Address, + owner: &Address, + amount: token::Amount, +) -> storage::Result<()> +where + S: StorageRead + StorageWrite, +{ + update_balance(storage, token, owner, |cur_amount| { + cur_amount + .checked_add(amount) + .ok_or(AmountError::Overflow) + .into_storage_result() + }) +} + +/// Decrement the balance of a given token and owner. +pub fn decrement_balance( + storage: &mut S, + token: &Address, + owner: &Address, + amount: token::Amount, +) -> storage::Result<()> +where + S: StorageRead + StorageWrite, +{ + update_balance(storage, token, owner, |cur_amount| { + cur_amount + .checked_sub(amount) + .ok_or(AmountError::Insufficient) + .into_storage_result() + }) +} + /// Read the total network supply of a given token. pub fn read_total_supply( storage: &S, @@ -41,8 +95,58 @@ where S: StorageRead, { let key = minted_balance_key(token); - let balance = storage.read::(&key)?.unwrap_or_default(); - Ok(balance) + let total_supply = storage.read::(&key)?.unwrap_or_default(); + Ok(total_supply) +} + +/// Update the total network supply of a given token. +pub fn update_total_supply( + storage: &mut S, + token: &Address, + f: F, +) -> storage::Result<()> +where + S: StorageRead + StorageWrite, + F: FnOnce(token::Amount) -> storage::Result, +{ + let key = minted_balance_key(token); + let total_supply = storage.read::(&key)?.unwrap_or_default(); + let new_supply = f(total_supply)?; + storage.write(&key, new_supply) +} + +/// Increment the total network supply of a given token. +pub fn increment_total_supply( + storage: &mut S, + token: &Address, + amount: token::Amount, +) -> storage::Result<()> +where + S: StorageRead + StorageWrite, +{ + update_total_supply(storage, token, |cur_supply| { + cur_supply + .checked_add(amount) + .ok_or(AmountError::Overflow) + .into_storage_result() + }) +} + +/// Decrement the total network supply of a given token. +pub fn decrement_total_supply( + storage: &mut S, + token: &Address, + amount: token::Amount, +) -> storage::Result<()> +where + S: StorageRead + StorageWrite, +{ + update_total_supply(storage, token, |cur_supply| { + cur_supply + .checked_sub(amount) + .ok_or(AmountError::Insufficient) + .into_storage_result() + }) } /// Get the effective circulating total supply of native tokens. @@ -163,22 +267,11 @@ pub fn credit_tokens( where S: StorageRead + StorageWrite, { - let balance_key = balance_key(token, dest); - let cur_balance = read_balance(storage, token, dest)?; - let new_balance = cur_balance - .checked_add(amount) - .ok_or_else(|| storage::Error::new_const("Token balance overflow"))?; - - let total_supply_key = minted_balance_key(token); - let cur_supply = storage - .read::(&total_supply_key)? - .unwrap_or_default(); - let new_supply = cur_supply.checked_add(amount).ok_or_else(|| { - storage::Error::new_const("Token total supply overflow") - })?; + // Increment the destination balance + increment_balance(storage, token, dest, amount)?; - storage.write(&balance_key, new_balance)?; - storage.write(&total_supply_key, new_supply) + // Increment the total supply + increment_total_supply(storage, token, amount) } /// Burn a specified amount of tokens from some address. If the burn amount is @@ -205,15 +298,8 @@ where source_balance }; - let old_total_supply = read_total_supply(storage, token)?; - let new_total_supply = old_total_supply - .checked_sub(amount_to_burn) - .ok_or_else(|| { - storage::Error::new_const("Total token supply underflowed") - })?; - - let total_supply_key = minted_balance_key(token); - storage.write(&total_supply_key, new_total_supply) + // Decrement the total supply + decrement_total_supply(storage, token, amount_to_burn) } /// Add denomination info if it exists in storage.