diff --git a/zebra-chain/src/orchard_zsa/asset_state.rs b/zebra-chain/src/orchard_zsa/asset_state.rs index e7413ad7604..242a8019722 100644 --- a/zebra-chain/src/orchard_zsa/asset_state.rs +++ b/zebra-chain/src/orchard_zsa/asset_state.rs @@ -1,6 +1,9 @@ //! Defines and implements the issued asset state types -use std::{collections::HashMap, sync::Arc}; +use std::{ + collections::{HashMap, HashSet}, + sync::Arc, +}; use orchard::issuance::IssueAction; pub use orchard::note::AssetBase; @@ -16,42 +19,108 @@ pub struct AssetState { pub is_finalized: bool, /// The circulating supply that has been issued for an asset. - pub total_supply: u128, + pub total_supply: u64, } /// A change to apply to the issued assets map. -// TODO: -// - Reference ZIP -// - Make this an enum of _either_ a finalization _or_ a supply change -// (applying the finalize flag for each issuance note will cause unexpected panics). +// TODO: Reference ZIP #[derive(Copy, Clone, Default, Debug, PartialEq, Eq, PartialOrd, Ord)] pub struct AssetStateChange { /// Whether the asset should be finalized such that no more of it can be issued. pub is_finalized: bool, - /// The change in supply from newly issued assets or burned assets. - pub supply_change: i128, + /// The change in supply from newly issued assets or burned assets, if any. + pub supply_change: SupplyChange, +} + +#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord)] +/// An asset supply change to apply to the issued assets map. +pub enum SupplyChange { + Issuance(u64), + Burn(u64), +} + +impl Default for SupplyChange { + fn default() -> Self { + Self::Issuance(0) + } +} + +impl SupplyChange { + fn apply_to(self, total_supply: u64) -> Option { + match self { + SupplyChange::Issuance(amount) => total_supply.checked_add(amount), + SupplyChange::Burn(amount) => total_supply.checked_sub(amount), + } + } + + fn as_i128(self) -> i128 { + match self { + SupplyChange::Issuance(amount) => i128::from(amount), + SupplyChange::Burn(amount) => -i128::from(amount), + } + } + + fn add(&mut self, rhs: Self) -> bool { + if let Some(result) = self + .as_i128() + .checked_add(rhs.as_i128()) + .and_then(|signed| match signed { + 0.. => signed.try_into().ok().map(Self::Issuance), + ..0 => signed.try_into().ok().map(Self::Burn), + }) + { + *self = result; + true + } else { + false + } + } +} + +impl std::ops::Neg for SupplyChange { + type Output = Self; + + fn neg(self) -> Self::Output { + match self { + Self::Issuance(amount) => Self::Burn(amount), + Self::Burn(amount) => Self::Issuance(amount), + } + } } impl AssetState { /// 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 { + pub fn apply_change(self, change: AssetStateChange) -> Option { + self.apply_finalization(change.is_finalized)? + .apply_supply_change(change.supply_change) + } + + fn apply_finalization(mut self, is_finalized: bool) -> Option { if self.is_finalized { - return None; + None + } else { + self.is_finalized = is_finalized; + Some(self) } + } - self.is_finalized |= change.is_finalized; - self.total_supply = self.total_supply.checked_add_signed(change.supply_change)?; + fn apply_supply_change(mut self, supply_change: SupplyChange) -> Option { + self.total_supply = supply_change.apply_to(self.total_supply)?; 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"); + *self = self + .revert_finalization(change.is_finalized) + .apply_supply_change(-change.supply_change) + .expect("reverted change should be validated"); + } + + fn revert_finalization(mut self, is_finalized: bool) -> Self { + self.is_finalized &= !is_finalized; + self } } @@ -62,50 +131,78 @@ impl From> for IssuedAssets { } impl AssetStateChange { - fn from_note(is_finalized: bool, note: orchard::Note) -> (AssetBase, Self) { + fn new( + asset_base: AssetBase, + supply_change: SupplyChange, + is_finalized: bool, + ) -> (AssetBase, Self) { ( - note.asset(), + asset_base, Self { is_finalized, - supply_change: note.value().inner().into(), + supply_change, }, ) } - fn from_notes( - is_finalized: bool, - notes: &[orchard::Note], - ) -> impl Iterator + '_ { - notes - .iter() - .map(move |note| Self::from_note(is_finalized, *note)) + fn from_transaction<'a>( + tx: &'a Arc, + ) -> impl Iterator + 'a { + Self::from_burns(tx.orchard_burns()) + .chain(Self::from_issue_actions(tx.orchard_issue_actions())) } fn from_issue_actions<'a>( actions: impl Iterator + 'a, ) -> impl Iterator + 'a { - actions.flat_map(|action| Self::from_notes(action.is_finalized(), action.notes())) + actions.flat_map(Self::from_issue_action) } - fn from_burn(burn: &BurnItem) -> (AssetBase, Self) { - ( - burn.asset(), - Self { - is_finalized: false, - supply_change: -i128::from(burn.amount()), - }, + fn from_issue_action<'a>( + action: &'a IssueAction, + ) -> impl Iterator + 'a { + let supply_changes = Self::from_notes(action.notes()); + let finalize_changes = action + .is_finalized() + .then(|| { + action + .notes() + .iter() + .map(orchard::Note::asset) + .collect::>() + }) + .unwrap_or_default() + .into_iter() + .map(|asset_base| Self::new(asset_base, SupplyChange::Issuance(0), true)); + + supply_changes.chain(finalize_changes) + } + + fn from_notes(notes: &[orchard::Note]) -> impl Iterator + '_ { + notes.iter().copied().map(Self::from_note) + } + + fn from_note(note: orchard::Note) -> (AssetBase, Self) { + Self::new( + note.asset(), + SupplyChange::Issuance(note.value().inner()), + false, ) } fn from_burns(burns: &[BurnItem]) -> impl Iterator + '_ { burns.iter().map(Self::from_burn) } -} -impl std::ops::AddAssign for AssetStateChange { - fn add_assign(&mut self, rhs: Self) { - self.is_finalized |= rhs.is_finalized; - self.supply_change += rhs.supply_change; + fn from_burn(burn: &BurnItem) -> (AssetBase, Self) { + Self::new(burn.asset(), SupplyChange::Burn(burn.amount()), false) + } + + /// 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) -> bool { + self.is_finalized |= change.is_finalized; + self.supply_change.add(change.supply_change) } } @@ -151,45 +248,33 @@ impl IssuedAssetsChange { Self(HashMap::new()) } - fn update<'a>(&mut self, changes: impl Iterator + 'a) { + fn update<'a>( + &mut self, + changes: impl Iterator + 'a, + ) -> bool { for (asset_base, change) in changes { - *self.0.entry(asset_base).or_default() += change; - } - } - - /// Accepts a slice of [`Arc`]s. - /// - /// Returns a tuple, ([`IssuedAssetsChange`], [`IssuedAssetsChange`]), where - /// the first item is from burns and the second one is for issuance. - pub fn from_transactions(transactions: &[Arc]) -> (Self, Self) { - let mut burn_change = Self::new(); - let mut issuance_change = Self::new(); - - for transaction in transactions { - burn_change.update(AssetStateChange::from_burns(transaction.orchard_burns())); - issuance_change.update(AssetStateChange::from_issue_actions( - transaction.orchard_issue_actions(), - )); + if !self.0.entry(asset_base).or_default().apply_change(change) { + return false; + } } - (burn_change, issuance_change) + true } /// Accepts a slice of [`Arc`]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]) -> Self { + pub fn from_transactions(transactions: &[Arc]) -> Option { 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(), - )); + if !issued_assets_change.update(AssetStateChange::from_transaction(transaction)) { + return None; + } } - issued_assets_change + Some(issued_assets_change) } /// Consumes self and accepts a closure for looking up previous asset states.