From 878a428c3ca7304e365a97c94ee3c15b37b34216 Mon Sep 17 00:00:00 2001
From: Kris Nuttycombe <kris@nutty.land>
Date: Mon, 25 Jul 2022 19:29:56 -0600
Subject: [PATCH] Factor out witness updates from scan_output

---
 zcash_client_backend/src/welding_rig.rs | 48 +++++++++++--------------
 1 file changed, 21 insertions(+), 27 deletions(-)

diff --git a/zcash_client_backend/src/welding_rig.rs b/zcash_client_backend/src/welding_rig.rs
index 2b651d680..e6bce1c86 100644
--- a/zcash_client_backend/src/welding_rig.rs
+++ b/zcash_client_backend/src/welding_rig.rs
@@ -18,7 +18,7 @@ use zcash_primitives::{
     zip32::{AccountId, ExtendedFullViewingKey},
 };
 
-use crate::proto::compact_formats::{CompactBlock, CompactSaplingOutput};
+use crate::proto::compact_formats::CompactBlock;
 use crate::wallet::{WalletShieldedOutput, WalletShieldedSpend, WalletTx};
 
 /// Scans a [`CompactSaplingOutput`] with a set of [`ScanningKey`]s.
@@ -35,29 +35,11 @@ fn scan_output<P: consensus::Parameters, K: ScanningKey>(
     params: &P,
     height: BlockHeight,
     index: usize,
-    output: CompactSaplingOutput,
+    output: CompactOutputDescription,
     vks: &[(&AccountId, &K)],
     spent_from_accounts: &HashSet<AccountId>,
     tree: &mut CommitmentTree<Node>,
-    existing_witnesses: &mut [&mut IncrementalWitness<Node>],
-    block_witnesses: &mut [&mut IncrementalWitness<Node>],
-    new_witnesses: &mut [&mut IncrementalWitness<Node>],
 ) -> Option<WalletShieldedOutput<K::Nf>> {
-    let output = CompactOutputDescription::try_from(output).ok()?;
-
-    // Increment tree and witnesses
-    let node = Node::new(output.cmu.to_repr());
-    for witness in existing_witnesses {
-        witness.append(node).unwrap();
-    }
-    for witness in block_witnesses {
-        witness.append(node).unwrap();
-    }
-    for witness in new_witnesses {
-        witness.append(node).unwrap();
-    }
-    tree.append(node).unwrap();
-
     for (account, vk) in vks.iter() {
         let (note, to) = match vk.try_decryption(params, height, &output) {
             Some(ret) => ret,
@@ -283,26 +265,38 @@ pub fn scan_block<P: consensus::Parameters, K: ScanningKey>(
                 })
                 .collect();
 
-            for (idx, c_out) in tx.outputs.into_iter().enumerate() {
+            for (index, output) in tx.outputs.into_iter().enumerate() {
                 // Grab mutable references to new witnesses from previous outputs
                 // in this transaction so that we can update them. Scoped so we
                 // don't hold mutable references to shielded_outputs for too long.
-                let mut new_witnesses: Vec<_> = shielded_outputs
+                let new_witnesses: Vec<_> = shielded_outputs
                     .iter_mut()
                     .map(|output| &mut output.witness)
                     .collect();
 
+                let output = CompactOutputDescription::try_from(output).ok().unwrap();
+
+                // Increment tree and witnesses
+                let node = Node::new(output.cmu.to_repr());
+                for witness in &mut *existing_witnesses {
+                    witness.append(node).unwrap();
+                }
+                for witness in &mut block_witnesses {
+                    witness.append(node).unwrap();
+                }
+                for witness in new_witnesses {
+                    witness.append(node).unwrap();
+                }
+                tree.append(node).unwrap();
+
                 if let Some(output) = scan_output(
                     params,
                     block_height,
-                    idx,
-                    c_out,
+                    index,
+                    output,
                     vks,
                     &spent_from_accounts,
                     tree,
-                    existing_witnesses,
-                    &mut block_witnesses,
-                    &mut new_witnesses,
                 ) {
                     shielded_outputs.push(output);
                 }