add program output struct

risc0-selective-privacy-poc/core/src/types.rs

@@ -11,6 +11,12 @@ pub type Key = [u32; 8];
 pub type AuthenticationPath = [[u32; 8]; 32];
 pub type ProgramId = [u32; 8];
 
+#[derive(Serialize, Deserialize)]
+pub struct ProgramOutput {
+    pub accounts_pre: Vec<Account>,
+    pub accounts_post: Vec<Account>,
+}
+
 #[derive(Serialize, Deserialize)]
 pub struct PrivacyExecutionOutput {
     pub public_accounts_pre: Vec<Account>,

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

@@ -1,4 +1,3 @@
-
 use crate::mocked_components::client::MockedClient;
 
 pub mod client;

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

@@ -1,4 +1,7 @@
-use core::{account::Account, types::Address};
+use core::{
+    account::Account,
+    types::{Address, ProgramOutput},
+};
 
 use super::MockedSequencer;
 
@@ -16,20 +19,17 @@ impl MockedSequencer {
             .collect::<Result<_, _>>()?;
 
         // Execute the program
-        let inputs_outputs = nssa::execute_onchain::<P>(&input_accounts, instruction_data)?;
+        let program_output = nssa::execute_onchain::<P>(&input_accounts, instruction_data)?;
 
         // Perform consistency checks
-        if !self.inputs_outputs_are_consistent(&input_accounts, &inputs_outputs) {
+        if !self.program_output_is_valid(&input_accounts, &program_output) {
             return Err(());
         }
 
         // Update the accounts states
-        inputs_outputs
-            .into_iter()
-            .skip(input_accounts.len())
-            .for_each(|account_post_state| {
-                self.accounts.insert(account_post_state.address, account_post_state);
-            });
+        program_output.accounts_post.into_iter().for_each(|account_post_state| {
+            self.accounts.insert(account_post_state.address, account_post_state);
+        });
         Ok(())
     }
 
@@ -37,22 +37,24 @@ impl MockedSequencer {
     /// `input_accounts` are the accounts provided as inputs to the program.
     /// `inputs_outputs` is the program output, which should consist of the accounts pre and
     /// post-states.
-    fn inputs_outputs_are_consistent(&self, input_accounts: &[Account], inputs_outputs: &[Account]) -> bool {
+    fn program_output_is_valid(&self, input_accounts: &[Account], program_output: &ProgramOutput) -> bool {
         let num_inputs = input_accounts.len();
 
-        // Fail if the number of accounts pre and post-states is inconsistent with the number of
-        // inputs.
-        if inputs_outputs.len() != num_inputs * 2 {
+        // Fail if the number of accounts pre and post-states is differ
+        if program_output.accounts_pre.len() != program_output.accounts_post.len() {
             return false;
         }
 
         // Fail if the accounts pre-states do not coincide with the input accounts.
-        let (claimed_accounts_pre, accounts_post) = inputs_outputs.split_at(num_inputs);
-        if claimed_accounts_pre != input_accounts {
+        if program_output.accounts_pre != input_accounts {
             return false;
         }
 
-        for (account_pre, account_post) in input_accounts.iter().zip(accounts_post) {
+        for (account_pre, account_post) in program_output
+            .accounts_pre
+            .iter()
+            .zip(program_output.accounts_post.iter())
+        {
             // Fail if the program modified the addresses of the input accounts
             if account_pre.address != account_post.address {
                 return false;
@@ -69,7 +71,7 @@ impl MockedSequencer {
         }
 
         let total_balance_pre: u128 = input_accounts.iter().map(|account| account.balance).sum();
-        let total_balance_post: u128 = accounts_post.iter().map(|account| account.balance).sum();
+        let total_balance_post: u128 = program_output.accounts_post.iter().map(|account| account.balance).sum();
         // Fail if the execution didn't preserve the total supply.
         if total_balance_pre != total_balance_post {
             return false;

risc0-selective-privacy-poc/program_methods/guest/src/bin/outer.rs

@@ -1,7 +1,7 @@
 use core::{
     account::Account,
     compute_nullifier, hash, is_in_tree,
-    types::{Nonce, PrivacyExecutionOutput, ProgramId},
+    types::{Nonce, PrivacyExecutionOutput, ProgramId, ProgramOutput},
     visibility::AccountVisibility,
 };
 use risc0_zkvm::{guest::env, serde::to_vec};
@@ -29,18 +29,18 @@ use risc0_zkvm::{guest::env, serde::to_vec};
 /// - The commitments for the ouput private accounts.
 /// - The commitment tree root used for the authentication path verifications.
 fn main() {
-    let num_inputs: u32 = env::read();
     // Read inputs and outputs
-    let mut inputs_outputs: Vec<Account> = env::read();
-    assert_eq!(inputs_outputs.len() as u32, num_inputs * 2);
+    let mut inner_program_output: ProgramOutput = env::read();
+    let num_inputs = inner_program_output.accounts_pre.len();
+    assert_eq!(inner_program_output.accounts_post.len(), num_inputs);
 
     // Read visibilities
     let account_visibilities: Vec<AccountVisibility> = env::read();
-    assert_eq!(account_visibilities.len() as u32, num_inputs);
+    assert_eq!(account_visibilities.len(), num_inputs);
 
     // Read nonces for outputs
     let output_nonces: Vec<Nonce> = env::read();
-    assert_eq!(output_nonces.len() as u32, num_inputs);
+    assert_eq!(output_nonces.len(), num_inputs);
 
     // Read root and program id.
     let commitment_tree_root: [u32; 8] = env::read();
@@ -48,14 +48,10 @@ fn main() {
 
     // Verify pre states and post states of accounts are consistent
     // with the execution of the `program_id` program
-    env::verify(program_id, &to_vec(&inputs_outputs).unwrap()).unwrap();
-
-    // Split inputs_outputs into two separate vectors
-    let (inputs, mut outputs) = {
-        let outputs = inputs_outputs.split_off(num_inputs as usize);
-        (inputs_outputs, outputs)
-    };
+    env::verify(program_id, &to_vec(&inner_program_output).unwrap()).unwrap();
 
+    let inputs = inner_program_output.accounts_pre;
+    let mut outputs = inner_program_output.accounts_post;
     let mut nullifiers = Vec::new();
     for (visibility, input_account) in account_visibilities.iter().zip(inputs.iter()) {
         match visibility {

risc0-selective-privacy-poc/program_methods/guest/src/bin/pinata.rs

@@ -1,4 +1,4 @@
-use core::{account::Account, hash};
+use core::{account::Account, hash, types::ProgramOutput};
 use risc0_zkvm::guest::env;
 
 const TARGET_HASH: [u32; 8] = [
@@ -36,10 +36,10 @@ fn main() {
     pinata_account_post.balance -= PINATA_PRIZE;
     winner_account_post.balance += PINATA_PRIZE;
 
-    env::commit(&vec![
-        pinata_account,
-        winner_account,
-        pinata_account_post,
-        winner_account_post,
-    ]);
+    let output = ProgramOutput {
+        accounts_pre: vec![pinata_account, winner_account],
+        accounts_post: vec![pinata_account_post, winner_account_post],
+    };
+
+    env::commit(&output);
 }

risc0-selective-privacy-poc/program_methods/guest/src/bin/transfer.rs

@@ -1,4 +1,4 @@
-use core::account::Account;
+use core::{account::Account, types::ProgramOutput};
 use risc0_zkvm::guest::env;
 
 /// A transfer of balance program.
@@ -22,5 +22,10 @@ fn main() {
     sender_post.balance -= balance_to_move;
     receiver_post.balance += balance_to_move;
 
-    env::commit(&vec![sender, receiver, sender_post, receiver_post]);
+    let output = ProgramOutput {
+        accounts_pre: vec![sender, receiver],
+        accounts_post: vec![sender_post, receiver_post],
+    };
+
+    env::commit(&output);
 }

risc0-selective-privacy-poc/program_methods/guest/src/bin/transfer_multiple.rs

@@ -1,4 +1,4 @@
-use core::account::Account;
+use core::{account::Account, types::ProgramOutput};
 use risc0_zkvm::guest::env;
 
 /// A transfer of balance program with one sender and multiple recipients.
@@ -27,21 +27,18 @@ fn main() {
 
     // Create accounts post states, with updated balances
     let mut sender_post = sender.clone();
-    let mut receivers_post = recipients.clone();
+    let mut recipients_post = recipients.clone();
 
     // Transfer balances
     sender_post.balance -= total_balance_to_move;
-    for (receiver, balance_for_receiver) in receivers_post.iter_mut().zip(target_balances) {
+    for (receiver, balance_for_receiver) in recipients_post.iter_mut().zip(target_balances) {
         receiver.balance += balance_for_receiver;
     }
 
-    // Flatten pre and post states for output
-    let inputs_outputs: Vec<Account> = vec![sender]
-        .into_iter()
-        .chain(recipients)
-        .chain(vec![sender_post])
-        .chain(receivers_post)
-        .collect();
+    let output = ProgramOutput {
+        accounts_pre: vec![sender].into_iter().chain(recipients).collect(),
+        accounts_post: vec![sender_post].into_iter().chain(recipients_post).collect(),
+    };
 
-    env::commit(&inputs_outputs);
+    env::commit(&output);
 }

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

@@ -1,6 +1,6 @@
 use core::{
     account::Account,
-    types::{Commitment, Nonce, Nullifier},
+    types::{Commitment, Nonce, Nullifier, ProgramOutput},
     visibility::AccountVisibility,
 };
 use program_methods::{OUTER_ELF, OUTER_ID};
@@ -33,7 +33,7 @@ fn write_inputs<P: Program>(
 fn execute_and_prove_inner<P: Program>(
     input_accounts: &[Account],
     instruction_data: P::InstructionData,
-) -> Result<(Receipt, Vec<Account>), ()> {
+) -> Result<Receipt, ()> {
     // Write inputs to the program
     let mut env_builder = ExecutorEnv::builder();
     write_inputs::<P>(input_accounts, instruction_data, &mut env_builder)?;
@@ -42,25 +42,20 @@ fn execute_and_prove_inner<P: Program>(
     // Prove the program
     let prover = default_prover();
     let prove_info = prover.prove(env, P::PROGRAM_ELF).map_err(|_| ())?;
-    let receipt = prove_info.receipt;
-
-    // Get proof and (inputs and) outputs
-    let inputs_outputs: Vec<Account> = receipt.journal.decode().map_err(|_| ())?;
-
-    Ok((receipt, inputs_outputs))
+    Ok(prove_info.receipt)
 }
 
 /// Builds the private outputs from the results of the execution of an inner program.
 /// Populates the nonces with the ones provided.
 fn build_private_outputs_from_inner_results(
-    inputs_outputs: &[Account],
+    inner_program_output: &ProgramOutput,
     num_inputs: usize,
     visibilities: &[AccountVisibility],
     nonces: &[Nonce],
 ) -> Vec<Account> {
-    inputs_outputs
+    inner_program_output
+        .accounts_post
         .iter()
-        .skip(num_inputs)
         .zip(visibilities)
         .zip(nonces)
         .filter(|((_, visibility), _)| matches!(visibility, AccountVisibility::Private(_)))
@@ -77,7 +72,7 @@ fn build_private_outputs_from_inner_results(
 pub fn execute_onchain<P: Program>(
     input_accounts: &[Account],
     instruction_data: P::InstructionData,
-) -> Result<Vec<Account>, ()> {
+) -> Result<ProgramOutput, ()> {
     // Write inputs to the program
     let mut env_builder = ExecutorEnv::builder();
     write_inputs::<P>(input_accounts, instruction_data, &mut env_builder)?;
@@ -88,9 +83,7 @@ pub fn execute_onchain<P: Program>(
     let session_info = executor.execute(env, P::PROGRAM_ELF).map_err(|_| ())?;
 
     // Get (inputs and) outputs
-    let inputs_outputs: Vec<Account> = session_info.journal.decode().map_err(|_| ())?;
-
-    Ok(inputs_outputs)
+    session_info.journal.decode().map_err(|_| ())
 }
 
 /// Executes and proves the inner program `P` and executes and proves the outer program on top of it.
@@ -104,7 +97,8 @@ pub fn execute_offchain<P: Program>(
 ) -> 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)?;
+    let inner_receipt = execute_and_prove_inner::<P>(inputs, instruction_data)?;
+    let inner_program_output: ProgramOutput = inner_receipt.journal.decode().map_err(|_| ())?;
 
     // Sample fresh random nonces for the outputs of this execution
     let output_nonces: Vec<_> = (0..num_inputs).map(|_| new_random_nonce()).collect();
@@ -112,8 +106,7 @@ pub fn execute_offchain<P: Program>(
     // Prove outer program.
     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(&inner_program_output).unwrap();
     env_builder.write(&visibilities).unwrap();
     env_builder.write(&output_nonces).unwrap();
     env_builder.write(&commitment_tree_root).unwrap();
@@ -124,7 +117,7 @@ pub fn execute_offchain<P: Program>(
 
     // Build private accounts.
     let private_outputs =
-        build_private_outputs_from_inner_results(&inputs_outputs, num_inputs, visibilities, &output_nonces);
+        build_private_outputs_from_inner_results(&inner_program_output, num_inputs, visibilities, &output_nonces);
     Ok((prove_info.receipt, private_outputs))
 }