From af9edcea2512f94a4a2afe4c911f5ad2b3c5bcb3 Mon Sep 17 00:00:00 2001
From: Kris Nuttycombe <kris@nutty.land>
Date: Thu, 3 Feb 2022 13:45:41 -0700
Subject: [PATCH] Define transaction::TransactionData::map_authorization

---
 .../src/extensions/transparent.rs             | 10 +++++
 .../src/transaction/components/orchard.rs     |  5 +++
 .../src/transaction/components/sapling.rs     | 39 +++++++++++++++++++
 .../src/transaction/components/transparent.rs | 23 +++++++++++
 .../src/transaction/components/tze.rs         | 22 +++++++++++
 zcash_primitives/src/transaction/mod.rs       | 29 ++++++++++++++
 6 files changed, 128 insertions(+)

diff --git a/zcash_primitives/src/extensions/transparent.rs b/zcash_primitives/src/extensions/transparent.rs
index 7ca6a3d6c..a105b73f9 100644
--- a/zcash_primitives/src/extensions/transparent.rs
+++ b/zcash_primitives/src/extensions/transparent.rs
@@ -76,6 +76,16 @@ pub struct Witness<T> {
     pub payload: T,
 }
 
+impl<T> Witness<T> {
+    pub fn map_payload<U, F: FnOnce(T) -> U>(self, f: F) -> Witness<U> {
+        Witness {
+            extension_id: self.extension_id,
+            mode: self.mode,
+            payload: f(self.payload),
+        }
+    }
+}
+
 impl Witness<AuthData> {
     /// Produce the intermediate format for an extension-specific witness
     /// type.
diff --git a/zcash_primitives/src/transaction/components/orchard.rs b/zcash_primitives/src/transaction/components/orchard.rs
index a65c7c65d..4d30b64dd 100644
--- a/zcash_primitives/src/transaction/components/orchard.rs
+++ b/zcash_primitives/src/transaction/components/orchard.rs
@@ -28,6 +28,11 @@ impl Authorization for Unauthorized {
     type SpendAuth = ();
 }
 
+pub trait MapAuth<A: Authorization, B: Authorization> {
+    fn map_spend_auth(&self, s: A::SpendAuth) -> B::SpendAuth;
+    fn map_authorization(&self, a: A) -> B;
+}
+
 /// Reads an [`orchard::Bundle`] from a v5 transaction format.
 pub fn read_v5_bundle<R: Read>(
     mut reader: R,
diff --git a/zcash_primitives/src/transaction/components/sapling.rs b/zcash_primitives/src/transaction/components/sapling.rs
index 2352f6723..8413b897b 100644
--- a/zcash_primitives/src/transaction/components/sapling.rs
+++ b/zcash_primitives/src/transaction/components/sapling.rs
@@ -47,6 +47,12 @@ impl Authorization for Authorized {
     type AuthSig = redjubjub::Signature;
 }
 
+pub trait MapAuth<A: Authorization, B: Authorization> {
+    fn map_proof(&self, p: A::Proof) -> B::Proof;
+    fn map_auth_sig(&self, s: A::AuthSig) -> B::AuthSig;
+    fn map_authorization(&self, a: A) -> B;
+}
+
 #[derive(Debug, Clone)]
 pub struct Bundle<A: Authorization> {
     pub shielded_spends: Vec<SpendDescription<A>>,
@@ -55,6 +61,39 @@ pub struct Bundle<A: Authorization> {
     pub authorization: A,
 }
 
+impl<A: Authorization> Bundle<A> {
+    pub fn map_authorization<B: Authorization, F: MapAuth<A, B>>(self, f: F) -> Bundle<B> {
+        Bundle {
+            shielded_spends: self
+                .shielded_spends
+                .into_iter()
+                .map(|d| SpendDescription {
+                    cv: d.cv,
+                    anchor: d.anchor,
+                    nullifier: d.nullifier,
+                    rk: d.rk,
+                    zkproof: f.map_proof(d.zkproof),
+                    spend_auth_sig: f.map_auth_sig(d.spend_auth_sig),
+                })
+                .collect(),
+            shielded_outputs: self
+                .shielded_outputs
+                .into_iter()
+                .map(|o| OutputDescription {
+                    cv: o.cv,
+                    cmu: o.cmu,
+                    ephemeral_key: o.ephemeral_key,
+                    enc_ciphertext: o.enc_ciphertext,
+                    out_ciphertext: o.out_ciphertext,
+                    zkproof: f.map_proof(o.zkproof),
+                })
+                .collect(),
+            value_balance: self.value_balance,
+            authorization: f.map_authorization(self.authorization),
+        }
+    }
+}
+
 #[derive(Clone)]
 pub struct SpendDescription<A: Authorization> {
     pub cv: jubjub::ExtendedPoint,
diff --git a/zcash_primitives/src/transaction/components/transparent.rs b/zcash_primitives/src/transaction/components/transparent.rs
index 59059dbb9..e52e0cdb4 100644
--- a/zcash_primitives/src/transaction/components/transparent.rs
+++ b/zcash_primitives/src/transaction/components/transparent.rs
@@ -22,6 +22,11 @@ impl Authorization for Authorized {
     type ScriptSig = Script;
 }
 
+pub trait MapAuth<A: Authorization, B: Authorization> {
+    fn map_script_sig(&self, s: A::ScriptSig) -> B::ScriptSig;
+    fn map_authorization(&self, s: A) -> B;
+}
+
 #[derive(Debug, Clone, PartialEq)]
 pub struct Bundle<A: Authorization> {
     pub vin: Vec<TxIn<A>>,
@@ -29,6 +34,24 @@ pub struct Bundle<A: Authorization> {
     pub authorization: A,
 }
 
+impl<A: Authorization> Bundle<A> {
+    pub fn map_authorization<B: Authorization, F: MapAuth<A, B>>(self, f: F) -> Bundle<B> {
+        Bundle {
+            vin: self
+                .vin
+                .into_iter()
+                .map(|txin| TxIn {
+                    prevout: txin.prevout,
+                    script_sig: f.map_script_sig(txin.script_sig),
+                    sequence: txin.sequence,
+                })
+                .collect(),
+            vout: self.vout,
+            authorization: f.map_authorization(self.authorization),
+        }
+    }
+}
+
 #[derive(Clone, Debug, PartialEq)]
 pub struct OutPoint {
     hash: [u8; 32],
diff --git a/zcash_primitives/src/transaction/components/tze.rs b/zcash_primitives/src/transaction/components/tze.rs
index c72dadc2b..326950610 100644
--- a/zcash_primitives/src/transaction/components/tze.rs
+++ b/zcash_primitives/src/transaction/components/tze.rs
@@ -28,6 +28,11 @@ impl Authorization for Authorized {
     type Witness = tze::AuthData;
 }
 
+pub trait MapAuth<A: Authorization, B: Authorization> {
+    fn map_witness(&self, s: A::Witness) -> B::Witness;
+    fn map_authorization(&self, s: A) -> B;
+}
+
 #[derive(Debug, Clone, PartialEq)]
 pub struct Bundle<A: Authorization> {
     pub vin: Vec<TzeIn<A::Witness>>,
@@ -35,6 +40,23 @@ pub struct Bundle<A: Authorization> {
     pub authorization: A,
 }
 
+impl<A: Authorization> Bundle<A> {
+    pub fn map_authorization<B: Authorization, F: MapAuth<A, B>>(self, f: F) -> Bundle<B> {
+        Bundle {
+            vin: self
+                .vin
+                .into_iter()
+                .map(|tzein| TzeIn {
+                    prevout: tzein.prevout,
+                    witness: tzein.witness.map_payload(|p| f.map_witness(p)),
+                })
+                .collect(),
+            vout: self.vout,
+            authorization: f.map_authorization(self.authorization),
+        }
+    }
+}
+
 #[derive(Clone, Debug, PartialEq)]
 pub struct OutPoint {
     txid: TxId,
diff --git a/zcash_primitives/src/transaction/mod.rs b/zcash_primitives/src/transaction/mod.rs
index 25755c38d..f4f3bce04 100644
--- a/zcash_primitives/src/transaction/mod.rs
+++ b/zcash_primitives/src/transaction/mod.rs
@@ -370,6 +370,35 @@ impl<A: Authorization> TransactionData<A> {
             digester.digest_tze(self.tze_bundle.as_ref()),
         )
     }
+
+    pub fn map_authorization<B: Authorization>(
+        self,
+        f_transparent: impl transparent::MapAuth<A::TransparentAuth, B::TransparentAuth>,
+        f_sapling: impl sapling::MapAuth<A::SaplingAuth, B::SaplingAuth>,
+        mut f_orchard: impl orchard_serialization::MapAuth<A::OrchardAuth, B::OrchardAuth>,
+        #[cfg(feature = "zfuture")] f_tze: impl tze::MapAuth<A::TzeAuth, B::TzeAuth>,
+    ) -> TransactionData<B> {
+        TransactionData {
+            version: self.version,
+            consensus_branch_id: self.consensus_branch_id,
+            lock_time: self.lock_time,
+            expiry_height: self.expiry_height,
+            transparent_bundle: self
+                .transparent_bundle
+                .map(|b| b.map_authorization(f_transparent)),
+            sprout_bundle: self.sprout_bundle,
+            sapling_bundle: self.sapling_bundle.map(|b| b.map_authorization(f_sapling)),
+            orchard_bundle: self.orchard_bundle.map(|b| {
+                b.authorize(
+                    &mut f_orchard,
+                    |f, _, s| f.map_spend_auth(s),
+                    |f, a| f.map_authorization(a),
+                )
+            }),
+            #[cfg(feature = "zfuture")]
+            tze_bundle: self.tze_bundle.map(|b| b.map_authorization(f_tze)),
+        }
+    }
 }
 
 impl<A: Authorization> std::fmt::Debug for TransactionData<A> {