Refactor Orchard ShieldedData to support multiple action groups

zebra-chain/Cargo.toml

@@ -66,6 +66,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};
 
 #[cfg(feature = "tx_v6")]

zebra-chain/src/orchard/shielded_data.rs

@@ -28,7 +28,8 @@ use orchard::{note::AssetBase, value::ValueSum};
 
 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.
 #[derive(Clone, Debug, PartialEq, Eq, Deserialize, Serialize)]
 #[cfg_attr(
     not(feature = "tx_v6"),
@@ -41,40 +42,77 @@ use super::{OrchardVanilla, ShieldedDataFlavor};
         deserialize = "Flavor::BurnType: serde::Deserialize<'de>"
     ))
 )]
-pub struct ShieldedData<Flavor: ShieldedDataFlavor> {
+pub struct ActionGroup<Flavor: 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,
+
+    /// Block height after which this Bundle's Actions are invalid by consensus.
+    /// Denoted as `nAGExpiryHeight` in the spec.
+    #[cfg(feature = "tx_v6")]
+    pub expiry_height: u32,
+
+    /// Assets intended for burning
+    /// Denoted as `vAssetBurn` in the spec (ZIP 230).
+    #[cfg(feature = "tx_v6")]
+    pub burn: Flavor::BurnType,
+
     /// The aggregated zk-SNARK proof for all the actions in this transaction.
     /// Denoted as `proofsOrchard` in the spec.
     pub proof: Halo2Proof,
     /// The Orchard Actions, in the order they appear in the transaction.
     /// Denoted as `vActionsOrchard` and `vSpendAuthSigsOrchard` in the spec.
     pub actions: AtLeastOne<AuthorizedAction<Flavor>>,
+}
+
+impl<Flavor: ShieldedDataFlavor> ActionGroup<Flavor> {
+    /// 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<Flavor>> {
+        self.actions.actions()
+    }
+}
+
+/// A bundle of [`Action`] descriptions and signature data.
+#[derive(Clone, Debug, PartialEq, Eq, Deserialize, Serialize)]
+#[cfg_attr(
+    not(feature = "tx_v6"),
+    serde(bound(serialize = "Flavor::EncryptedNote: serde::Serialize"))
+)]
+#[cfg_attr(
+    feature = "tx_v6",
+    serde(bound(
+        serialize = "Flavor::EncryptedNote: serde::Serialize, Flavor::BurnType: serde::Serialize",
+        deserialize = "Flavor::BurnType: serde::Deserialize<'de>"
+    ))
+)]
+pub struct ShieldedData<Flavor: ShieldedDataFlavor> {
+    /// Action Group descriptions.
+    /// Denoted as `vActionGroupsOrchard` in the spec (ZIP 230).
+    pub action_groups: AtLeastOne<ActionGroup<Flavor>>,
+    /// 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>,
-
-    #[cfg(feature = "tx_v6")]
-    /// Assets intended for burning
-    /// Denoted as `vAssetBurn` in the spec (ZIP 230).
-    pub burn: Flavor::BurnType,
 }
 
-impl<Flavor: ShieldedDataFlavor> fmt::Display for ShieldedData<Flavor> {
+impl<Flavor: ShieldedDataFlavor> fmt::Display for ActionGroup<Flavor> {
     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");
 
         fmter.field("actions", &self.actions.len());
-        fmter.field("value_balance", &self.value_balance);
         fmter.field("flags", &self.flags);
 
+        #[cfg(feature = "tx_v6")]
+        fmter.field("expiry_height", &self.expiry_height);
+
+        #[cfg(feature = "tx_v6")]
+        fmter.field("burn", &self.burn.as_ref().len());
+
         fmter.field("proof_len", &self.proof.zcash_serialized_size());
 
         fmter.field("shared_anchor", &self.shared_anchor);
@@ -83,11 +121,25 @@ impl<Flavor: ShieldedDataFlavor> fmt::Display for ShieldedData<Flavor> {
     }
 }
 
+impl<Flavor: ShieldedDataFlavor> fmt::Display for ShieldedData<Flavor> {
+    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()
+    }
+}
+
 impl<Flavor: ShieldedDataFlavor> ShieldedData<Flavor> {
     /// 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<Flavor>> {
-        self.actions.actions()
+        self.authorized_actions()
+            .map(|authorized_action| &authorized_action.action)
     }
 
     /// Collect the [`Nullifier`]s for this transaction.
@@ -139,7 +191,11 @@ impl<Flavor: ShieldedDataFlavor> ShieldedData<Flavor> {
                 (ValueSum::default() + i64::from(self.value_balance)).unwrap(),
                 AssetBase::native(),
             );
-            let burn_value_commitment = compute_burn_value_commitment(self.burn.as_ref());
+            let burn_value_commitment = self
+                .action_groups
+                .iter()
+                .map(|action_group| compute_burn_value_commitment(action_group.burn.as_ref()))
+                .sum::<ValueCommitment>();
             cv - cv_balance - burn_value_commitment
         };
 
@@ -160,6 +216,29 @@ impl<Flavor: ShieldedDataFlavor> ShieldedData<Flavor> {
     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<Flavor>> {
+        self.action_groups
+            .iter()
+            .flat_map(|action_group| action_group.actions.iter())
+    }
 }
 
 impl<Flavor: ShieldedDataFlavor> AtLeastOne<AuthorizedAction<Flavor>> {

zebra-chain/src/orchard/shielded_data_flavor.rs

@@ -53,6 +53,8 @@ pub trait ShieldedDataFlavor: OrchardFlavor {
     #[cfg(feature = "tx_v6")]
     type BurnType: Clone
         + Debug
+        + PartialEq
+        + Eq
         + ZcashDeserialize
         + ZcashSerialize
         + AsRef<[BurnItem]>

zebra-chain/src/primitives/zcash_note_encryption.rs

@@ -70,8 +70,8 @@ pub fn decrypts_successfully(transaction: &Transaction, network: &Network, heigh
 }
 
 /// Checks if all actions in an Orchard bundle decrypt successfully.
-fn orchard_bundle_decrypts_successfully<A: Authorization, V, D: OrchardPrimitives>(
-    bundle: &Bundle<A, V, D>,
+fn orchard_bundle_decrypts_successfully<A: Authorization, V, P: OrchardPrimitives>(
+    bundle: &Bundle<A, V, P>,
 ) -> bool {
     bundle.actions().iter().all(|act| {
         zcash_note_encryption::try_output_recovery_with_ovk(