make merkle depth an adjustable parameter in the code

src/bin/tx.rs

@@ -67,7 +67,6 @@ impl MemoryState {
                     .append(node)
                     .expect("append to witness");
             }
-            assert_eq!(self.own_coins.len(), 0);
 
             if let Some((note, secret)) = self.try_decrypt_note(enc_note) {
                 // We need to keep track of the witness for this coin.
@@ -196,7 +195,7 @@ fn main() {
     assert_eq!(state.own_coins.len(), 1);
     //let (coin, note, secret, witness) = &mut state.own_coins[0];
 
-    let auth_path = {
+    let merkle_path = {
         let tree = &mut state.tree;
         //let coin: &Coin = &state.own_coins[0].0;
         //let witness = &mut state.own_coins[0].3;
@@ -218,16 +217,9 @@ fn main() {
 
         assert_eq!(state.merkle_roots.len(), 16);
 
-        // TODO: Some stupid glue code. Need to put this somewhere else.
+        // Just test the path is good
         let merkle_path = witness.path().unwrap();
-        let auth_path: Vec<(bls12_381::Scalar, bool)> = merkle_path
-            .auth_path
-            .iter()
-            .map(|(node, b)| ((*node).into(), *b))
-            .collect();
-
         let node = Node::from_coin(&coin);
-
         let root = tree.root();
         drop(tree);
         drop(witness);
@@ -235,7 +227,7 @@ fn main() {
         let root = root.into();
         assert!(state.is_valid_merkle(&root));
 
-        auth_path
+        merkle_path
     };
 
     // Step 3: wallet1 sends payment to wallet2
@@ -248,18 +240,12 @@ fn main() {
 
     // Make a spend tx
 
-    // Get the coin we're spending from the previous tx
-    let coin = {
-        let tx = tx::Transaction::decode(&tx_data[..]).unwrap();
-        tx.outputs[0].revealed.coin
-    };
-
     // Construct a new tx spending the coin
     // We need the decrypted note and our private key
     let builder = tx::TransactionBuilder {
         clear_inputs: vec![],
         inputs: vec![tx::TransactionBuilderInputInfo {
-            merkle_path: auth_path,
+            merkle_path,
             secret: secret.clone(),
             note: state.own_coins[0].1.clone(),
         }],
@@ -281,5 +267,6 @@ fn main() {
         let tx = tx::Transaction::decode(&tx_data[..]).unwrap();
         assert!(state.is_valid_merkle(&tx.inputs[0].revealed.merkle_root));
         let update = state_transition(&state, tx).expect("step 3 state transition failed");
+        state.apply(update);
     }
 }

src/circuit/spend_contract.rs

@@ -13,20 +13,16 @@ use ff::{Field, PrimeField};
 use group::Curve;
 use zcash_proofs::circuit::{ecc, pedersen_hash};
 
+use crate::crypto::node::SAPLING_COMMITMENT_TREE_DEPTH;
+
 pub struct SpendContract {
     pub value: Option<u64>,
     pub randomness_value: Option<jubjub::Fr>,
     pub serial: Option<jubjub::Fr>,
     pub randomness_coin: Option<jubjub::Fr>,
     pub secret: Option<jubjub::Fr>,
-    pub branch_0: Option<bls12_381::Scalar>,
-    pub is_right_0: Option<bool>,
-    pub branch_1: Option<bls12_381::Scalar>,
-    pub is_right_1: Option<bool>,
-    pub branch_2: Option<bls12_381::Scalar>,
-    pub is_right_2: Option<bool>,
-    pub branch_3: Option<bls12_381::Scalar>,
-    pub is_right_3: Option<bool>,
+    pub branch: [Option<bls12_381::Scalar>; SAPLING_COMMITMENT_TREE_DEPTH],
+    pub is_right: [Option<bool>; SAPLING_COMMITMENT_TREE_DEPTH],
     pub signature_secret: Option<jubjub::Fr>,
 }
 impl Circuit<bls12_381::Scalar> for SpendContract {
@@ -253,6 +249,54 @@ impl Circuit<bls12_381::Scalar> for SpendContract {
         // Line 168: ec_get_u current cm
         let mut current = cm.get_u().clone();
 
+        for i in 0..SAPLING_COMMITMENT_TREE_DEPTH {
+            // Line 174: alloc_scalar branch param:branch_0
+            let branch = num::AllocatedNum::alloc(
+                cs.namespace(|| "Line 174: alloc_scalar branch param:branch_0"),
+                || Ok(*self.branch[i].get()?),
+            )?;
+
+            // Line 177: alloc_bit is_right param:is_right_0
+            let is_right = boolean::Boolean::from(boolean::AllocatedBit::alloc(
+                cs.namespace(|| "Line 177: alloc_bit is_right param:is_right_0"),
+                self.is_right[i],
+            )?);
+
+            // Line 180: conditionally_reverse left right current branch is_right
+            let (left, right) = num::AllocatedNum::conditionally_reverse(
+                cs.namespace(|| "Line 180: conditionally_reverse left right current branch is_right"),
+                &current,
+                &branch,
+                &is_right,
+            )?;
+
+            // Line 183: scalar_as_binary left left
+            let left = left.to_bits_le(cs.namespace(|| "Line 183: scalar_as_binary left left"))?;
+
+            // Line 184: scalar_as_binary right right
+            let right = right.to_bits_le(cs.namespace(|| "Line 184: scalar_as_binary right right"))?;
+
+            // Line 185: alloc_binary preimage
+            let mut preimage = vec![];
+
+            // Line 186: binary_extend preimage left
+            preimage.extend(left);
+
+            // Line 187: binary_extend preimage right
+            preimage.extend(right);
+
+            // Line 188: pedersen_hash cm preimage MERKLE_0
+            let mut cm = pedersen_hash::pedersen_hash(
+                cs.namespace(|| "Line 188: pedersen_hash cm preimage MERKLE_0"),
+                pedersen_hash::Personalization::MerkleTree(i),
+                &preimage,
+            )?;
+
+            // Line 190: ec_get_u current cm
+            current = cm.get_u().clone();
+        }
+
+        /*
         // Line 174: alloc_scalar branch param:branch_0
         let branch = num::AllocatedNum::alloc(
             cs.namespace(|| "Line 174: alloc_scalar branch param:branch_0"),
@@ -432,6 +476,7 @@ impl Circuit<bls12_381::Scalar> for SpendContract {
 
         // Line 250: ec_get_u current cm
         let mut current = cm.get_u().clone();
+        */
 
         // Line 253: emit_scalar current
         current.inputize(cs.namespace(|| "Line 253: emit_scalar current"))?;

src/crypto/node.rs

@@ -6,7 +6,7 @@ use std::io;
 
 use super::{coin::Coin, merkle::Hashable};
 
-pub const SAPLING_COMMITMENT_TREE_DEPTH: usize = 4;
+pub const SAPLING_COMMITMENT_TREE_DEPTH: usize = 6;
 
 /// Compute a parent node in the Sapling commitment tree given its two children.
 pub fn merkle_hash(depth: usize, lhs: &[u8; 32], rhs: &[u8; 32]) -> bls12_381::Scalar {

src/crypto/spend_proof.rs

@@ -8,7 +8,7 @@ use rand::rngs::OsRng;
 use std::io;
 use std::time::Instant;
 
-use super::node::merkle_hash;
+use super::node::{SAPLING_COMMITMENT_TREE_DEPTH, merkle_hash};
 use crate::circuit::spend_contract::SpendContract;
 use crate::error::Result;
 use crate::serial::{Decodable, Encodable};
@@ -183,14 +183,8 @@ pub fn setup_spend_prover() -> groth16::Parameters<Bls12> {
             randomness_coin: None,
             secret: None,
 
-            branch_0: None,
-            is_right_0: None,
-            branch_1: None,
-            is_right_1: None,
-            branch_2: None,
-            is_right_2: None,
-            branch_3: None,
-            is_right_3: None,
+            branch: [None; SAPLING_COMMITMENT_TREE_DEPTH],
+            is_right: [None; SAPLING_COMMITMENT_TREE_DEPTH],
 
             signature_secret: None,
         };
@@ -210,11 +204,16 @@ pub fn create_spend_proof(
     merkle_path: Vec<(bls12_381::Scalar, bool)>,
     signature_secret: jubjub::Fr,
 ) -> (groth16::Proof<Bls12>, SpendRevealedValues) {
-    assert_eq!(merkle_path.len(), 4);
     assert_eq!(
         merkle_path.len(),
-        super::node::SAPLING_COMMITMENT_TREE_DEPTH
+        SAPLING_COMMITMENT_TREE_DEPTH
     );
+    let mut branch: [_; SAPLING_COMMITMENT_TREE_DEPTH] = Default::default();
+    let mut is_right: [_; SAPLING_COMMITMENT_TREE_DEPTH] = Default::default();
+    for (i, (branch_i, is_right_i)) in merkle_path.iter().enumerate() {
+        branch[i] = Some(branch_i.clone());
+        is_right[i] = Some(is_right_i.clone());
+    }
     let c = SpendContract {
         value: Some(value),
         randomness_value: Some(randomness_value),
@@ -222,14 +221,9 @@ pub fn create_spend_proof(
         randomness_coin: Some(randomness_coin),
         secret: Some(secret),
 
-        branch_0: Some(merkle_path[0].0),
-        is_right_0: Some(merkle_path[0].1),
-        branch_1: Some(merkle_path[1].0),
-        is_right_1: Some(merkle_path[1].1),
-        branch_2: Some(merkle_path[2].0),
-        is_right_2: Some(merkle_path[2].1),
-        branch_3: Some(merkle_path[3].0),
-        is_right_3: Some(merkle_path[3].1),
+        branch,
+        is_right,
+
         signature_secret: Some(signature_secret),
     };
 

src/tx/builder.rs

@@ -7,7 +7,7 @@ use super::{
     partial::{PartialTransaction, PartialTransactionClearInput, PartialTransactionInput},
     Transaction, TransactionClearInput, TransactionInput, TransactionOutput,
 };
-use crate::crypto::{create_mint_proof, create_spend_proof, note::Note, schnorr};
+use crate::crypto::{merkle::MerklePath, node::Node, create_mint_proof, create_spend_proof, note::Note, schnorr};
 use crate::serial::Encodable;
 
 pub struct TransactionBuilder {
@@ -22,7 +22,7 @@ pub struct TransactionBuilderClearInputInfo {
 }
 
 pub struct TransactionBuilderInputInfo {
-    pub merkle_path: Vec<(bls12_381::Scalar, bool)>,
+    pub merkle_path: MerklePath<Node>,
     pub secret: jubjub::Fr,
     pub note: Note,
 }
@@ -84,6 +84,13 @@ impl TransactionBuilder {
 
             // make proof
 
+            // TODO: Some stupid glue code. Need to sort this out
+            let auth_path: Vec<(bls12_381::Scalar, bool)> = input.merkle_path
+                .auth_path
+                .iter()
+                .map(|(node, b)| ((*node).into(), *b))
+                .collect();
+
             let (proof, revealed) = create_spend_proof(
                 &spend_params,
                 input.note.value,
@@ -91,7 +98,7 @@ impl TransactionBuilder {
                 input.note.serial,
                 input.note.coin_blind,
                 input.secret,
-                input.merkle_path.clone(),
+                auth_path,
                 signature_secret.clone(),
             );