Tracks issued_assets map in non-finalized state

zebra-state/src/service/check.rs

@@ -28,6 +28,7 @@ use crate::service::non_finalized_state::Chain;
 
 pub(crate) mod anchors;
 pub(crate) mod difficulty;
+pub(crate) mod issuance;
 pub(crate) mod nullifier;
 pub(crate) mod utxo;
 

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

@@ -0,0 +1,41 @@
+use std::{collections::HashMap, sync::Arc};
+
+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>(
+    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);
+    }
+}
+
+#[tracing::instrument(skip(issued_assets, transaction))]
+pub(crate) fn add_to_non_finalized_chain(
+    issued_assets: &mut HashMap<AssetBase, AssetState>,
+    transaction: &Arc<Transaction>,
+) -> Result<(), ValidateContextError> {
+    modify_non_finalized_chain(issued_assets, transaction, |prev, change| {
+        prev.with_change(change)
+    });
+
+    Ok(())
+}
+
+#[tracing::instrument(skip(issued_assets, transaction))]
+pub(crate) fn remove_from_non_finalized_chain(
+    issued_assets: &mut HashMap<AssetBase, AssetState>,
+    transaction: &Arc<Transaction>,
+) {
+    modify_non_finalized_chain(issued_assets, transaction, |prev, change| {
+        prev.with_revert(change)
+    });
+}

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

@@ -5,6 +5,8 @@
 //! [`crate::constants::state_database_format_version_in_code()`] must be incremented
 //! each time the database format (column, serialization, etc) changes.
 
+use std::sync::Arc;
+
 use bincode::Options;
 
 use zebra_chain::{
@@ -13,6 +15,7 @@ use zebra_chain::{
     orchard_zsa::{self, BurnItem},
     sapling, sprout,
     subtree::{NoteCommitmentSubtreeData, NoteCommitmentSubtreeIndex},
+    transaction::Transaction,
 };
 
 use crate::service::finalized_state::disk_format::{FromDisk, IntoDisk};
@@ -40,6 +43,17 @@ impl AssetState {
         }
     }
 
