From 159fd52d254e5d3d42fdf5911176c2653694648f Mon Sep 17 00:00:00 2001 From: Sergio Chouhy Date: Sat, 19 Jul 2025 18:37:21 -0300 Subject: [PATCH] add program output struct --- risc0-selective-privacy-poc/core/src/types.rs | 6 ++++ .../examples/mocked_components/mod.rs | 1 - .../sequencer/process_public_execution.rs | 36 ++++++++++--------- .../program_methods/guest/src/bin/outer.rs | 22 +++++------- .../program_methods/guest/src/bin/pinata.rs | 14 ++++---- .../program_methods/guest/src/bin/transfer.rs | 9 +++-- .../guest/src/bin/transfer_multiple.rs | 19 +++++----- risc0-selective-privacy-poc/src/lib.rs | 31 +++++++--------- 8 files changed, 68 insertions(+), 70 deletions(-) diff --git a/risc0-selective-privacy-poc/core/src/types.rs b/risc0-selective-privacy-poc/core/src/types.rs index f742c30..a2bda29 100644 --- a/risc0-selective-privacy-poc/core/src/types.rs +++ b/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, + pub accounts_post: Vec, +} + #[derive(Serialize, Deserialize)] pub struct PrivacyExecutionOutput { pub public_accounts_pre: Vec, diff --git a/risc0-selective-privacy-poc/examples/mocked_components/mod.rs b/risc0-selective-privacy-poc/examples/mocked_components/mod.rs index 20f37c0..7ba4703 100644 --- a/risc0-selective-privacy-poc/examples/mocked_components/mod.rs +++ b/risc0-selective-privacy-poc/examples/mocked_components/mod.rs @@ -1,4 +1,3 @@ - use crate::mocked_components::client::MockedClient; pub mod client; diff --git a/risc0-selective-privacy-poc/examples/mocked_components/sequencer/process_public_execution.rs b/risc0-selective-privacy-poc/examples/mocked_components/sequencer/process_public_execution.rs index c6c1366..70c14b4 100644 --- a/risc0-selective-privacy-poc/examples/mocked_components/sequencer/process_public_execution.rs +++ b/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::>()?; // Execute the program - let inputs_outputs = nssa::execute_onchain::

(&input_accounts, instruction_data)?; + let program_output = nssa::execute_onchain::

(&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; diff --git a/risc0-selective-privacy-poc/program_methods/guest/src/bin/outer.rs b/risc0-selective-privacy-poc/program_methods/guest/src/bin/outer.rs index bf75dbb..5871895 100644 --- a/risc0-selective-privacy-poc/program_methods/guest/src/bin/outer.rs +++ b/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 = 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 = 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 = 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 { diff --git a/risc0-selective-privacy-poc/program_methods/guest/src/bin/pinata.rs b/risc0-selective-privacy-poc/program_methods/guest/src/bin/pinata.rs index 48c5a50..49c1e84 100644 --- a/risc0-selective-privacy-poc/program_methods/guest/src/bin/pinata.rs +++ b/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); } diff --git a/risc0-selective-privacy-poc/program_methods/guest/src/bin/transfer.rs b/risc0-selective-privacy-poc/program_methods/guest/src/bin/transfer.rs index 1e1a961..d600d67 100644 --- a/risc0-selective-privacy-poc/program_methods/guest/src/bin/transfer.rs +++ b/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); } diff --git a/risc0-selective-privacy-poc/program_methods/guest/src/bin/transfer_multiple.rs b/risc0-selective-privacy-poc/program_methods/guest/src/bin/transfer_multiple.rs index dc495f1..bdf1dda 100644 --- a/risc0-selective-privacy-poc/program_methods/guest/src/bin/transfer_multiple.rs +++ b/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 = 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); } diff --git a/risc0-selective-privacy-poc/src/lib.rs b/risc0-selective-privacy-poc/src/lib.rs index 08c17df..b28d4d8 100644 --- a/risc0-selective-privacy-poc/src/lib.rs +++ b/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( fn execute_and_prove_inner( input_accounts: &[Account], instruction_data: P::InstructionData, -) -> Result<(Receipt, Vec), ()> { +) -> Result { // Write inputs to the program let mut env_builder = ExecutorEnv::builder(); write_inputs::

(input_accounts, instruction_data, &mut env_builder)?; @@ -42,25 +42,20 @@ fn execute_and_prove_inner( // 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 = 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 { - 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( input_accounts: &[Account], instruction_data: P::InstructionData, -) -> Result, ()> { +) -> Result { // Write inputs to the program let mut env_builder = ExecutorEnv::builder(); write_inputs::

(input_accounts, instruction_data, &mut env_builder)?; @@ -88,9 +83,7 @@ pub fn execute_onchain( let session_info = executor.execute(env, P::PROGRAM_ELF).map_err(|_| ())?; // Get (inputs and) outputs - let inputs_outputs: Vec = 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( ) -> Result<(Receipt, Vec), ()> { // Prove inner program and get post state of the accounts let num_inputs = inputs.len(); - let (inner_receipt, inputs_outputs) = execute_and_prove_inner::

(inputs, instruction_data)?; + let inner_receipt = execute_and_prove_inner::

(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( // 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( // 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)) }