Adds issued assets map to non-finalized chains

zebra-chain/src/orchard_zsa/asset_state.rs

@@ -5,7 +5,7 @@ use std::{collections::HashMap, sync::Arc};
 use orchard::issuance::IssueAction;
 pub use orchard::note::AssetBase;
 
-use crate::block::Block;
+use crate::transaction::Transaction;
 
 use super::BurnItem;
 
@@ -30,9 +30,9 @@ pub struct AssetStateChange {
 }
 
 impl AssetState {
-    /// Updates and returns self with the provided [`AssetStateChange`] if the change is valid, or
-    /// returns None otherwise.
-    pub fn with_change(mut self, change: AssetStateChange) -> Option<Self> {
+    /// Updates and returns self with the provided [`AssetStateChange`] if
+    /// the change is valid, or returns None otherwise.
+    pub fn apply_change(mut self, change: AssetStateChange) -> Option<Self> {
         if self.is_finalized {
             return None;
         }
@@ -41,6 +41,15 @@ impl AssetState {
         self.total_supply = self.total_supply.checked_add_signed(change.supply_change)?;
         Some(self)
     }
+
+    /// Reverts the provided [`AssetStateChange`].
+    pub fn revert_change(&mut self, change: AssetStateChange) {
+        self.is_finalized &= !change.is_finalized;
+        self.total_supply = self
+            .total_supply
+            .checked_add_signed(-change.supply_change)
+            .expect("reversions must not overflow");
+    }
 }
 
 impl From<HashMap<AssetBase, AssetState>> for IssuedAssets {
@@ -108,6 +117,11 @@ impl IssuedAssets {
         Self(HashMap::new())
     }
 
+    /// Returns an iterator of the inner HashMap.
+    pub fn iter(&self) -> impl Iterator<Item = (&AssetBase, &AssetState)> {
+        self.0.iter()
+    }
+
     fn update<'a>(&mut self, issued_assets: impl Iterator<Item = (AssetBase, AssetState)> + 'a) {
         for (asset_base, asset_state) in issued_assets {
             self.0.insert(asset_base, asset_state);
@@ -140,15 +154,15 @@ impl IssuedAssetsChange {
         }
     }
 
-    /// Accepts a reference to an [`Arc<Block>`].
+    /// Accepts a slice of [`Arc<Transaction>`]s.
     ///
     /// Returns a tuple, ([`IssuedAssetsChange`], [`IssuedAssetsChange`]), where
     /// the first item is from burns and the second one is for issuance.
-    pub fn from_block(block: &Arc<Block>) -> (Self, Self) {
+    pub fn from_transactions(transactions: &[Arc<Transaction>]) -> (Self, Self) {
         let mut burn_change = Self::new();
         let mut issuance_change = Self::new();
 
-        for transaction in &block.transactions {
+        for transaction in transactions {
             burn_change.update(AssetStateChange::from_burns(transaction.orchard_burns()));
             issuance_change.update(AssetStateChange::from_issue_actions(
                 transaction.orchard_issue_actions(),
@@ -158,6 +172,23 @@ impl IssuedAssetsChange {
         (burn_change, issuance_change)
     }
 
+    /// Accepts a slice of [`Arc<Transaction>`]s.
+    ///
+    /// Returns an [`IssuedAssetsChange`] representing all of the changes to the issued assets
+    /// map that should be applied for the provided transactions.
+    pub fn combined_from_transactions(transactions: &[Arc<Transaction>]) -> Self {
+        let mut issued_assets_change = Self::new();
+
+        for transaction in transactions {
+            issued_assets_change.update(AssetStateChange::from_burns(transaction.orchard_burns()));
+            issued_assets_change.update(AssetStateChange::from_issue_actions(
+                transaction.orchard_issue_actions(),
+            ));
+        }
+
+        issued_assets_change
+    }
+
     /// Consumes self and accepts a closure for looking up previous asset states.
     ///
     /// Applies changes in self to the previous asset state.
@@ -170,7 +201,7 @@ impl IssuedAssetsChange {
             (
                 asset_base,
                 f(asset_base)
-                    .with_change(change)
+                    .apply_change(change)
                     .expect("must be valid change"),
             )
         }));

zebra-state/src/arbitrary.rs

@@ -32,7 +32,7 @@ impl Prepare for Arc<Block> {
         let transaction_hashes: Arc<[_]> = block.transactions.iter().map(|tx| tx.hash()).collect();
         let new_outputs =
             transparent::new_ordered_outputs_with_height(&block, height, &transaction_hashes);
-        let (burns, issuance) = IssuedAssetsChange::from_block(&block);
+        let (burns, issuance) = IssuedAssetsChange::from_transactions(&block.transactions);
 
         SemanticallyVerifiedBlock {
             block,

zebra-state/src/request.rs

@@ -313,7 +313,7 @@ pub struct FinalizedBlock {
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub enum IssuedAssetsOrChanges {
     /// A map of updated issued assets.
-    State(IssuedAssets),
+    Updated(IssuedAssets),
 
     /// A map of changes to apply to the issued assets map.
     Change(IssuedAssetsChange),
@@ -499,7 +499,7 @@ impl SemanticallyVerifiedBlock {
             .expect("semantically verified block should have a coinbase height");
         let transaction_hashes: Arc<[_]> = block.transactions.iter().map(|tx| tx.hash()).collect();
         let new_outputs = transparent::new_ordered_outputs(&block, &transaction_hashes);
-        let (burns, issuance) = IssuedAssetsChange::from_block(&block);
+        let (burns, issuance) = IssuedAssetsChange::from_transactions(&block.transactions);
 
         Self {
             block,
@@ -537,7 +537,7 @@ impl From<Arc<Block>> for SemanticallyVerifiedBlock {
             .expect("semantically verified block should have a coinbase height");
         let transaction_hashes: Arc<[_]> = block.transactions.iter().map(|tx| tx.hash()).collect();
         let new_outputs = transparent::new_ordered_outputs(&block, &transaction_hashes);
-        let (burns, issuance) = IssuedAssetsChange::from_block(&block);
+        let (burns, issuance) = IssuedAssetsChange::from_transactions(&block.transactions);
 
         Self {
             block,
@@ -556,8 +556,6 @@ impl From<Arc<Block>> for SemanticallyVerifiedBlock {
 
 impl From<ContextuallyVerifiedBlock> for SemanticallyVerifiedBlock {
     fn from(valid: ContextuallyVerifiedBlock) -> Self {
-        let (burns, issuance) = IssuedAssetsChange::from_block(&valid.block);
-
         Self {
             block: valid.block,
             hash: valid.hash,
@@ -571,29 +569,21 @@ impl From<ContextuallyVerifiedBlock> for SemanticallyVerifiedBlock {
                     .constrain::<NonNegative>()
                     .expect("deferred balance in a block must me non-negative"),
             ),
-            issued_assets_changes: IssuedAssetsOrChanges::BurnAndIssuanceChanges {
-                burns,
-                issuance,
-            },
+            issued_assets_changes: IssuedAssetsOrChanges::Updated(valid.issued_assets),
         }
     }
 }
 
 impl From<FinalizedBlock> for SemanticallyVerifiedBlock {
     fn from(finalized: FinalizedBlock) -> Self {
-        let (burns, issuance) = IssuedAssetsChange::from_block(&finalized.block);
-
         Self {
             block: finalized.block,
             hash: finalized.hash,
             height: finalized.height,
             new_outputs: finalized.new_outputs,
             transaction_hashes: finalized.transaction_hashes,
             deferred_balance: finalized.deferred_balance,
-            issued_assets_changes: IssuedAssetsOrChanges::BurnAndIssuanceChanges {
-                burns,
-                issuance,
-            },
+            issued_assets_changes: finalized.issued_assets,
         }
     }
 }

zebra-state/src/service/check/issuance.rs

@@ -27,7 +27,7 @@ pub fn valid_burns_and_issuance(
             .issued_asset(&asset_base)
             .or_else(|| finalized_state.issued_asset(&asset_base))
             .ok_or(ValidateContextError::InvalidBurn)?
-            .with_change(burn_change)
+            .apply_change(burn_change)
             .ok_or(ValidateContextError::InvalidBurn)?;
 
         issued_assets
@@ -49,7 +49,7 @@ pub fn valid_burns_and_issuance(
         let _ = issued_assets.insert(
             asset_base,
             asset_state
-                .with_change(issuance_change)
+                .apply_change(issuance_change)
                 .ok_or(ValidateContextError::InvalidIssuance)?,
         );
     }

zebra-state/src/service/finalized_state/zebra_db/block/tests/vectors.rs

@@ -130,7 +130,8 @@ fn test_block_db_round_trip_with(
                 .collect();
             let new_outputs =
                 new_ordered_outputs_with_height(&original_block, Height(0), &transaction_hashes);
-            let (burns, issuance) = IssuedAssetsChange::from_block(&original_block);
+            let (burns, issuance) =
+                IssuedAssetsChange::from_transactions(&original_block.transactions);
 
             CheckpointVerifiedBlock(SemanticallyVerifiedBlock {
                 block: original_block.clone(),

zebra-state/src/service/finalized_state/zebra_db/shielded.rs

@@ -520,7 +520,7 @@ impl DiskWriteBatch {
         let mut batch = zebra_db.issued_assets_cf().with_batch_for_writing(self);
 
         let updated_issued_assets = match issued_assets_or_changes.clone().combine() {
-            IssuedAssetsOrChanges::State(issued_assets) => issued_assets,
+            IssuedAssetsOrChanges::Updated(issued_assets) => issued_assets,
             IssuedAssetsOrChanges::Change(issued_assets_change) => issued_assets_change
                 .apply_with(|asset_base| zebra_db.issued_asset(&asset_base).unwrap_or_default()),
             IssuedAssetsOrChanges::BurnAndIssuanceChanges { .. } => {

zebra-state/src/service/non_finalized_state/chain.rs

@@ -16,7 +16,7 @@ use zebra_chain::{
     block::{self, Height},
     history_tree::HistoryTree,
     orchard,
-    orchard_zsa::{AssetBase, AssetState},
+    orchard_zsa::{AssetBase, AssetState, IssuedAssets, IssuedAssetsChange},
     parallel::tree::NoteCommitmentTrees,
     parameters::Network,
     primitives::Groth16Proof,
@@ -952,6 +952,36 @@ impl Chain {
         self.issued_assets.get(asset_base).cloned()
     }
 
+    /// Remove the History tree index at `height`.
+    fn revert_issued_assets(
+        &mut self,
+        position: RevertPosition,
+        issued_assets: &IssuedAssets,
+        transactions: &[Arc<Transaction>],
+    ) {
+        if position == RevertPosition::Root {
+            trace!(?position, "removing unmodified issued assets");
+            for (asset_base, &asset_state) in issued_assets.iter() {
+                if self
+                    .issued_asset(asset_base)
+                    .expect("issued assets for chain should include those in all blocks")
+                    == asset_state
+                {
+                    self.issued_assets.remove(asset_base);
+                }
+            }
+        } else {
+            trace!(?position, "reverting changes to issued assets");
+            for (asset_base, change) in IssuedAssetsChange::combined_from_transactions(transactions)
+            {
+                self.issued_assets
+                    .entry(asset_base)
+                    .or_default()
+                    .revert_change(change);
+            }
+        }
+    }
+
     /// Adds the Orchard `tree` to the tree and anchor indexes at `height`.
     ///
     /// `height` can be either:
@@ -1454,6 +1484,9 @@ impl Chain {
 
         self.add_history_tree(height, history_tree);
 
+        self.issued_assets
+            .extend(contextually_valid.issued_assets.clone());
+
         Ok(())
     }
 
@@ -1682,6 +1715,7 @@ impl UpdateWith<ContextuallyVerifiedBlock> for Chain {
             spent_outputs,
             transaction_hashes,
             chain_value_pool_change,
+            issued_assets,
         ) = (
             contextually_valid.block.as_ref(),
             contextually_valid.hash,
@@ -1690,6 +1724,7 @@ impl UpdateWith<ContextuallyVerifiedBlock> for Chain {
             &contextually_valid.spent_outputs,
             &contextually_valid.transaction_hashes,
             &contextually_valid.chain_value_pool_change,
+            &contextually_valid.issued_assets,
         );
 
         // remove the blocks hash from `height_by_hash`
@@ -1788,6 +1823,9 @@ impl UpdateWith<ContextuallyVerifiedBlock> for Chain {
         // TODO: move this to the history tree UpdateWith.revert...()?
         self.remove_history_tree(position, height);
 
+        // revert the issued assets map, if needed
+        self.revert_issued_assets(position, issued_assets, &block.transactions);
+
         // revert the chain value pool balances, if needed
         self.revert_chain_with(chain_value_pool_change, position);
     }