add test. refactor

nssa/core/src/lib.rs

@@ -28,7 +28,8 @@ pub fn verify_membership_proof(
     proof: &MembershipProof,
     digest: &CommitmentSetDigest,
 ) -> bool {
-    todo!()
+    // TODO: implement
+    true
 }
 
 pub type IncomingViewingPublicKey = [u8; 32];
@@ -93,6 +94,7 @@ pub struct PrivacyPreservingCircuitInput {
 }
 
 #[derive(Serialize, Deserialize)]
+#[cfg_attr(any(feature = "host", test), derive(Debug, PartialEq, Eq))]
 pub struct PrivacyPreservingCircuitOutput {
     pub public_pre_states: Vec<AccountWithMetadata>,
     pub public_post_states: Vec<Account>,
@@ -113,3 +115,58 @@ impl PrivacyPreservingCircuitOutput {
         result
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use risc0_zkvm::serde::from_slice;
+
+    use crate::{
+        EncryptedAccountData, PrivacyPreservingCircuitOutput,
+        account::{Account, AccountWithMetadata, Commitment, Nullifier, NullifierPublicKey},
+    };
+
+    #[test]
+    fn test_privacy_preserving_circuit_output_to_bytes_is_compatible_with_from_slice() {
+        let output = PrivacyPreservingCircuitOutput {
+            public_pre_states: vec![
+                AccountWithMetadata {
+                    account: Account {
+                        program_owner: [1, 2, 3, 4, 5, 6, 7, 8],
+                        balance: 12345678901234567890,
+                        data: b"test data".to_vec(),
+                        nonce: 18446744073709551614,
+                    },
+                    is_authorized: true,
+                },
+                AccountWithMetadata {
+                    account: Account {
+                        program_owner: [9, 9, 9, 8, 8, 8, 7, 7],
+                        balance: 123123123456456567112,
+                        data: b"test data".to_vec(),
+                        nonce: 9999999999999999999999,
+                    },
+                    is_authorized: false,
+                },
+            ],
+            public_post_states: vec![Account {
+                program_owner: [1, 2, 3, 4, 5, 6, 7, 8],
+                balance: 100,
+                data: b"post state data".to_vec(),
+                nonce: 18446744073709551615,
+            }],
+            encrypted_private_post_states: vec![EncryptedAccountData(0)],
+            new_commitments: vec![Commitment::new(
+                &NullifierPublicKey::from(&[1; 32]),
+                &Account::default(),
+            )],
+            new_nullifiers: vec![Nullifier::new(
+                &Commitment::new(&NullifierPublicKey::from(&[2; 32]), &Account::default()),
+                &[1; 32],
+            )],
+            commitment_set_digest: [0, 1, 0, 1, 0, 1, 0, 1],
+        };
+        let bytes = output.to_bytes();
+        let output_from_slice: PrivacyPreservingCircuitOutput = from_slice(&bytes).unwrap();
+        assert_eq!(output, output_from_slice);
+    }
+}

nssa/program_methods/guest/src/bin/privacy_preserving_circuit.rs

@@ -38,19 +38,6 @@ fn main() {
         panic!();
     }
 
-    let n_private_accounts = visibility_mask.iter().filter(|&&flag| flag != 0).count();
-    if private_account_nonces.len() != n_private_accounts {
-        panic!();
-    }
-    if private_account_keys.len() != n_private_accounts {
-        panic!();
-    }
-
-    let n_auth_private_accounts = visibility_mask.iter().filter(|&&flag| flag == 1).count();
-    if private_account_auth.len() != n_auth_private_accounts {
-        panic!();
-    }
-
     // These lists will be the public outputs of this circuit
     // and will be populated next.
     let mut public_pre_states: Vec<AccountWithMetadata> = Vec::new();
@@ -104,9 +91,12 @@ fn main() {
                     let nullifier = Nullifier::new(&commitment_pre, nsk);
                     new_nullifiers.push(nullifier);
                 } else {
-                    // Private account marked as empty
-                    if pre_states[i].account != Account::default() || pre_states[i].is_authorized {
-                        panic!("Invalid empty private account pre-state");
+                    if pre_states[i].account != Account::default() {
+                        panic!("Found new private account with non default values.");
+                    }
+
+                    if pre_states[i].is_authorized {
+                        panic!("Found new private account marked as authorized.");
                     }
                 }
 
@@ -127,6 +117,18 @@ fn main() {
         }
     }
 
+    if private_nonces_iter.next().is_some() {
+        panic!("Too many nonces.");
+    }
+
+    if private_keys_iter.next().is_some() {
+        panic!("Too many private accounts keys.");
+    }
+
+    if private_auth_iter.next().is_some() {
+        panic!("Too many private account authentication keys.");
+    }
+
     let output = PrivacyPreservingCircuitOutput {
         public_pre_states,
         public_post_states,

nssa/src/error.rs

@@ -42,4 +42,7 @@ pub enum NssaError {
 
     #[error("Circuit output deserialization error: {0}")]
     CircuitOutputDeserializationError(String),
+
+    #[error("Invalid privacy preserving execution circuit proof")]
+    InvalidPrivacyPreservingProof,
 }

nssa/src/privacy_preserving_transaction/mod.rs

@@ -31,8 +31,6 @@ pub mod circuit {
         }
     }
 
-    /// Executes and proves the program `P`.
-    /// Returns the proof
     fn execute_and_prove_program(
         program: &Program,
         pre_states: &[AccountWithMetadata],
@@ -128,14 +126,15 @@ mod tests {
     use super::*;
 
     #[test]
-    fn test() {
+    fn prove_privacy_preserving_execution_circuit_public_and_private_pre_accounts() {
         let sender = AccountWithMetadata {
             account: Account {
                 balance: 100,
                 ..Account::default()
             },
-            is_authorized: false,
+            is_authorized: true,
         };
+
         let recipient = AccountWithMetadata {
             account: Account::default(),
             is_authorized: false,
@@ -155,7 +154,7 @@ mod tests {
         let private_account_auth = vec![];
         let visibility_mask = vec![0, 2];
         let commitment_set_digest = [99; 8];
-        let program = Program::simple_balance_transfer();
+        let program = Program::authenticated_transfer_program();
         let (proof, output) = prove_privacy_preserving_execution_circuit(
             &pre_states,
             &instruction_data,
@@ -177,7 +176,61 @@ mod tests {
         assert_eq!(output.new_nullifiers.len(), 0);
         assert_eq!(output.commitment_set_digest, commitment_set_digest);
         assert_eq!(output.encrypted_private_post_states.len(), 1);
-        // TODO: replace with real assert when encryption is implemented
+        // TODO: replace with real assertion when encryption is implemented
         assert_eq!(output.encrypted_private_post_states[0].to_bytes(), vec![0]);
     }
+
+    #[test]
+    fn prove_privacy_preserving_execution_circuit_fully_private() {
+        let sender = AccountWithMetadata {
+            account: Account {
+                balance: 100,
+                ..Account::default()
+            },
+            is_authorized: true,
+        };
+
+        let recipient = AccountWithMetadata {
+            account: Account::default(),
+            is_authorized: false,
+        };
+
+        let balance_to_move: u128 = 37;
+
+        let expected_sender_pre = sender.clone();
+        let pre_states = vec![sender, 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]),
+            (NullifierPublicKey::from(&[2; 32]), [4; 32], [5; 32]),
+        ];
+        //  TODO: Replace dummy authentication path when implemented
+        let private_account_auth = vec![(private_key, vec![])];
+        let visibility_mask = vec![1, 2];
+        let commitment_set_digest = [99; 8];
+        let program = Program::authenticated_transfer_program();
+        let (proof, output) = prove_privacy_preserving_execution_circuit(
+            &pre_states,
+            &instruction_data,
+            &private_account_keys,
+            &private_account_auth,
+            &visibility_mask,
+            commitment_set_digest,
+            &program,
+        )
+        .unwrap();
+
+        assert!(proof.is_valid_for(&output));
+
+        assert_eq!(output.public_post_states.len(), 0);
+        assert_eq!(output.public_pre_states.len(), 0);
+        assert_eq!(output.new_commitments.len(), 2);
+        assert_eq!(output.new_nullifiers.len(), 1);
+        assert_eq!(output.commitment_set_digest, commitment_set_digest);
+        assert_eq!(output.encrypted_private_post_states.len(), 2);
+        // TODO: replace with real assertion when encryption is implemented
+        assert_eq!(output.encrypted_private_post_states[0].to_bytes(), vec![0]);
+        assert_eq!(output.encrypted_private_post_states[1].to_bytes(), vec![0]);
+    }
 }

nssa/src/privacy_preserving_transaction/transaction.rs

@@ -1,7 +1,7 @@
 use std::collections::{HashMap, HashSet};
 
 use nssa_core::account::{Account, AccountWithMetadata};
-use nssa_core::{CommitmentSetDigest, EncryptedAccountData};
+use nssa_core::{CommitmentSetDigest, EncryptedAccountData, PrivacyPreservingCircuitOutput};
 
 use crate::error::NssaError;
 use crate::privacy_preserving_transaction::circuit::Proof;
@@ -146,7 +146,18 @@ fn check_privacy_preserving_circuit_proof_is_valid(
     new_nullifiers: &[nssa_core::account::Nullifier],
     commitment_set_digest: CommitmentSetDigest,
 ) -> Result<(), NssaError> {
-    todo!()
+    let output = PrivacyPreservingCircuitOutput {
+        public_pre_states: public_pre_states.to_vec(),
+        public_post_states: public_post_states.to_vec(),
+        encrypted_private_post_states: encrypted_private_post_states.to_vec(),
+        new_commitments: new_commitments.to_vec(),
+        new_nullifiers: new_nullifiers.to_vec(),
+        commitment_set_digest,
+    };
+    proof
+        .is_valid_for(&output)
+        .then_some(())
+        .ok_or(NssaError::InvalidPrivacyPreservingProof)
 }
 
 use std::hash::Hash;