From 8feb17db1a3d373b5f84871c832b522fc7dc9d0e Mon Sep 17 00:00:00 2001
From: Sergio Chouhy <sergio.chouhy@gmail.com>
Date: Wed, 16 Jul 2025 17:13:26 -0300
Subject: [PATCH] refactor

---
 risc0-selective-privacy-poc/src/lib.rs        |  2 -
 .../src/private_execution.rs                  | 32 +--------------
 risc0-selective-privacy-poc/src/program.rs    | 39 ++++++++++++++++++-
 3 files changed, 38 insertions(+), 35 deletions(-)

diff --git a/risc0-selective-privacy-poc/src/lib.rs b/risc0-selective-privacy-poc/src/lib.rs
index ce39e62..7d9f777 100644
--- a/risc0-selective-privacy-poc/src/lib.rs
+++ b/risc0-selective-privacy-poc/src/lib.rs
@@ -9,9 +9,7 @@ struct TransferProgram;
 
 impl Program for TransferProgram {
     const PROGRAM_ID: [u32; 8] = TRANSFER_ID;
-
     const PROGRAM_ELF: &[u8] = TRANSFER_ELF;
-
     // Amount to transfer
     type InstructionData = u128;
 }
diff --git a/risc0-selective-privacy-poc/src/private_execution.rs b/risc0-selective-privacy-poc/src/private_execution.rs
index 1576cef..1cb3ceb 100644
--- a/risc0-selective-privacy-poc/src/private_execution.rs
+++ b/risc0-selective-privacy-poc/src/private_execution.rs
@@ -10,7 +10,7 @@ use toy_example_core::{
 };
 use transfer_methods::{TRANSFER_ELF, TRANSFER_ID};
 
-use crate::program::{execute_and_prove, Program};
+use crate::program::{execute_and_prove, prove_privacy_execution, Program};
 use crate::TransferProgram;
 
 pub fn new_random_nonce() -> Nonce {
@@ -23,36 +23,6 @@ fn mint_fresh_account(address: Address) -> Account {
     Account::new(address, nonce)
 }
 
-fn prove_privacy_execution<P: Program>(
-    inputs: &[Account],
-    instruction_data: &P::InstructionData,
-    visibilities: &[InputVisibiility],
-    commitment_tree_root: [u32; 8],
-) -> Result<ProveInfo, ()> {
-    // 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)?;
-
-    // Sample fresh random nonces for the outputs of this execution
-    let output_nonces: Vec<_> = (0..num_inputs).map(|_| new_random_nonce()).collect();
-
-    // Prove outer program.
-    // This computes the nullifiers for the input accounts and commitments for the output accounts.
-    let mut env_builder = ExecutorEnv::builder();
-    env_builder.add_assumption(inner_receipt);
-    env_builder.write(&(num_inputs as u32)).unwrap();
-    env_builder.write(&inputs_outputs).unwrap();
-    env_builder.write(&visibilities).unwrap();
-    env_builder.write(&output_nonces).unwrap();
-    env_builder.write(&commitment_tree_root).unwrap();
-    env_builder.write(&TRANSFER_ID).unwrap();
-    let env = env_builder.build().unwrap();
-
-    let prover = default_prover();
-    let prove_info = prover.prove(env, OUTER_ELF).unwrap();
-    Ok(prove_info)
-}
-
 /// A private execution of the transfer function.
 /// This actually "burns" a sender private account and "mints" two new private accounts:
 /// one for the recipient with the transferred balance, and another owned by the sender with the remaining balance.
diff --git a/risc0-selective-privacy-poc/src/program.rs b/risc0-selective-privacy-poc/src/program.rs
index bf31cf1..ade94c0 100644
--- a/risc0-selective-privacy-poc/src/program.rs
+++ b/risc0-selective-privacy-poc/src/program.rs
@@ -1,6 +1,11 @@
-use risc0_zkvm::{default_executor, default_prover, ExecutorEnv, ExecutorEnvBuilder, Receipt};
+use outer_methods::OUTER_ELF;
+use risc0_zkvm::{
+    default_executor, default_prover, ExecutorEnv, ExecutorEnvBuilder, ProveInfo, Receipt,
+};
 use serde::{Deserialize, Serialize};
-use toy_example_core::account::Account;
+use toy_example_core::{account::Account, input::InputVisibiility};
+
+use crate::private_execution::new_random_nonce;
 
 pub(crate) trait Program {
     const PROGRAM_ID: [u32; 8];
@@ -58,3 +63,33 @@ pub(crate) fn execute<P: Program>(
 
     Ok(inputs_outputs)
 }
+
+pub(crate) fn prove_privacy_execution<P: Program>(
+    inputs: &[Account],
+    instruction_data: &P::InstructionData,
+    visibilities: &[InputVisibiility],
+    commitment_tree_root: [u32; 8],
+) -> Result<ProveInfo, ()> {
+    // 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)?;
+
+    // Sample fresh random nonces for the outputs of this execution
+    let output_nonces: Vec<_> = (0..num_inputs).map(|_| new_random_nonce()).collect();
+
+    // Prove outer program.
+    // This computes the nullifiers for the input accounts and commitments for the output accounts.
+    let mut env_builder = ExecutorEnv::builder();
+    env_builder.add_assumption(inner_receipt);
+    env_builder.write(&(num_inputs as u32)).unwrap();
+    env_builder.write(&inputs_outputs).unwrap();
+    env_builder.write(&visibilities).unwrap();
+    env_builder.write(&output_nonces).unwrap();
+    env_builder.write(&commitment_tree_root).unwrap();
+    env_builder.write(&P::PROGRAM_ID).unwrap();
+    let env = env_builder.build().unwrap();
+
+    let prover = default_prover();
+    let prove_info = prover.prove(env, OUTER_ELF).unwrap();
+    Ok(prove_info)
+}