fix test with valid merkle proofs

nssa/Cargo.toml

@@ -13,6 +13,7 @@ sha2 = "0.10.9"
 secp256k1 = "0.31.1"
 rand = "0.8"
 borsh = "1.5.7"
+bytemuck = "1.13"
 
 [dev-dependencies]
 test-program-methods = { path = "test_program_methods" }

nssa/core/src/lib.rs

@@ -1,4 +1,7 @@
-use risc0_zkvm::serde::to_vec;
+use risc0_zkvm::{
+    serde::to_vec,
+    sha::{Impl, Sha256},
+};
 use serde::{Deserialize, Serialize};
 
 #[cfg(feature = "host")]
@@ -21,15 +24,35 @@ pub mod program;
 #[cfg(feature = "host")]
 pub mod error;
 
-pub type CommitmentSetDigest = [u32; 8];
-pub type MembershipProof = Vec<[u8; 32]>;
+pub type CommitmentSetDigest = [u8; 32];
+pub type MembershipProof = (usize, Vec<[u8; 32]>);
 pub fn verify_membership_proof(
     commitment: &Commitment,
     proof: &MembershipProof,
     digest: &CommitmentSetDigest,
 ) -> bool {
-    // TODO: implement
-    true
+    let value_bytes = commitment.to_byte_array();
+    let mut result: [u8; 32] = Impl::hash_bytes(&value_bytes)
+        .as_bytes()
+        .try_into()
+        .unwrap();
+    let mut level_index = proof.0;
+    for node in &proof.1 {
+        let is_left_child = level_index & 1 == 0;
+        if is_left_child {
+            let mut bytes = [0u8; 64];
+            bytes[..32].copy_from_slice(&result);
+            bytes[32..].copy_from_slice(node);
+            result = Impl::hash_bytes(&bytes).as_bytes().try_into().unwrap();
+        } else {
+            let mut bytes = [0u8; 64];
+            bytes[..32].copy_from_slice(node);
+            bytes[32..].copy_from_slice(&result);
+            result = Impl::hash_bytes(&bytes).as_bytes().try_into().unwrap();
+        }
+        level_index >>= 1;
+    }
+    &result == digest
 }
 
 pub type IncomingViewingPublicKey = [u8; 32];
@@ -158,7 +181,7 @@ mod tests {
                 &Commitment::new(&NullifierPublicKey::from(&[2; 32]), &Account::default()),
                 &[1; 32],
             )],
-            commitment_set_digest: [0, 1, 0, 1, 0, 1, 0, 1],
+            commitment_set_digest: [0xab; 32],
         };
         let bytes = output.to_bytes();
         let output_from_slice: PrivacyPreservingCircuitOutput = from_slice(&bytes).unwrap();

nssa/src/merkle_tree.rs

@@ -93,7 +93,7 @@ impl MerkleTree {
         let base_length = self.capacity;
         let mut layer_node = hash_value(&value);
         let mut layer_index = new_index + base_length - 1;
-        self.node_map.insert(layer_index, layer_node);
+        self.set_node(layer_index, layer_node);
 
         let mut layer = 0;
         let mut top_layer = self.depth();
@@ -159,6 +159,10 @@ impl MerkleTree {
         }
         Some((*value_index, result))
     }
+
+    pub(crate) fn contains(&self, value: &[u8; 32]) -> bool {
+        self.index_map.contains_key(value)
+    }
 }
 
 // Reference implementation
@@ -518,7 +522,6 @@ mod tests {
         assert!(tree.get_authentication_path_for(&value).is_none());
     }
 
