Modifiy zebra-chain structures to support multiple action groups

zebra-chain/Cargo.toml

@@ -17,6 +17,7 @@ categories = ["asynchronous", "cryptography::cryptocurrencies", "encoding"]
 [features]
 #default = []
 default = ["tx-v6"]
+#default = ["tx-v6", "zsa-swap"]
 
 # Production features that activate extra functionality
 
@@ -67,6 +68,9 @@ tx-v6 = [
     "nonempty"
 ]
 
+# Support for ZSA asset swaps
+zsa-swap = []
+
 [dependencies]
 
 # Cryptography

zebra-chain/src/orchard.rs

@@ -23,7 +23,7 @@ pub use address::Address;
 pub use commitment::{CommitmentRandomness, NoteCommitment, ValueCommitment};
 pub use keys::Diversifier;
 pub use note::{EncryptedNote, Note, Nullifier, WrappedNoteKey};
-pub use shielded_data::{AuthorizedAction, Flags, ShieldedData};
+pub use shielded_data::{ActionGroup, AuthorizedAction, Flags, ShieldedData};
 pub use shielded_data_flavor::{OrchardVanilla, ShieldedDataFlavor};
 
 pub(crate) use shielded_data::ActionCommon;

zebra-chain/src/orchard/shielded_data.rs

@@ -22,19 +22,18 @@ use crate::{
 
 use super::{OrchardVanilla, ShieldedDataFlavor};
 
-/// A bundle of [`Action`] descriptions and signature data.
+// FIXME: wrap all ActionGroup usages with tx-v6 feature flag?
+/// Action Group description.
+#[cfg(feature = "tx-v6")]
 #[derive(Clone, Debug, PartialEq, Eq, Deserialize, Serialize)]
 #[serde(bound(
     serialize = "FL::EncryptedNote: serde::Serialize, FL::BurnType: serde::Serialize",
     deserialize = "FL::BurnType: serde::Deserialize<'de>"
 ))]
-pub struct ShieldedData<FL: ShieldedDataFlavor> {
+pub struct ActionGroup<FL: ShieldedDataFlavor> {
     /// The orchard flags for this transaction.
     /// Denoted as `flagsOrchard` in the spec.
     pub flags: Flags,
-    /// The net value of Orchard spends minus outputs.
-    /// Denoted as `valueBalanceOrchard` in the spec.
-    pub value_balance: Amount,
     /// The shared anchor for all `Spend`s in this transaction.
     /// Denoted as `anchorOrchard` in the spec.
     pub shared_anchor: tree::Root,
@@ -44,28 +43,67 @@ pub struct ShieldedData<FL: ShieldedDataFlavor> {
     /// The Orchard Actions, in the order they appear in the transaction.
     /// Denoted as `vActionsOrchard` and `vSpendAuthSigsOrchard` in the spec.
     pub actions: AtLeastOne<AuthorizedAction<FL>>,
-    /// A signature on the transaction `sighash`.
-    /// Denoted as `bindingSigOrchard` in the spec.
-    pub binding_sig: Signature<Binding>,
 
     #[cfg(feature = "tx-v6")]
     /// Assets intended for burning
     /// Denoted as `vAssetBurn` in the spec (ZIP 230).
     pub burn: FL::BurnType,
 }
 
-impl<FL: ShieldedDataFlavor> fmt::Display for ShieldedData<FL> {
+impl<FL: ShieldedDataFlavor> ActionGroup<FL> {
+    /// Iterate over the [`Action`]s for the [`AuthorizedAction`]s in this
+    /// action group, in the order they appear in it.
+    pub fn actions(&self) -> impl Iterator<Item = &Action<FL>> {
+        self.actions.actions()
+    }
+}
+
+/// A bundle of [`Action`] descriptions and signature data.
+#[derive(Clone, Debug, PartialEq, Eq, Deserialize, Serialize)]
+#[serde(bound(
+    serialize = "FL::EncryptedNote: serde::Serialize, FL::BurnType: serde::Serialize",
+    deserialize = "FL::BurnType: serde::Deserialize<'de>"
+))]
+pub struct ShieldedData<FL: ShieldedDataFlavor> {
+    /// Action Group descriptions.
+    /// Denoted as `vActionGroupsOrchard` in the spec (ZIP 230).
+    pub action_groups: AtLeastOne<ActionGroup<FL>>,
+    /// Denoted as `valueBalanceOrchard` in the spec.
+    pub value_balance: Amount,
+    /// A signature on the transaction `sighash`.
+    /// Denoted as `bindingSigOrchard` in the spec.
+    pub binding_sig: Signature<Binding>,
+}
+
+impl<FL: ShieldedDataFlavor> fmt::Display for ActionGroup<FL> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        let mut fmter = f.debug_struct("orchard::ShieldedData");
+        let mut fmter = f.debug_struct("orchard::ActionGroup");
+
+        // FIXME: reorder fields here according the struct/spec?
 
         fmter.field("actions", &self.actions.len());
