Refactor Sapling data and use it in V4

zebra-chain/src/sapling.rs

@@ -13,11 +13,13 @@ mod tests;
 // XXX clean up these modules
 
 pub mod keys;
+pub mod shielded_data;
 pub mod tree;
 
 pub use address::Address;
 pub use commitment::{CommitmentRandomness, NoteCommitment, ValueCommitment};
 pub use keys::Diversifier;
 pub use note::{EncryptedNote, Note, Nullifier, WrappedNoteKey};
 pub use output::Output;
+pub use shielded_data::{AnchorVariant, PerSpendAnchor, SharedAnchor, ShieldedData};
 pub use spend::Spend;

zebra-chain/src/sapling/arbitrary.rs

@@ -3,9 +3,9 @@ use proptest::{arbitrary::any, array, collection::vec, prelude::*};
 
 use crate::primitives::Groth16Proof;
 
-use super::{keys, note, tree, NoteCommitment, Output, Spend, ValueCommitment};
+use super::{keys, note, tree, NoteCommitment, Output, PerSpendAnchor, Spend, ValueCommitment};
 
-impl Arbitrary for Spend {
+impl Arbitrary for Spend<PerSpendAnchor> {
     type Parameters = ();
 
     fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
@@ -16,18 +16,20 @@ impl Arbitrary for Spend {
             any::<Groth16Proof>(),
             vec(any::<u8>(), 64),
         )
-            .prop_map(|(anchor, nullifier, rpk_bytes, proof, sig_bytes)| Self {
-                anchor,
-                cv: ValueCommitment(AffinePoint::identity()),
-                nullifier,
-                rk: redjubjub::VerificationKeyBytes::from(rpk_bytes),
-                zkproof: proof,
-                spend_auth_sig: redjubjub::Signature::from({
-                    let mut b = [0u8; 64];
-                    b.copy_from_slice(sig_bytes.as_slice());
-                    b
-                }),
-            })
+            .prop_map(
+                |(per_spend_anchor, nullifier, rpk_bytes, proof, sig_bytes)| Self {
+                    per_spend_anchor,
+                    cv: ValueCommitment(AffinePoint::identity()),
+                    nullifier,
+                    rk: redjubjub::VerificationKeyBytes::from(rpk_bytes),
+                    zkproof: proof,
+                    spend_auth_sig: redjubjub::Signature::from({
+                        let mut b = [0u8; 64];
+                        b.copy_from_slice(sig_bytes.as_slice());
+                        b
+                    }),
+                },
+            )
             .boxed()
     }
 

zebra-chain/src/transaction/shielded_data.rs → zebra-chain/src/sapling/shielded_data.rs

@@ -1,12 +1,60 @@
+//! Sapling shielded data for `V4` and `V5` `Transaction`s.
+//!
+//! Zebra uses a generic shielded data type for `V4` and `V5` transactions.
+//! The `value_balance` change is handled using the default zero value.
+//! The anchor change is handled using the `AnchorVariant` type trait.
+
 use futures::future::Either;
 
 use crate::{
     amount::Amount,
     primitives::redjubjub::{Binding, Signature},
-    sapling::{Nullifier, Output, Spend, ValueCommitment},
+    sapling::{tree, Nullifier, Output, Spend, ValueCommitment},
     serialization::serde_helpers,
 };
 
+use serde::{de::DeserializeOwned, Serialize};
+use std::{
+    cmp::{Eq, PartialEq},
+    fmt::Debug,
+};
+
+/// Per-Spend Sapling anchors, used in Transaction V4 and the
+/// `spends_per_anchor` method.
+#[derive(Copy, Clone, Debug, Deserialize, Serialize, PartialEq, Eq)]
+pub struct PerSpendAnchor {}
+
+/// Shared Sapling anchors, used in Transaction V5.
+#[derive(Copy, Clone, Debug, Deserialize, Serialize, PartialEq, Eq)]
+pub struct SharedAnchor {}
+
+impl AnchorVariant for PerSpendAnchor {
+    type Shared = ();
+    type PerSpend = tree::Root;
+}
+
+impl AnchorVariant for SharedAnchor {
+    type Shared = tree::Root;
+    type PerSpend = ();
+}
+
+/// A type trait to handle structural differences between V4 and V5 Sapling
+/// Transaction anchors.
+///
+/// In Transaction V4, anchors are per-Spend. In Transaction V5, there is a
+/// single transaction anchor for all Spends in a transaction.
+pub trait AnchorVariant {
+    /// The type of the shared anchor.
+    ///
+    /// `()` means "not present in this transaction version".
+    type Shared: Clone + Debug + DeserializeOwned + Serialize + Eq + PartialEq;
+
+    /// The type of the per-spend anchor.
+    ///
+    /// `()` means "not present in this transaction version".
+    type PerSpend: Clone + Debug + DeserializeOwned + Serialize + Eq + PartialEq;
+}
+
 /// A bundle of [`Spend`] and [`Output`] descriptions and signature data.
 ///
 /// Spend and Output descriptions are optional, but Zcash transactions must
@@ -15,8 +63,29 @@ use crate::{
 /// description with the required signature data, so that an
 /// `Option<ShieldedData>` correctly models the presence or absence of any
 /// shielded data.
+///
+/// # Differences between Transaction Versions
+///
+/// The Sapling `value_balance` field is optional in `Transaction::V5`, but
+/// required in `Transaction::V4`. In both cases, if there is no `ShieldedData`,
+/// then the field value must be zero. Therefore, only need to store
+/// `value_balance` when there is some Sapling `ShieldedData`.
+///
+/// In `Transaction::V4`, each `Spend` has its own anchor. In `Transaction::V5`,
+/// there is a single `shared_anchor` for the entire transaction. This
+/// structural difference is modeled using the `AnchorVariant` type trait.
+/// A type of `()` means "not present in this transaction version".
 #[derive(Clone, Debug, Serialize, Deserialize)]
-pub struct ShieldedData {
+pub struct ShieldedData<AnchorV>
+where
+    AnchorV: AnchorVariant + Clone,
+{
+    /// The net value of Sapling spend transfers minus output transfers.
+    pub value_balance: Amount,
+    /// The shared anchor for all `Spend`s in this transaction.
+    ///
+    /// A type of `()` means "not present in this transaction version".
+    pub shared_anchor: AnchorV::Shared,
     /// Either a spend or output description.
     ///
     /// Storing this separately ensures that it is impossible to construct
@@ -27,12 +96,12 @@ pub struct ShieldedData {
     /// methods provide iterators over all of the [`Spend`]s and
     /// [`Output`]s.
     #[serde(with = "serde_helpers::Either")]
-    pub first: Either<Spend, Output>,
+    pub first: Either<Spend<AnchorV>, Output>,
     /// The rest of the [`Spend`]s for this transaction.
     ///
     /// Note that the [`ShieldedData::spends`] method provides an iterator
     /// over all spend descriptions.
-    pub rest_spends: Vec<Spend>,
+    pub rest_spends: Vec<Spend<AnchorV>>,
     /// The rest of the [`Output`]s for this transaction.
     ///
     /// Note that the [`ShieldedData::outputs`] method provides an iterator
@@ -42,9 +111,37 @@ pub struct ShieldedData {
     pub binding_sig: Signature<Binding>,
 }
 
-impl ShieldedData {
+impl<AnchorV> ShieldedData<AnchorV>
+where
+    AnchorV: AnchorVariant + Clone,
+    Spend<PerSpendAnchor>: From<(Spend<AnchorV>, AnchorV::Shared)>,
+{
     /// Iterate over the [`Spend`]s for this transaction.
-    pub fn spends(&self) -> impl Iterator<Item = &Spend> {
+    ///
+    /// Returns `Spend<PerSpendAnchor>` regardless of the underlying transaction
+    /// version, to allow generic verification over V4 and V5 transactions.
+    ///
+    /// # Correctness
+    ///
+    /// Do not use this function for serialization.
+    pub fn spends_per_anchor(&self) -> impl Iterator<Item = Spend<PerSpendAnchor>> + '_ {
+        self.spends()
+            .cloned()
+            .map(move |spend| Spend::<PerSpendAnchor>::from((spend, self.shared_anchor.clone())))
+    }
+}
+
+impl<AnchorV> ShieldedData<AnchorV>
+where
+    AnchorV: AnchorVariant + Clone,
+{
+    /// Iterate over the [`Spend`]s for this transaction, returning them as
+    /// their generic type.
+    ///
+    /// # Correctness
+    ///
+    /// Use this function for serialization.
+    pub fn spends(&self) -> impl Iterator<Item = &Spend<AnchorV>> {
         match self.first {
             Either::Left(ref spend) => Some(spend),
             Either::Right(_) => None,
@@ -96,13 +193,11 @@ impl ShieldedData {
     /// descriptions of the transaction, and the balancing value.
     ///
     /// https://zips.z.cash/protocol/protocol.pdf#saplingbalance
-    pub fn binding_verification_key(
-        &self,
-        value_balance: Amount,
-    ) -> redjubjub::VerificationKeyBytes<Binding> {
+    pub fn binding_verification_key(&self) -> redjubjub::VerificationKeyBytes<Binding> {
         let cv_old: ValueCommitment = self.spends().map(|spend| spend.cv).sum();
         let cv_new: ValueCommitment = self.outputs().map(|output| output.cv).sum();
-        let cv_balance: ValueCommitment = ValueCommitment::new(jubjub::Fr::zero(), value_balance);
+        let cv_balance: ValueCommitment =
+            ValueCommitment::new(jubjub::Fr::zero(), self.value_balance);
 
         let key_bytes: [u8; 32] = (cv_old - cv_new - cv_balance).into();
 
@@ -114,8 +209,14 @@ impl ShieldedData {
 // of a ShieldedData with at least one spend and at least one output, depending
 // on which goes in the `first` slot.  This is annoying but a smallish price to
 // pay for structural validity.
+//
+// A `ShieldedData<PerSpendAnchor>` can never be equal to a
+// `ShieldedData<SharedAnchor>`, even if they have the same effects.
 
-impl std::cmp::PartialEq for ShieldedData {
+impl<AnchorV> std::cmp::PartialEq for ShieldedData<AnchorV>
+where
+    AnchorV: AnchorVariant + Clone + PartialEq,
+{
     fn eq(&self, other: &Self) -> bool {
         // First check that the lengths match, so we know it is safe to use zip,
         // which truncates to the shorter of the two iterators.
@@ -126,11 +227,14 @@ impl std::cmp::PartialEq for ShieldedData {
             return false;
         }
 
-        // Now check that the binding_sig, spends, outputs match.
-        self.binding_sig == other.binding_sig
+        // Now check that all the fields match
+        self.value_balance == other.value_balance
+            && self.shared_anchor == other.shared_anchor
+            && self.binding_sig == other.binding_sig
             && self.spends().zip(other.spends()).all(|(a, b)| a == b)
             && self.outputs().zip(other.outputs()).all(|(a, b)| a == b)
     }
 }
 
-impl std::cmp::Eq for ShieldedData {}
+impl<AnchorV> std::cmp::Eq for ShieldedData<AnchorV> where AnchorV: AnchorVariant + Clone + PartialEq
+{}