-
     #[test]
     fn test_authentication_path_5() {
         let values = vec![[1; 32], [2; 32], [3; 32], [4; 32], [5; 32]];

nssa/src/privacy_preserving_transaction/mod.rs

@@ -113,14 +113,16 @@ pub mod circuit {
 mod tests {
     use nssa_core::{
         EncryptedAccountData,
-        account::{Account, AccountWithMetadata, NullifierPublicKey, NullifierSecretKey},
+        account::{
+            Account, AccountWithMetadata, Commitment, NullifierPublicKey, NullifierSecretKey,
+        },
     };
     use risc0_zkvm::{InnerReceipt, Journal, Receipt};
 
     use crate::{
-        Address, V01State,
+        Address, V01State, merkle_tree::MerkleTree,
         privacy_preserving_transaction::circuit::prove_privacy_preserving_execution_circuit,
-        program::Program,
+        program::Program, state::CommitmentSet,
     };
 
     use super::*;
@@ -153,7 +155,7 @@ mod tests {
         let private_account_keys = vec![(NullifierPublicKey::from(&[1; 32]), [2; 32], [3; 32])];
         let private_account_auth = vec![];
         let visibility_mask = vec![0, 2];
-        let commitment_set_digest = [99; 8];
+        let commitment_set_digest = [99; 32];
         let program = Program::authenticated_transfer_program();
         let (proof, output) = prove_privacy_preserving_execution_circuit(
             &pre_states,
@@ -190,6 +192,10 @@ mod tests {
             is_authorized: true,
         };
 
+        let private_key = [1; 32];
+        let Npk = NullifierPublicKey::from(&private_key);
+        let commitment = Commitment::new(&Npk, &sender.account);
+
         let recipient = AccountWithMetadata {
             account: Account::default(),
             is_authorized: false,
@@ -197,19 +203,22 @@ mod tests {
 
         let balance_to_move: u128 = 37;
 
+        let commitment_set = CommitmentSet(MerkleTree::new(vec![commitment.to_byte_array()]));
+
         let expected_sender_pre = sender.clone();
-        let pre_states = vec![sender, recipient];
+        let pre_states = vec![sender.clone(), recipient];
         let instruction_data = Program::serialize_instruction(balance_to_move).unwrap();
-        let private_key = [1; 32];
         let private_account_keys = vec![
-            (NullifierPublicKey::from(&private_key), [2; 32], [3; 32]),
+            (Npk.clone(), [2; 32], [3; 32]),
             (NullifierPublicKey::from(&[2; 32]), [4; 32], [5; 32]),
         ];
-        //  TODO: Replace dummy authentication path when implemented
-        let private_account_auth = vec![(private_key, vec![])];
+        let private_account_auth = vec![(
+            private_key,
+            commitment_set.get_proof_for(&commitment).unwrap(),
+        )];
         let visibility_mask = vec![1, 2];
-        let commitment_set_digest = [99; 8];
         let program = Program::authenticated_transfer_program();
+        let mut commitment_set_digest = commitment_set.digest();
         let (proof, output) = prove_privacy_preserving_execution_circuit(
             &pre_states,
             &instruction_data,

nssa/src/state.rs

@@ -1,29 +1,32 @@
 use crate::{
-    address::Address, error::NssaError,
+    address::Address, error::NssaError, merkle_tree::MerkleTree,
     privacy_preserving_transaction::PrivacyPreservingTransaction, program::Program,
     public_transaction::PublicTransaction,
 };
 use nssa_core::{
-    CommitmentSetDigest,
-    account::{Account, Commitment, Nullifier},
-    program::{DEFAULT_PROGRAM_ID, ProgramId},
+    account::{Account, Commitment, Nullifier}, program::{ProgramId, DEFAULT_PROGRAM_ID}, CommitmentSetDigest, MembershipProof
 };
 use std::collections::{HashMap, HashSet};
 
-struct CommitmentSet(HashSet<Commitment>);
+pub(crate) struct CommitmentSet(pub(crate) MerkleTree);
 
 impl CommitmentSet {
-    fn extend(&mut self, commitments: Vec<Commitment>) {
-        self.0.extend(commitments)
+    fn extend(&mut self, commitments: &[Commitment]) {
+        for commitment in commitments {
+            self.0.insert(commitment.to_byte_array());
+        }
     }
 
-    fn digest(&self) -> CommitmentSetDigest {
-        // TODO: implement
-        [0; 8]
+    pub fn digest(&self) -> CommitmentSetDigest {
+        self.0.root()
+    }
+
+    pub fn get_proof_for(&self, commitment: &Commitment) -> Option<MembershipProof> {
+        self.0.get_authentication_path_for(&commitment.to_byte_array())
     }
 
     fn contains(&self, commitment: &Commitment) -> bool {
-        self.0.contains(commitment)
+        self.0.contains(&commitment.to_byte_array())
     }
 }
 
@@ -54,7 +57,10 @@ impl V01State {
 
         let mut this = Self {
             public_state,
-            private_state: (CommitmentSet(HashSet::new()), NullifierSet::new()),
+            private_state: (
+                CommitmentSet(MerkleTree::with_capacity(32)),
+                NullifierSet::new(),
+            ),
             builtin_programs: HashMap::new(),
         };
 
@@ -101,7 +107,7 @@ impl V01State {
         let message = tx.message();
 
         // 2. Add new commitments
-        self.private_state.0.extend(message.new_commitments.clone());
+        self.private_state.0.extend(&message.new_commitments);
 
         // 3. Add new nullifiers
         self.private_state.1.extend(message.new_nullifiers.clone());