refactor client

risc0-selective-privacy-poc/examples/happy_path.rs

@@ -17,7 +17,7 @@ fn main() {
     let mut sequencer = MockedSequencer::new();
     let addresses = sequencer.addresses();
     println!("addresses: {:?}", addresses);
-    println!("1️⃣ 🚀 Initial balances");
+    println!("🚀 Initial balances");
     sequencer.print();
 
     // A public execution of the Transfer Program
@@ -26,7 +26,7 @@ fn main() {
     sequencer
         .invoke_public_execution::<TransferProgram>(&[sender_addr, receiver_addr], 10)
         .unwrap();
-    println!("2️⃣ 🚀 Balances after transfer");
+    println!("🚀 Balances after transfer");
     sequencer.print();
 
     // A shielded execution of the Transfer Program
@@ -36,25 +36,25 @@ fn main() {
     sequencer.print();
 
     // A private execution of the Transfer Program
-    let private_account_1 = MockedClient::transfer_private(
-        &private_account_2,
+    let [_, private_account_1] = MockedClient::transfer_private(
+        private_account_2,
         &ACCOUNTS_PRIVATE_KEYS[1], // <-- this is shifted 🫠
         &addresses[3],
         8,
         &mut sequencer,
     );
-    println!("3️⃣ 🚀 Balances after shielded execution");
+    println!("🚀 Balances after shielded execution");
     sequencer.print();
 
     // A deshielded execution of the Transfer Program
     MockedClient::transfer_deshielded(
-        &private_account_1,
+        private_account_1,
         &ACCOUNTS_PRIVATE_KEYS[0],
         &addresses[0],
         1,
         &mut sequencer,
     );
-    println!("4️⃣ 🚀 Balances after deshielded execution");
+    println!("🚀 Balances after deshielded execution");
     sequencer.print();
 
     // A public execution of the Pinata program
@@ -62,6 +62,6 @@ fn main() {
     sequencer
         .invoke_public_execution::<PinataProgram>(&[addresses[0], addresses[3]], preimage)
         .unwrap();
-    println!("5️⃣ 🚀 Balances after public piñata execution");
+    println!("🚀 Balances after public piñata execution");
     sequencer.print();
 }

risc0-selective-privacy-poc/examples/mocked_components/client/mod.rs

@@ -20,4 +20,29 @@ impl MockedClient {
         let nonce = [0; 8];
         Account::new(address, nonce)
     }
+
+    fn prove_and_send_to_sequencer<P: nssa::Program>(
+        input_accounts: &[Account],
+        instruction_data: P::InstructionData,
+        visibilities: &[InputVisibiility],
+        commitment_tree_root: [u32; 8],
+        sequencer: &mut MockedSequencer,
+    ) -> Vec<Account> {
+        let (receipt, private_outputs) = nssa::invoke_privacy_execution::<P>(
+            input_accounts,
+            instruction_data,
+            visibilities,
+            commitment_tree_root,
+        )
+        .unwrap();
+        let output: (Vec<Account>, Vec<Nullifier>, Vec<Commitment>, [u32; 8]) =
+            receipt.journal.decode().unwrap();
+
+        // Send to te sequencer
+        sequencer
+            .invoke_privacy_execution(receipt, &output.0, &output.1, &output.2)
+            .unwrap();
+
+        private_outputs
+    }
 }

risc0-selective-privacy-poc/examples/mocked_components/client/transfer_deshielded.rs

@@ -4,7 +4,7 @@ use core::types::{Address, Commitment, Key, Nullifier};
 
 use nssa::program::TransferProgram;
 
-use super::{MockedSequencer, MockedClient};
+use super::{MockedClient, MockedSequencer};
 
 impl MockedClient {
     /// A shielded execution of the Transfer program
@@ -19,25 +19,16 @@ impl MockedClient {
         let commitment_tree_root = sequencer.get_commitment_tree_root();
         let receiver_addr = to_address;
         let mut receiver_account = Self::fresh_account_for_mint(*receiver_addr);
-        let visibilities = vec![InputVisibiility::Public, InputVisibiility::Private(None)];
-        let (receipt, nonces) = nssa::invoke_privacy_execution::<TransferProgram>(
-            &[sender_account, receiver_account.clone()],
+        let visibilities = [InputVisibiility::Public, InputVisibiility::Private(None)];
+        let private_outputs = Self::prove_and_send_to_sequencer::<TransferProgram>(
+            &[sender_account, receiver_account],
             balance_to_move,
             &visibilities,
             commitment_tree_root,
-        )
-        .unwrap();
-
-        // Assemble the private account
-        receiver_account.nonce = nonces[1];
-        receiver_account.balance = balance_to_move;
-        let output: (Vec<Account>, Vec<Nullifier>, Vec<Commitment>, [u32; 8]) =
-            receipt.journal.decode().unwrap();
-
-        // Send to te sequencer
-        sequencer
-            .invoke_privacy_execution(receipt, &output.0, &output.1, &output.2)
-            .unwrap();
-        receiver_account
+            sequencer,
+        );
+        let [receiver_private_account] = private_outputs.try_into().unwrap();
+        receiver_private_account
     }
+
 }

risc0-selective-privacy-poc/examples/mocked_components/client/transfer_private.rs

@@ -9,12 +9,12 @@ use super::{MockedClient, MockedSequencer};
 impl MockedClient {
     /// A private execution of the Transfer program
     pub fn transfer_private(
-        from_account: &Account,
+        from_account: Account,
         from_account_pk: &Key,
         to_address: &Address,
         balance_to_move: u128,
         sequencer: &mut MockedSequencer,
-    ) -> Account {
+    ) -> [Account; 2] {
         // All of this is executed locally by the sender
         let commitment_tree_root = sequencer.get_commitment_tree_root();
         let receiver_addr = to_address;
@@ -25,24 +25,15 @@ impl MockedClient {
             InputVisibiility::Private(Some((from_account_pk.clone(), sender_commitment_auth_path))),
             InputVisibiility::Private(None),
         ];
-        let (receipt, nonces) = nssa::invoke_privacy_execution::<TransferProgram>(
-            &[from_account.clone(), receiver_account.clone()],
+
+        let private_outputs = Self::prove_and_send_to_sequencer::<TransferProgram>(
+            &[from_account, receiver_account],
             balance_to_move,
             &visibilities,
             commitment_tree_root,
-        )
-        .unwrap();
-        let output: (Vec<Account>, Vec<Nullifier>, Vec<Commitment>, [u32; 8]) =
-            receipt.journal.decode().unwrap();
+            sequencer,
+        );
 
-        // Send to te sequencer
-        sequencer
-            .invoke_privacy_execution(receipt, &output.0, &output.1, &output.2)
-            .unwrap();
-
-        // Assemble the private account
-        receiver_account.nonce = nonces[1];
-        receiver_account.balance = balance_to_move;
-        receiver_account
+        private_outputs.try_into().unwrap()
     }
 }

risc0-selective-privacy-poc/examples/mocked_components/client/transfer_shielded.rs

@@ -8,7 +8,7 @@ use super::{MockedClient, MockedSequencer};
 
 impl MockedClient {
     pub fn transfer_deshielded(
-        from_account: &Account,
+        from_account: Account,
         from_account_pk: &Key,
         to_address: &Address,
         balance_to_move: u128,
@@ -24,20 +24,12 @@ impl MockedClient {
             InputVisibiility::Private(Some((from_account_pk.clone(), sender_commitment_auth_path))),
             InputVisibiility::Public,
         ];
-        let (receipt, nonces) = nssa::invoke_privacy_execution::<TransferProgram>(
-            &[from_account.clone(), to_account],
+        let private_outputs = Self::prove_and_send_to_sequencer::<TransferProgram>(
+            &[from_account, to_account],
             balance_to_move,
             &visibilities,
             commitment_tree_root,
-        )
-        .unwrap();
-
-        let output: (Vec<Account>, Vec<Nullifier>, Vec<Commitment>, [u32; 8]) =
-            receipt.journal.decode().unwrap();
-
-        // Send to te sequencer
-        sequencer
-            .invoke_privacy_execution(receipt, &output.0, &output.1, &output.2)
-            .unwrap();
+            sequencer,
+        );
     }
 }

risc0-selective-privacy-poc/examples/mocked_components/sequencer/mod.rs

@@ -5,9 +5,7 @@ use core::{
 };
 use std::collections::{BTreeMap, HashSet};
 
-use nssa;
 use program_methods::{PINATA_ID, TRANSFER_ID, TRANSFER_MULTIPLE_ID};
-use risc0_zkvm::Receipt;
 use sparse_merkle_tree::SparseMerkleTree;
 
 use super::ACCOUNTS_PRIVATE_KEYS;

risc0-selective-privacy-poc/src/lib.rs

@@ -1,7 +1,7 @@
 use core::{
     account::Account,
     input::InputVisibiility,
-    types::{Commitment, Nonce, Nullifier},
+    types::{AuthenticationPath, Commitment, Key, Nonce, Nullifier},
 };
 use program_methods::{OUTER_ELF, OUTER_ID};
 use rand::{rngs::OsRng, Rng};
@@ -66,12 +66,32 @@ pub fn execute<P: Program>(
     Ok(inputs_outputs)
 }
 
+pub fn extract_private_outputs_from_inner_results(
+    inputs_outputs: &[Account],
+    num_inputs: usize,
+    visibilities: &[InputVisibiility],
+    nonces: &[Nonce],
+) -> Vec<Account> {
+    inputs_outputs
+        .iter()
+        .skip(num_inputs)
+        .zip(visibilities)
+        .zip(nonces)
+        .filter(|((_, visibility), _)| matches!(visibility, InputVisibiility::Private(_)))
+        .map(|((account, _), nonce)| {
+            let mut this = account.clone();
+            this.nonce = *nonce;
+            this
+        })
+        .collect()
+}
+
 pub fn invoke_privacy_execution<P: Program>(
     inputs: &[Account],
     instruction_data: P::InstructionData,
     visibilities: &[InputVisibiility],
     commitment_tree_root: [u32; 8],
-) -> Result<(Receipt, Vec<Nonce>), ()> {
+) -> Result<(Receipt, Vec<Account>), ()> {
     // Prove inner program and get post state of the accounts
     let num_inputs = inputs.len();
     let (inner_receipt, inputs_outputs) = execute_and_prove_inner::<P>(inputs, instruction_data)?;
@@ -93,7 +113,13 @@ pub fn invoke_privacy_execution<P: Program>(
 
     let prover = default_prover();
     let prove_info = prover.prove(env, OUTER_ELF).unwrap();
-    Ok((prove_info.receipt, output_nonces))
+    let private_outputs = extract_private_outputs_from_inner_results(
+        &inputs_outputs,
+        num_inputs,
+        &visibilities,
+        &output_nonces,
+    );
+    Ok((prove_info.receipt, private_outputs))
 }
 
 pub fn verify_privacy_execution(