From 1f5eeba2a34cffcc3607a8b3b90a4cbbc074c42a Mon Sep 17 00:00:00 2001 From: Sergio Chouhy Date: Thu, 17 Jul 2025 13:34:12 -0300 Subject: [PATCH] wip --- risc0-selective-privacy-poc/Cargo.toml | 1 + .../examples/sequencer.rs | 21 +++++- .../examples/sequencer_mock/mod.rs | 67 ++++++++++++++----- .../sparse_merkle_tree/src/lib.rs | 4 ++ risc0-selective-privacy-poc/src/lib.rs | 4 +- 5 files changed, 80 insertions(+), 17 deletions(-) diff --git a/risc0-selective-privacy-poc/Cargo.toml b/risc0-selective-privacy-poc/Cargo.toml index 055271f..06b2f83 100644 --- a/risc0-selective-privacy-poc/Cargo.toml +++ b/risc0-selective-privacy-poc/Cargo.toml @@ -16,3 +16,4 @@ sparse-merkle-tree = { path = "./sparse_merkle_tree/" } cuda = ["risc0-zkvm/cuda"] default = [] prove = ["risc0-zkvm/prove"] + diff --git a/risc0-selective-privacy-poc/examples/sequencer.rs b/risc0-selective-privacy-poc/examples/sequencer.rs index da98a47..e1e2506 100644 --- a/risc0-selective-privacy-poc/examples/sequencer.rs +++ b/risc0-selective-privacy-poc/examples/sequencer.rs @@ -1,7 +1,26 @@ +use nssa::program::TransferProgram; + use crate::sequencer_mock::MockedSequencer; mod sequencer_mock; fn main() { - let sequencer = MockedSequencer::new(); + let mut sequencer = MockedSequencer::new(); + let addresses = sequencer.addresses(); + println!("Initial balances"); + sequencer.print(); + + // A public execution of the Transfer Program + let sender_addr = addresses[0]; + let receiver_addr = addresses[1]; + sequencer + .invoke_public::(&[sender_addr, receiver_addr], 10) + .unwrap(); + + println!("Balances after transfer"); + sequencer.print(); + + + // A private execution of the Transfer Program + } diff --git a/risc0-selective-privacy-poc/examples/sequencer_mock/mod.rs b/risc0-selective-privacy-poc/examples/sequencer_mock/mod.rs index 53b9985..b74e73f 100644 --- a/risc0-selective-privacy-poc/examples/sequencer_mock/mod.rs +++ b/risc0-selective-privacy-poc/examples/sequencer_mock/mod.rs @@ -2,26 +2,26 @@ use core::{ account::Account, types::{Address, Key, Nullifier, ProgramId}, }; -use std::collections::{HashMap, HashSet}; +use std::collections::{BTreeMap, HashSet}; use nssa; use program_methods::{PINATA_ID, TRANSFER_ID, TRANSFER_MULTIPLE_ID}; use sparse_merkle_tree::SparseMerkleTree; pub struct MockedSequencer { - accounts: HashMap, + accounts: BTreeMap, commitment_tree: SparseMerkleTree, nullifier_set: HashSet, deployed_program_ids: HashSet, } -const ACCOUNTS_PRIVATE_KEYS: [Key; 5] = [[1; 8], [2; 8], [3; 8], [4; 8], [5; 8]]; -const ACCOUNTS_INITIAL_BALANCES: [u128; 5] = [100, 3, 900, 44, 0]; +const ACCOUNTS_PRIVATE_KEYS: [Key; 3] = [[1; 8], [2; 8], [3; 8]]; +const ACCOUNTS_INITIAL_BALANCES: [u128; 3] = [100, 1337, 0]; const DEPLOYED_PROGRAM_IDS: [ProgramId; 3] = [TRANSFER_ID, TRANSFER_MULTIPLE_ID, PINATA_ID]; impl MockedSequencer { pub fn new() -> Self { - let mut accounts: HashMap = ACCOUNTS_PRIVATE_KEYS + let mut accounts: BTreeMap = ACCOUNTS_PRIVATE_KEYS .iter() .cloned() .zip(ACCOUNTS_INITIAL_BALANCES) @@ -46,7 +46,7 @@ impl MockedSequencer { accounts, commitment_tree, nullifier_set, - deployed_program_ids: DEPLOYED_PROGRAM_IDS.iter().collect(), + deployed_program_ids: DEPLOYED_PROGRAM_IDS.iter().cloned().collect(), } } @@ -54,19 +54,32 @@ impl MockedSequencer { self.accounts.get(address).cloned() } - pub fn invoke_program_public( + pub fn invoke_public( &mut self, - program_id: ProgramId, input_account_addresses: &[Address], - instruction_data: &P::InstructionData, + instruction_data: P::InstructionData, ) -> Result<(), ()> { + // Fetch accounts let input_accounts: Vec = input_account_addresses .iter() - .map(|address| self.get_account(address).ok_or(|_| ())?) - .collect(); - let inputs_outputs = nssa::execute(input_accounts, instruction_data)?; + .map(|address| self.get_account(address).ok_or(())) + .collect::>()?; - self.inputs_outputs_are_consistent(&input_accounts, &inputs_outputs)? + // Execute + let inputs_outputs = nssa::execute::

(&input_accounts, &instruction_data)?; + + // Consistency checks + self.inputs_outputs_are_consistent(&input_accounts, &inputs_outputs)?; + + // Update accounts + inputs_outputs + .into_iter() + .skip(input_accounts.len()) + .for_each(|account_post_state| { + self.accounts + .insert(account_post_state.address, account_post_state); + }); + Ok(()) } fn inputs_outputs_are_consistent( @@ -91,12 +104,36 @@ impl MockedSequencer { if account_pre.nonce != account_post.nonce { return Err(()); } + // Redundant with previous checks, but better make it explicit. + if !self.accounts.contains_key(&account_post.address) { + return Err(()); + } } - let accounts_pre_total_balance = input_accounts.iter().map(|account| account.balance).sum(); - let accounts_post_total_balance = accounts_post.iter().map(|account| account.balance).sum(); + let accounts_pre_total_balance: u128 = + input_accounts.iter().map(|account| account.balance).sum(); + let accounts_post_total_balance: u128 = + accounts_post.iter().map(|account| account.balance).sum(); if accounts_pre_total_balance != accounts_post_total_balance { return Err(()); } return Ok(()); } + + pub fn addresses(&self) -> Vec

{ + self.accounts.keys().cloned().collect() + } + + pub fn print(&self) { + println!("{:<20} | {:>10}", "Address (first u32)", "Balance"); + println!("{:-<20}-+-{:-<10}", "", ""); + + for account in self.accounts.values() { + println!("{:<20x} | {:>10}", account.address[0], account.balance); + } + println!("{:-<20}-+-{:-<10}", "", ""); + println!("Commitments: {:?}", self.commitment_tree.values()); + println!("Nullifiers: {:?}", self.nullifier_set); + println!(""); + println!(""); + } } diff --git a/risc0-selective-privacy-poc/sparse_merkle_tree/src/lib.rs b/risc0-selective-privacy-poc/sparse_merkle_tree/src/lib.rs index 8d09d9c..3822116 100644 --- a/risc0-selective-privacy-poc/sparse_merkle_tree/src/lib.rs +++ b/risc0-selective-privacy-poc/sparse_merkle_tree/src/lib.rs @@ -111,6 +111,10 @@ impl SparseMerkleTree { path } + pub fn values(&self) -> HashSet { + self.values.clone() + } + pub fn verify_value_is_in_set(value: u32, path: [[u8; 32]; 32], root: [u8; 32]) -> bool { let mut hash = ONE_HASH; let mut current_index = value; diff --git a/risc0-selective-privacy-poc/src/lib.rs b/risc0-selective-privacy-poc/src/lib.rs index a44ae55..c7f4f15 100644 --- a/risc0-selective-privacy-poc/src/lib.rs +++ b/risc0-selective-privacy-poc/src/lib.rs @@ -3,13 +3,15 @@ use core::{ input::InputVisibiility, types::{Commitment, Nonce, Nullifier}, }; -use program::Program; use rand::{rngs::OsRng, Rng}; use risc0_zkvm::{default_executor, default_prover, ExecutorEnv, ExecutorEnvBuilder, Receipt}; use program_methods::{OUTER_ELF, OUTER_ID}; pub mod program; +pub use program::Program; + + pub fn new_random_nonce() -> Nonce { let mut rng = OsRng; std::array::from_fn(|_| rng.gen())