From b44653e686e9b7ae16f2228f034535c5dbbdc67a Mon Sep 17 00:00:00 2001
From: Jack Grigg <jack@z.cash>
Date: Mon, 26 Aug 2019 11:59:07 +0100
Subject: [PATCH] Move cmu and epk parsing onto CompactOutput struct

---
 zcash_client_backend/src/proto/mod.rs   | 31 ++++++++++++++++++++++++-
 zcash_client_backend/src/welding_rig.rs | 25 ++++----------------
 2 files changed, 34 insertions(+), 22 deletions(-)

diff --git a/zcash_client_backend/src/proto/mod.rs b/zcash_client_backend/src/proto/mod.rs
index 50d562b..0ab1b6d 100644
--- a/zcash_client_backend/src/proto/mod.rs
+++ b/zcash_client_backend/src/proto/mod.rs
@@ -1,6 +1,12 @@
 //! Generated code for handling light client protobuf structs.
 
-use zcash_primitives::block::{BlockHash, BlockHeader};
+use ff::{PrimeField, PrimeFieldRepr};
+use pairing::bls12_381::{Bls12, Fr, FrRepr};
+use zcash_primitives::{
+    block::{BlockHash, BlockHeader},
+    jubjub::{edwards, PrimeOrder},
+    JUBJUB,
+};
 
 pub mod compact_formats;
 
@@ -52,3 +58,26 @@ impl compact_formats::CompactBlock {
         }
     }
 }
+
+impl compact_formats::CompactOutput {
+    /// Returns the note commitment for this output.
+    ///
+    /// A convenience method that parses [`CompactOutput.cmu`].
+    ///
+    /// [`CompactOutput.cmu`]: #structfield.cmu
+    pub fn cmu(&self) -> Result<Fr, ()> {
+        let mut repr = FrRepr::default();
+        repr.read_le(&self.cmu[..]).map_err(|_| ())?;
+        Fr::from_repr(repr).map_err(|_| ())
+    }
+
+    /// Returns the ephemeral public key for this output.
+    ///
+    /// A convenience method that parses [`CompactOutput.epk`].
+    ///
+    /// [`CompactOutput.epk`]: #structfield.epk
+    pub fn epk(&self) -> Result<edwards::Point<Bls12, PrimeOrder>, ()> {
+        let p = edwards::Point::<Bls12, _>::read(&self.epk[..], &JUBJUB).map_err(|_| ())?;
+        p.as_prime_order(&JUBJUB).ok_or(())
+    }
+}
diff --git a/zcash_client_backend/src/welding_rig.rs b/zcash_client_backend/src/welding_rig.rs
index 307dcea..0996607 100644
--- a/zcash_client_backend/src/welding_rig.rs
+++ b/zcash_client_backend/src/welding_rig.rs
@@ -1,17 +1,15 @@
 //! Tools for scanning a compact representation of the Zcash block chain.
 
-use ff::{PrimeField, PrimeFieldRepr};
-use pairing::bls12_381::{Bls12, Fr, FrRepr};
+use ff::PrimeField;
 use std::collections::HashSet;
 use subtle::{ConditionallySelectable, ConstantTimeEq, CtOption};
 use zcash_primitives::{
-    jubjub::{edwards, fs::Fs},
+    jubjub::fs::Fs,
     merkle_tree::{CommitmentTree, IncrementalWitness},
     note_encryption::try_sapling_compact_note_decryption,
     sapling::Node,
     transaction::TxId,
     zip32::ExtendedFullViewingKey,
-    JUBJUB,
 };
 
 use crate::proto::compact_formats::{CompactBlock, CompactOutput};
@@ -33,23 +31,8 @@ fn scan_output(
     block_witnesses: &mut [&mut IncrementalWitness<Node>],
     new_witnesses: &mut [IncrementalWitness<Node>],
 ) -> Option<(WalletShieldedOutput, IncrementalWitness<Node>)> {
-    let mut repr = FrRepr::default();
-    if repr.read_le(&output.cmu[..]).is_err() {
-        return None;
-    }
-    let cmu = match Fr::from_repr(repr) {
-        Ok(cmu) => cmu,
-        Err(_) => return None,
-    };
-
-    let epk = match edwards::Point::<Bls12, _>::read(&output.epk[..], &JUBJUB) {
-        Ok(p) => match p.as_prime_order(&JUBJUB) {
-            Some(epk) => epk,
-            None => return None,
-        },
-        Err(_) => return None,
-    };
-
+    let cmu = output.cmu().ok()?;
+    let epk = output.epk().ok()?;
     let ct = output.ciphertext;
 
     // Increment tree and witnesses