diff --git a/crates/trans_token/src/storage.rs b/crates/trans_token/src/storage.rs index 4324f0be29..a8cce6ebe3 100644 --- a/crates/trans_token/src/storage.rs +++ b/crates/trans_token/src/storage.rs @@ -32,6 +32,23 @@ 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) -> token::Amount, +{ + let key = balance_key(token, owner); + let balance = storage.read::(&key)?.unwrap_or_default(); + let new_balance = f(balance); + storage.write(&key, new_balance) +} + /// Read the total network supply of a given token. pub fn read_total_supply( storage: &S, @@ -41,8 +58,24 @@ 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) -> token::Amount, +{ + 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) } /// Get the effective circulating total supply of native tokens. @@ -163,22 +196,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") - })?; + // Update the destination balance + update_balance(storage, token, dest, |cur_amount| cur_amount + amount)?; - storage.write(&balance_key, new_balance)?; - storage.write(&total_supply_key, new_supply) + // Update the total supply + update_total_supply(storage, token, |cur_supply| cur_supply + amount) } /// Burn a specified amount of tokens from some address. If the burn amount is @@ -205,15 +227,12 @@ 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) + // Update total supply + update_total_supply(storage, token, |cur_supply| { + cur_supply + .checked_sub(amount_to_burn) + .expect("Total token supply underflowed") + }) } /// Add denomination info if it exists in storage.