From e081031580d8042bc504cf3a6955eaec079c13d1 Mon Sep 17 00:00:00 2001
From: Sergio Chouhy <sergio.chouhy@gmail.com>
Date: Thu, 17 Jul 2025 09:20:03 -0300
Subject: [PATCH] add verification function

---
 .../examples/private_execution.rs             | 21 +++++-----
 .../outer_methods/guest/src/bin/outer.rs      |  1 +
 risc0-selective-privacy-poc/src/lib.rs        | 42 +++++++++++++++----
 3 files changed, 44 insertions(+), 20 deletions(-)

diff --git a/risc0-selective-privacy-poc/examples/private_execution.rs b/risc0-selective-privacy-poc/examples/private_execution.rs
index 2686596..d77e700 100644
--- a/risc0-selective-privacy-poc/examples/private_execution.rs
+++ b/risc0-selective-privacy-poc/examples/private_execution.rs
@@ -1,14 +1,12 @@
 use nssa::program::TransferProgram;
-use outer_methods::{OUTER_ELF, OUTER_ID};
-use risc0_zkvm::{default_prover, ExecutorEnv, ProveInfo, Receipt};
+use outer_methods::OUTER_ID;
 use sparse_merkle_tree::SparseMerkleTree;
 use toy_example_core::{
     account::Account,
     bytes_to_words,
     input::InputVisibiility,
-    types::{Address, AuthenticationPath, Commitment, Nonce, Nullifier},
+    types::{Address, AuthenticationPath, Commitment, Nullifier},
 };
-use transfer_methods::{TRANSFER_ELF, TRANSFER_ID};
 
 fn mint_fresh_account(address: Address) -> Account {
     let nonce = [0; 8];
@@ -48,7 +46,7 @@ fn main() {
         InputVisibiility::Private(Some((sender_private_key, auth_path))),
         InputVisibiility::Private(None),
     ];
-    let prove_info = nssa::prove_privacy_execution::<TransferProgram>(
+    let receipt = nssa::prove_privacy_execution::<TransferProgram>(
         &[sender, receiver],
         &balance_to_move,
         &visibilities,
@@ -56,13 +54,14 @@ fn main() {
     )
     .unwrap();
 
-    let receipt = prove_info.receipt;
-
-    // Sanity check
-    receipt.verify(OUTER_ID).unwrap();
-
-    let output: (Vec<Account>, Vec<Nullifier>, Vec<Commitment>) = receipt.journal.decode().unwrap();
+    let output: (Vec<Account>, Vec<Nullifier>, Vec<Commitment>, [u32; 8]) =
+        receipt.journal.decode().unwrap();
     println!("public_outputs: {:?}", output.0);
     println!("nullifiers: {:?}", output.1);
     println!("commitments: {:?}", output.2);
+    println!("commitment_tree_root: {:?}", output.3);
+
+    assert!(
+        nssa::verify_privacy_execution(receipt, &output.0, &output.1, &output.2, &output.3).is_ok()
+    );
 }
diff --git a/risc0-selective-privacy-poc/outer_methods/guest/src/bin/outer.rs b/risc0-selective-privacy-poc/outer_methods/guest/src/bin/outer.rs
index a8ba142..62b4b33 100644
--- a/risc0-selective-privacy-poc/outer_methods/guest/src/bin/outer.rs
+++ b/risc0-selective-privacy-poc/outer_methods/guest/src/bin/outer.rs
@@ -115,5 +115,6 @@ fn main() {
         public_inputs_outputs,
         nullifiers,
         private_output_commitments,
+        commitment_tree_root,
     ));
 }
diff --git a/risc0-selective-privacy-poc/src/lib.rs b/risc0-selective-privacy-poc/src/lib.rs
index 711663e..f8e095e 100644
--- a/risc0-selective-privacy-poc/src/lib.rs
+++ b/risc0-selective-privacy-poc/src/lib.rs
@@ -1,11 +1,12 @@
-use outer_methods::OUTER_ELF;
+use outer_methods::{OUTER_ELF, OUTER_ID};
 use program::Program;
 use rand::{rngs::OsRng, Rng};
-use risc0_zkvm::{
-    default_executor, default_prover, ExecutorEnv, ExecutorEnvBuilder, ProveInfo, Receipt,
+use risc0_zkvm::{default_executor, default_prover, ExecutorEnv, ExecutorEnvBuilder, Receipt};
+use toy_example_core::{
+    account::Account,
+    input::InputVisibiility,
+    types::{Commitment, Nonce, Nullifier},
 };
-use serde::{Deserialize, Serialize};
-use toy_example_core::{account::Account, input::InputVisibiility, types::Nonce};
 
 pub mod program;
 
@@ -13,7 +14,8 @@ pub fn new_random_nonce() -> Nonce {
     let mut rng = OsRng;
     std::array::from_fn(|_| rng.gen())
 }
-pub(crate) fn write_inputs<P: Program>(
+
+fn write_inputs<P: Program>(
     input_accounts: &[Account],
     instruction_data: &P::InstructionData,
     env_builder: &mut ExecutorEnvBuilder,
@@ -25,7 +27,7 @@ pub(crate) fn write_inputs<P: Program>(
     Ok(())
 }
 
-pub(crate) fn execute_and_prove<P: Program>(
+fn execute_and_prove<P: Program>(
     input_accounts: &[Account],
     instruction_data: &P::InstructionData,
 ) -> Result<(Receipt, Vec<Account>), ()> {
@@ -69,7 +71,7 @@ pub fn prove_privacy_execution<P: Program>(
     instruction_data: &P::InstructionData,
     visibilities: &[InputVisibiility],
     commitment_tree_root: [u32; 8],
-) -> Result<ProveInfo, ()> {
+) -> Result<Receipt, ()> {
     // Prove inner program and get post state of the accounts
     let num_inputs = inputs.len();
     let (inner_receipt, inputs_outputs) = execute_and_prove::<P>(inputs, instruction_data)?;
@@ -91,5 +93,27 @@ pub fn prove_privacy_execution<P: Program>(
 
     let prover = default_prover();
     let prove_info = prover.prove(env, OUTER_ELF).unwrap();
-    Ok(prove_info)
+    Ok(prove_info.receipt)
+}
+
+pub fn verify_privacy_execution(
+    receipt: Receipt,
+    public_accounts_inputs_outputs: &[Account],
+    nullifiers: &[Nullifier],
+    private_output_commitments: &[Commitment],
+    commitment_tree_root: &[u32; 8],
+) -> Result<(), ()> {
+    let output: (Vec<Account>, Vec<Nullifier>, Vec<Commitment>, [u32; 8]) =
+        receipt.journal.decode().unwrap();
+    let expected_output = (
+        public_accounts_inputs_outputs.to_vec(),
+        nullifiers.to_vec(),
+        private_output_commitments.to_vec(),
+        commitment_tree_root.to_owned(),
+    );
+    if output != expected_output {
+        return Err(());
+    } else {
+        receipt.verify(OUTER_ID).map_err(|_| ())
+    }
 }