+    /// 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 from_note(is_finalized: bool, note: orchard_zsa::Note) -> (AssetBase, Self) {
         (
             note.asset(),
@@ -72,6 +86,15 @@ impl AssetState {
     pub fn from_burns(burns: Vec<BurnItem>) -> impl Iterator<Item = (AssetBase, Self)> {
         burns.into_iter().map(Self::from_burn)
     }
+
+    pub fn from_transaction(
+        transaction: &Arc<Transaction>,
+    ) -> impl Iterator<Item = (AssetBase, Self)> + '_ {
+        transaction
+            .issue_actions()
+            .flat_map(|action| AssetState::from_notes(action.is_finalized(), action.notes()))
+            .chain(AssetState::from_burns(transaction.burns()))
+    }
 }
 
 impl IntoDisk for AssetState {

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

@@ -518,27 +518,21 @@ impl DiskWriteBatch {
 
         let updated_issued_assets = transaction
             .iter()
-            .map(|tx| (tx.issue_actions(), tx.burns()))
+            .flat_map(AssetState::from_transaction)
             .fold(
                 HashMap::new(),
-                |mut issued_assets: HashMap<orchard::AssetBase, AssetState>, (actions, burns)| {
-                    for (asset_base, asset_state_change) in actions
-                        .flat_map(|action| {
-                            AssetState::from_notes(action.is_finalized(), action.notes())
-                        })
-                        .chain(AssetState::from_burns(burns))
-                    {
-                        if let Some(asset_state) = issued_assets.get_mut(&asset_base) {
-                            assert!(!asset_state.is_finalized);
-                            *asset_state = asset_state.with_change(asset_state_change);
-                        } else {
-                            let prev_state = issued_assets_cf.zs_get(&asset_base);
-                            let new_state = prev_state.map_or(asset_state_change, |prev| {
-                                prev.with_change(asset_state_change)
-                            });
-                            issued_assets.insert(asset_base, new_state);
-                        };
-                    }
+                |mut issued_assets: HashMap<orchard::AssetBase, AssetState>,
+                 (asset_base, asset_state_change)| {
+                    if let Some(asset_state) = issued_assets.get_mut(&asset_base) {
+                        assert!(!asset_state.is_finalized);
+                        *asset_state = asset_state.with_change(asset_state_change);
+                    } else {
+                        let prev_state = issued_assets_cf.zs_get(&asset_base);
+                        let new_state = prev_state.map_or(asset_state_change, |prev| {
+                            prev.with_change(asset_state_change)
+                        });
+                        issued_assets.insert(asset_base, new_state);
+                    };
 
                     issued_assets
                 },

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

@@ -15,22 +15,26 @@ use zebra_chain::{
     amount::{Amount, NegativeAllowed, NonNegative},
     block::{self, Height},
     history_tree::HistoryTree,
-    orchard,
+    orchard::{self, AssetBase},
     parallel::tree::NoteCommitmentTrees,
     parameters::Network,
     primitives::Groth16Proof,
     sapling, sprout,
     subtree::{NoteCommitmentSubtree, NoteCommitmentSubtreeData, NoteCommitmentSubtreeIndex},
-    transaction::Transaction::*,
-    transaction::{self, Transaction},
+    transaction::{
+        self,
+        Transaction::{self, *},
+    },
     transparent,
     value_balance::ValueBalance,
     work::difficulty::PartialCumulativeWork,
 };
 
 use crate::{
-    request::Treestate, service::check, ContextuallyVerifiedBlock, HashOrHeight, OutputLocation,
-    TransactionLocation, ValidateContextError,
+    request::Treestate,
+    service::{check, finalized_state::AssetState},
+    ContextuallyVerifiedBlock, HashOrHeight, OutputLocation, TransactionLocation,
+    ValidateContextError,
 };
 
 use self::index::TransparentTransfers;
@@ -174,6 +178,10 @@ pub struct ChainInner {
     pub(crate) orchard_subtrees:
         BTreeMap<NoteCommitmentSubtreeIndex, NoteCommitmentSubtreeData<orchard::tree::Node>>,
 
+    /// A map of `issued_assets`,
+    // TODO: Add reference to ZIP
+    pub(crate) orchard_issued_assets: HashMap<AssetBase, AssetState>,
+
     // Nullifiers
     //
     /// The Sprout nullifiers revealed by `blocks`.
@@ -220,30 +228,8 @@ impl Chain {
         finalized_tip_chain_value_pools: ValueBalance<NonNegative>,
     ) -> Self {
         let inner = ChainInner {
-            blocks: Default::default(),
-            height_by_hash: Default::default(),
-            tx_loc_by_hash: Default::default(),
-            created_utxos: Default::default(),
-            spent_utxos: Default::default(),
-            sprout_anchors: MultiSet::new(),
-            sprout_anchors_by_height: Default::default(),
-            sprout_trees_by_anchor: Default::default(),
-            sprout_trees_by_height: Default::default(),
-            sapling_anchors: MultiSet::new(),
-            sapling_anchors_by_height: Default::default(),
-            sapling_trees_by_height: Default::default(),
-            sapling_subtrees: Default::default(),
-            orchard_anchors: MultiSet::new(),
-            orchard_anchors_by_height: Default::default(),
-            orchard_trees_by_height: Default::default(),
-            orchard_subtrees: Default::default(),
-            sprout_nullifiers: Default::default(),
-            sapling_nullifiers: Default::default(),
-            orchard_nullifiers: Default::default(),
-            partial_transparent_transfers: Default::default(),
-            partial_cumulative_work: Default::default(),
-            history_trees_by_height: Default::default(),
             chain_value_pools: finalized_tip_chain_value_pools,
+            ..ChainInner::default()
         };
 
         let mut chain = Self {
@@ -1507,7 +1493,7 @@ impl Chain {
                     &None,
                     &None,
                     #[cfg(feature ="tx-v6")]
-                    &None
+                    &None,
                     ),
                 V5 {
                     inputs,
@@ -1588,7 +1574,10 @@ impl Chain {
             self.update_chain_tip_with(sapling_shielded_data_shared_anchor)?;
             self.update_chain_tip_with(orchard_shielded_data_vanilla)?;
             #[cfg(feature = "tx-v6")]
-            self.update_chain_tip_with(orchard_shielded_data_zsa)?;
+            {
+                self.update_chain_tip_with(orchard_shielded_data_zsa)?;
+                self.update_chain_tip_with(transaction)?;
+            }
         }
 
         // update the chain value pool balances
@@ -1763,6 +1752,9 @@ impl UpdateWith<ContextuallyVerifiedBlock> for Chain {
             self.revert_chain_with(sapling_shielded_data_per_spend_anchor, position);
             self.revert_chain_with(sapling_shielded_data_shared_anchor, position);
             self.revert_chain_with(orchard_shielded_data, position);
+
+            #[cfg(feature = "tx-v6")]
+            self.revert_chain_with(transaction, position);
         }
 
         // TODO: move these to the shielded UpdateWith.revert...()?
@@ -2117,6 +2109,25 @@ impl<V: orchard::OrchardFlavorExt> UpdateWith<Option<orchard::ShieldedData<V>>>
     }
 }
 
+#[cfg(feature = "tx-v6")]
+impl UpdateWith<Arc<Transaction>> for Chain {
+    #[instrument(skip(self, transaction))]
+    fn update_chain_tip_with(
+        &mut self,
+        transaction: &Arc<Transaction>,
+    ) -> Result<(), ValidateContextError> {
+        check::issuance::add_to_non_finalized_chain(&mut self.orchard_issued_assets, transaction)
+    }
+
+    #[instrument(skip(self, transaction))]
+    fn revert_chain_with(&mut self, transaction: &Arc<Transaction>, _position: RevertPosition) {
+        check::issuance::remove_from_non_finalized_chain(
+            &mut self.orchard_issued_assets,
+            transaction,
+        );
+    }
+}
+
 impl UpdateWith<ValueBalance<NegativeAllowed>> for Chain {
     fn update_chain_tip_with(
         &mut self,

zebra-state/src/service/write.rs

@@ -141,6 +141,7 @@ pub fn write_blocks_from_channels(
 ) {
     let mut last_zebra_mined_log_height = None;
     let mut prev_finalized_note_commitment_trees = None;
+    // TODO: Add a `prev_issued_assets` here.
 
     // Write all the finalized blocks sent by the state,
     // until the state closes the finalized block channel's sender.