-        fmter.field("value_balance", &self.value_balance);
         fmter.field("flags", &self.flags);
 
         fmter.field("proof_len", &self.proof.zcash_serialized_size());
 
         fmter.field("shared_anchor", &self.shared_anchor);
 
+        #[cfg(feature = "tx-v6")]
+        fmter.field("burn", &self.burn.as_ref().len());
+
+        fmter.finish()
+    }
+}
+
+impl<FL: ShieldedDataFlavor> fmt::Display for ShieldedData<FL> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let mut fmter = f.debug_struct("orchard::ShieldedData");
+
+        fmter.field("action_groups", &self.action_groups);
+        fmter.field("value_balance", &self.value_balance);
+
+        // FIXME: format binding_sig as well?
+
         fmter.finish()
     }
 }
@@ -74,13 +112,14 @@ impl<FL: ShieldedDataFlavor> ShieldedData<FL> {
     /// Iterate over the [`Action`]s for the [`AuthorizedAction`]s in this
     /// transaction, in the order they appear in it.
     pub fn actions(&self) -> impl Iterator<Item = &Action<FL>> {
-        self.actions.actions()
+        self.authorized_actions()
+            .map(|authorized_action| &authorized_action.action)
     }
 
     /// Return an iterator for the [`ActionCommon`] copy of the Actions in this
     /// transaction, in the order they appear in it.
     pub fn action_commons(&self) -> impl Iterator<Item = ActionCommon> + '_ {
-        self.actions.actions().map(|action| action.into())
+        self.actions().map(|action| action.into())
     }
 
     /// Collect the [`Nullifier`]s for this transaction.
@@ -123,7 +162,14 @@ impl<FL: ShieldedDataFlavor> ShieldedData<FL> {
 
         // FIXME: use asset to create ValueCommitment here for burns and above for value_balance?
         #[cfg(feature = "tx-v6")]
-        let key_bytes: [u8; 32] = (cv - cv_balance - self.burn.clone().into()).into();
+        let key_bytes: [u8; 32] = (cv
+            - cv_balance
+            - self
+                .action_groups
+                .iter()
+                .map(|action_group| action_group.burn.clone().into())
+                .sum::<ValueCommitment>())
+        .into();
 
         key_bytes.into()
     }
