wip

ZcashFoundation · Nov 11, 2024 · c7da266 · c7da266
1 parent 0d0f904
commit c7da266
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 34 deletions.
diff --git 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<F: Fn(AssetState, AssetState) -> AssetState>(
+fn modify_non_finalized_chain<F: Fn(&mut AssetState, AssetState)>(
     issued_assets: &mut HashMap<AssetBase, AssetState>,
     transaction: &Arc<Transaction>,
     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<Transaction>,
 ) -> 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<Transaction>,
 ) {
     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
@@ -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<Item = (AssetBase, Self)> + '_ {
@@ -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<BurnItem>) -> impl Iterator<Item = (AssetBase, Self)> {
+    fn from_burns(burns: Vec<BurnItem>) -> impl Iterator<Item = (AssetBase, Self)> {
         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
@@ -523,13 +523,10 @@ impl DiskWriteBatch {
                 HashMap::new(),
                 |mut issued_assets: HashMap<orchard::AssetBase, AssetState>,
                  (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
                 },
             );