diff --git a/zebra-state/src/service/check/issuance.rs b/zebra-state/src/service/check/issuance.rs index d94b16896aa..a67b2069da1 100644 --- a/zebra-state/src/service/check/issuance.rs +++ b/zebra-state/src/service/check/issuance.rs @@ -5,16 +5,16 @@ use zebra_chain::{orchard::AssetBase, transaction::Transaction}; use crate::{service::finalized_state::AssetState, ValidateContextError}; #[tracing::instrument(skip(issued_assets, transaction, modifier))] -fn modify_non_finalized_chain AssetState>( +fn modify_non_finalized_chain( issued_assets: &mut HashMap, transaction: &Arc, modifier: F, ) { for (asset_base, change) in AssetState::from_transaction(transaction) { - let new_state = issued_assets - .get(&asset_base) - .map_or(change, |&prev| modifier(prev, change)); - issued_assets.insert(asset_base, new_state); + issued_assets + .entry(asset_base) + .and_modify(|prev| modifier(prev, change)) + .or_insert(change); } } @@ -24,7 +24,7 @@ pub(crate) fn add_to_non_finalized_chain( transaction: &Arc, ) -> Result<(), ValidateContextError> { modify_non_finalized_chain(issued_assets, transaction, |prev, change| { - prev.with_change(change) + prev.apply_change(change); }); Ok(()) @@ -36,6 +36,6 @@ pub(crate) fn remove_from_non_finalized_chain( transaction: &Arc, ) { modify_non_finalized_chain(issued_assets, transaction, |prev, change| { - prev.with_revert(change) + prev.revert_change(change); }); } diff --git a/zebra-state/src/service/finalized_state/disk_format/shielded.rs b/zebra-state/src/service/finalized_state/disk_format/shielded.rs index ef9acd16537..a903439b90f 100644 --- a/zebra-state/src/service/finalized_state/disk_format/shielded.rs +++ b/zebra-state/src/service/finalized_state/disk_format/shielded.rs @@ -33,28 +33,24 @@ pub struct AssetState { impl AssetState { /// Adds partial asset states - pub fn with_change(&self, change: Self) -> Self { - Self { - is_finalized: self.is_finalized || change.is_finalized, - total_supply: self - .total_supply - .checked_add(change.total_supply) - .expect("asset supply sum should not exceed u64 size"), - } + pub fn apply_change(&mut self, change: Self) { + self.is_finalized |= change.is_finalized; + self.total_supply = self + .total_supply + .checked_add(change.total_supply) + .expect("asset supply sum should not exceed u64 size"); } /// Adds partial asset states - pub fn with_revert(&self, change: Self) -> Self { - Self { - is_finalized: self.is_finalized && !change.is_finalized, - total_supply: self - .total_supply - .checked_sub(change.total_supply) - .expect("change should be less than total supply"), - } + pub fn revert_change(&mut self, change: Self) { + self.is_finalized &= !change.is_finalized; + self.total_supply = self + .total_supply + .checked_sub(change.total_supply) + .expect("validated change should be less than total supply"); } - pub fn from_note(is_finalized: bool, note: orchard_zsa::Note) -> (AssetBase, Self) { + fn from_note(is_finalized: bool, note: orchard_zsa::Note) -> (AssetBase, Self) { ( note.asset(), Self { @@ -64,7 +60,7 @@ impl AssetState { ) } - pub fn from_notes( + fn from_notes( is_finalized: bool, notes: &[orchard_zsa::Note], ) -> impl Iterator + '_ { @@ -73,7 +69,7 @@ impl AssetState { .map(move |note| Self::from_note(is_finalized, *note)) } - pub fn from_burn(burn: BurnItem) -> (AssetBase, Self) { + fn from_burn(burn: BurnItem) -> (AssetBase, Self) { ( burn.asset(), Self { @@ -83,7 +79,7 @@ impl AssetState { ) } - pub fn from_burns(burns: Vec) -> impl Iterator { + fn from_burns(burns: Vec) -> impl Iterator { burns.into_iter().map(Self::from_burn) } diff --git a/zebra-state/src/service/finalized_state/zebra_db/shielded.rs b/zebra-state/src/service/finalized_state/zebra_db/shielded.rs index 2de873498f6..7992a1c8ff6 100644 --- a/zebra-state/src/service/finalized_state/zebra_db/shielded.rs +++ b/zebra-state/src/service/finalized_state/zebra_db/shielded.rs @@ -523,13 +523,10 @@ impl DiskWriteBatch { HashMap::new(), |mut issued_assets: HashMap, (asset_base, asset_state_change)| { - let new_state = issued_assets - .get(&asset_base) - .copied() - .or_else(|| issued_assets_cf.zs_get(&asset_base)) - .map(|prev| prev.with_change(asset_state_change)) - .unwrap_or(asset_state_change); - issued_assets.insert(asset_base, new_state); + issued_assets + .entry(asset_base) + .and_modify(|prev| prev.apply_change(asset_state_change)) + .or_insert(asset_state_change); issued_assets }, );