@@ -140,6 +186,29 @@ impl<FL: ShieldedDataFlavor> ShieldedData<FL> {
     pub fn note_commitments(&self) -> impl Iterator<Item = &pallas::Base> {
         self.actions().map(|action| &action.cm_x)
     }
+
+    /// Makes a union of the flags for this transaction.
+    pub fn flags_union(&self) -> Flags {
+        self.action_groups
+            .iter()
+            .map(|action_group| &action_group.flags)
+            .fold(Flags::empty(), |result, flags| result.union(*flags))
+    }
+
+    /// Collect the shared anchors for this transaction.
+    pub fn shared_anchors(&self) -> impl Iterator<Item = tree::Root> + '_ {
+        self.action_groups
+            .iter()
+            .map(|action_group| action_group.shared_anchor)
+    }
+
+    /// Iterate over the [`AuthorizedAction`]s in this
+    /// transaction, in the order they appear in it.
+    pub fn authorized_actions(&self) -> impl Iterator<Item = &AuthorizedAction<FL>> {
+        self.action_groups
+            .iter()
+            .flat_map(|action_group| action_group.actions.iter())
+    }
 }
 
 impl<FL: ShieldedDataFlavor> AtLeastOne<AuthorizedAction<FL>> {

zebra-chain/src/orchard/shielded_data_flavor.rs

@@ -57,6 +57,8 @@ pub trait ShieldedDataFlavor: OrchardFlavor {
     type BurnType: Clone
         + Debug
         + Default
+        + PartialEq
+        + Eq
         + ZcashDeserialize
         + ZcashSerialize
         + Into<ValueCommitment>

zebra-chain/src/transaction.rs

@@ -80,8 +80,8 @@ macro_rules! orchard_shielded_data_iter {
 
 // FIXME: doc this
 // Move down
-macro_rules! orchard_shielded_data_field {
-    ($self:expr, $field:ident) => {
+macro_rules! orchard_shielded_data_method {
+    ($self:expr, $method:ident) => {
         match $self {
             // No Orchard shielded data
             Transaction::V1 { .. }
@@ -92,13 +92,13 @@ macro_rules! orchard_shielded_data_field {
             Transaction::V5 {
                 orchard_shielded_data,
                 ..
-            } => orchard_shielded_data.as_ref().map(|data| data.$field),
+            } => orchard_shielded_data.as_ref().map(|data| data.$method()),
 
             #[cfg(feature = "tx-v6")]
             Transaction::V6 {
                 orchard_shielded_data,
                 ..
-            } => orchard_shielded_data.as_ref().map(|data| data.$field),
+            } => orchard_shielded_data.as_ref().map(|data| data.$method()),
         }
     };
 }
@@ -354,11 +354,12 @@ impl Transaction {
     ///
     /// See [`Self::has_transparent_or_shielded_inputs`] for details.
     pub fn has_shielded_inputs(&self) -> bool {
+        // FIXME: is it correct to use orchard_flags_union here?
         self.joinsplit_count() > 0
             || self.sapling_spends_per_anchor().count() > 0
             || (self.orchard_actions().count() > 0
                 && self
-                    .orchard_flags()
+                    .orchard_flags_union()
                     .unwrap_or_else(orchard::Flags::empty)
                     .contains(orchard::Flags::ENABLE_SPENDS))
     }
@@ -372,11 +373,12 @@ impl Transaction {
     ///
     /// See [`Self::has_transparent_or_shielded_outputs`] for details.
     pub fn has_shielded_outputs(&self) -> bool {
+        // FIXME: is it correct to use orchard_flags_union here?
         self.joinsplit_count() > 0
             || self.sapling_outputs().count() > 0
             || (self.orchard_actions().count() > 0
                 && self
-                    .orchard_flags()
+                    .orchard_flags_union()
                     .unwrap_or_else(orchard::Flags::empty)
                     .contains(orchard::Flags::ENABLE_OUTPUTS))
     }
@@ -386,7 +388,7 @@ impl Transaction {
         if self.version() < 5 || self.orchard_actions().count() == 0 {
             return true;
         }
-        self.orchard_flags()
+        self.orchard_flags_union()
             .unwrap_or_else(orchard::Flags::empty)
             .intersects(orchard::Flags::ENABLE_SPENDS | orchard::Flags::ENABLE_OUTPUTS)
     }
@@ -1070,20 +1072,20 @@ impl Transaction {
 
     /// Access the [`orchard::Flags`] in this transaction, if there is any,
     /// regardless of version.
-    pub fn orchard_flags(&self) -> Option<orchard::shielded_data::Flags> {
-        orchard_shielded_data_field!(self, flags)
+    pub fn orchard_flags_union(&self) -> Option<orchard::shielded_data::Flags> {
+        orchard_shielded_data_method!(self, flags_union)
     }
 
     /// Access the [`orchard::tree::Root`] in this transaction, if there is any,
     /// regardless of version.
-    pub fn orchard_shared_anchor(&self) -> Option<orchard::tree::Root> {
-        orchard_shielded_data_field!(self, shared_anchor)
+    pub fn orchard_shared_anchors(&self) -> Box<dyn Iterator<Item = orchard::tree::Root> + '_> {
+        orchard_shielded_data_iter!(self, orchard::ShieldedData::shared_anchors)
     }
 
     /// Return if the transaction has any Orchard shielded data,
     /// regardless of version.
     pub fn has_orchard_shielded_data(&self) -> bool {
-        self.orchard_flags().is_some()
+        orchard_shielded_data_method!(self, value_balance).is_some()
     }
 
     // value balances
@@ -1458,7 +1460,7 @@ impl Transaction {
     /// <https://zebra.zfnd.org/dev/rfcs/0012-value-pools.html#definitions>
     pub fn orchard_value_balance(&self) -> ValueBalance<NegativeAllowed> {
         let orchard_value_balance =
-            orchard_shielded_data_field!(self, value_balance).unwrap_or_else(Amount::zero);
+            orchard_shielded_data_method!(self, value_balance).unwrap_or_else(Amount::zero);
 
         ValueBalance::from_orchard_amount(orchard_value_balance)
     }

zebra-chain/src/transaction/arbitrary.rs

@@ -14,7 +14,7 @@ use crate::{
     parameters::{Network, NetworkUpgrade},
     primitives::{Bctv14Proof, Groth16Proof, Halo2Proof, ZkSnarkProof},
     sapling::{self, AnchorVariant, PerSpendAnchor, SharedAnchor},
-    serialization::ZcashDeserializeInto,
+    serialization::{AtLeastOne, ZcashDeserializeInto},
     sprout, transparent,
     value_balance::{ValueBalance, ValueBalanceError},
     LedgerState,
@@ -811,17 +811,20 @@ impl<FL: orchard::ShieldedDataFlavor + 'static> Arbitrary for orchard::ShieldedD
                 let (flags, value_balance, shared_anchor, proof, actions, binding_sig, burn) =
                     props;
 
+                // FIXME: support multiple action groups
                 Self {
-                    flags,
+                    action_groups: AtLeastOne::from_one(orchard::ActionGroup {
+                        flags,
+                        shared_anchor,
+                        proof,
+                        actions: actions
+                            .try_into()
+                            .expect("arbitrary vector size range produces at least one action"),
+                        #[cfg(feature = "tx-v6")]
+                        burn,
+                    }),
                     value_balance,
-                    shared_anchor,
-                    proof,
-                    actions: actions
-                        .try_into()
-                        .expect("arbitrary vector size range produces at least one action"),
                     binding_sig: binding_sig.0,
-                    #[cfg(feature = "tx-v6")]
-                    burn,
                 }
             })
             .boxed()
@@ -1157,14 +1160,16 @@ pub fn insert_fake_orchard_shielded_data(
 
     // Place the dummy action inside the Orchard shielded data
     let dummy_shielded_data = orchard::ShieldedData::<orchard::OrchardVanilla> {
-        flags: orchard::Flags::empty(),
+        action_groups: AtLeastOne::from_one(orchard::ActionGroup {
+            flags: orchard::Flags::empty(),
+            shared_anchor: orchard::tree::Root::default(),
+            proof: Halo2Proof(vec![]),
+            actions: at_least_one![dummy_authorized_action],
+            #[cfg(feature = "tx-v6")]
+            burn: Default::default(),
+        }),
         value_balance: Amount::try_from(0).expect("invalid transaction amount"),
-        shared_anchor: orchard::tree::Root::default(),
-        proof: Halo2Proof(vec![]),
-        actions: at_least_one![dummy_authorized_action],
         binding_sig: Signature::from([0u8; 64]),
-        #[cfg(feature = "tx-v6")]
-        burn: Default::default(),
     };
 
     // Replace the shielded data in the transaction