add public pre and post states to the output of the outer program

risc0-selective-privacy-poc/Cargo.toml

@@ -5,13 +5,15 @@ edition = "2021"
 
 [dependencies]
 risc0-zkvm = "2.2"
-toy-example-core = {path = "core"}
-transfer-methods = {path = "transfer_methods"}
-outer-methods = {path = "outer_methods"}
+toy-example-core = { path = "core" }
+transfer-methods = { path = "transfer_methods" }
+outer-methods = { path = "outer_methods" }
 serde = "1.0"
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }
+rand = "0.8"
 
 [features]
 cuda = ["risc0-zkvm/cuda"]
 default = []
-prove = ["risc0-zkvm/prove"]
+prove = ["risc0-zkvm/prove"]
+

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

@@ -52,8 +52,3 @@ pub fn compute_nullifier(commitment: &[u32; 8], private_key: &[u32; 8]) -> [u32;
     bytes_to_hash[8..].copy_from_slice(private_key);
     hash(&bytes_to_hash)
 }
-
-/// Dummy implementation
-pub fn new_random_nonce() -> [u32; 8] {
-    [0xcc, 0xaa, 0xff, 0xee, 0xcc, 0xaa, 0xff, 0xff]
-}

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

@@ -1,8 +1,4 @@
-use risc0_zkvm::{
-    guest::env,
-    serde::to_vec,
-    sha::{Impl, Sha256},
-};
+use risc0_zkvm::{guest::env, serde::to_vec};
 use toy_example_core::{
     account::{compute_nullifier, hash, is_in_commitment_tree, Account},
     input::InputVisibiility,
@@ -26,25 +22,32 @@ use toy_example_core::{
 fn main() {
     let num_inputs: u32 = env::read();
     // Read inputs and outputs
-    let mut inputs_outputs = Vec::new();
-    for _ in 0..(2 * num_inputs) {
-        let account: Account = env::read();
-        inputs_outputs.push(account);
-    }
+    let mut inputs_outputs: Vec<Account> = env::read();
+    assert_eq!(inputs_outputs.len() as u32, num_inputs * 2);
 
     // Read visibilities
-    let mut input_visibilities = Vec::new();
-    for _ in 0..num_inputs {
-        let input_visibility: InputVisibiility = env::read();
-        input_visibilities.push(input_visibility);
-    }
+    let input_visibilities: Vec<InputVisibiility> = env::read();
+    assert_eq!(input_visibilities.len() as u32, num_inputs);
+
+    // Read nonces for outputs
+    let output_nonces: Vec<[u32; 8]> = env::read();
+    assert_eq!(output_nonces.len() as u32, num_inputs);
 
     let commitment_tree_root: [u32; 8] = env::read();
     let program_id: [u32; 8] = env::read();
 
-    let inputs = inputs_outputs.iter().take(num_inputs as usize);
+    // 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)
+    };
+
     let mut nullifiers = Vec::new();
-    for (visibility, input_account) in input_visibilities.iter().zip(inputs) {
+    for (visibility, input_account) in input_visibilities.iter().zip(inputs.iter()) {
         match visibility {
             InputVisibiility::Private(Some(private_key)) => {
                 // Prove ownership of input accounts by proving
@@ -55,7 +58,7 @@ fn main() {
                 let commitment = input_account.commitment();
                 assert!(is_in_commitment_tree(commitment, commitment_tree_root));
                 // Compute nullifier to nullify this private input account.
-                let nullifier = compute_nullifier(&commitment, &private_key);
+                let nullifier = compute_nullifier(&commitment, private_key);
                 nullifiers.push(nullifier);
             }
             InputVisibiility::Private(None) => {
@@ -67,22 +70,52 @@ fn main() {
             InputVisibiility::Public => continue,
         }
     }
-    let outputs = inputs_outputs.iter().skip(num_inputs as usize);
-    let output_commitments: Vec<_> = outputs.map(|account| account.commitment()).collect();
 
     // Assert `program_id` program didn't modify address fields
-    for (account_pre, account_post) in inputs_outputs
-        .iter()
-        .take(num_inputs as usize)
-        .zip(inputs_outputs.iter().skip(num_inputs as usize))
-    {
+    for (account_pre, account_post) in inputs.iter().zip(outputs.iter()) {
         assert_eq!(account_pre.address, account_post.address);
     }
 
-    // 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();
+    // Insert new nonces in outputs (including public ones (?!))
+    outputs
+        .iter_mut()
+        .zip(output_nonces)
+        .for_each(|(account, new_nonce)| account.nonce = new_nonce);
+
+    // Compute private outputs commitments
+    let mut private_outputs = Vec::new();
+    for (output, visibility) in outputs.iter().zip(input_visibilities.iter()) {
+        match visibility {
+            InputVisibiility::Public => continue,
+            InputVisibiility::Private(_) => private_outputs.push(output),
+        }
+    }
+
+    // Get the list of public inputs pre states and their post states
+    let mut public_inputs_outputs = Vec::new();
+    for (account, visibility) in inputs
+        .iter()
+        .chain(outputs.iter())
+        .zip(input_visibilities.iter().chain(input_visibilities.iter()))
+    {
+        match visibility {
+            InputVisibiility::Public => {
+                public_inputs_outputs.push(account);
+            }
+            InputVisibiility::Private(_) => continue,
+        }
+    }
+
+    // Compute commitments for every private output
+    let private_output_commitments: Vec<_> = private_outputs
+        .iter()
+        .map(|account| account.commitment())
+        .collect();
 
     // Output nullifier of consumed input accounts and commitments of new output private accounts
-    env::commit(&(nullifiers, output_commitments));
+    env::commit(&(
+        public_inputs_outputs,
+        nullifiers,
+        private_output_commitments,
+    ));
 }

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

@@ -1,13 +1,16 @@
 use outer_methods::{OUTER_ELF, OUTER_ID};
+use rand::{rngs::OsRng, Rng};
 use risc0_zkvm::{default_prover, ExecutorEnv, Receipt};
-use toy_example_core::{
-    account::{new_random_nonce, Account},
-    input::InputVisibiility,
-};
+use toy_example_core::{account::Account, input::InputVisibiility};
 use transfer_methods::{TRANSFER_ELF, TRANSFER_ID};
 
 const COMMITMENT_TREE_ROOT: [u32; 8] = [0xdd, 0xee, 0xaa, 0xdd, 0xbb, 0xee, 0xee, 0xff];
 
+pub fn new_random_nonce() -> [u32; 8] {
+    let mut rng = OsRng;
+    std::array::from_fn(|_| rng.gen())
+}
+
 fn mint_fresh_account(address: [u32; 8]) -> Account {
     let nonce = new_random_nonce();
     Account::new(address, nonce)
@@ -35,19 +38,16 @@ fn run_private_execution_of_transfer_program() {
     let receiver = mint_fresh_account(receiver_address);
 
     // Prove inner program and get post state of the accounts
-    let (inner_receipt, outputs) = prove_inner(&sender, &receiver, balance_to_move);
+    let (inner_receipt, inputs_outputs) = prove_inner(&sender, &receiver, balance_to_move);
 
     let visibilities = vec![
         InputVisibiility::Private(Some(sender_private_key)),
         InputVisibiility::Private(None),
     ];
 
-    let inputs_outputs = {
-        let mut vec = vec![sender, receiver];
-        vec.extend_from_slice(&outputs);
-        vec
-    };
     let num_inputs: u32 = inputs_outputs.len() as u32 / 2;
+    let output_nonces: Vec<_> = (0..num_inputs).map(|_| new_random_nonce()).collect();
+    println!("output nonces {output_nonces:?}");
 
     // Prove outer program.
     // This computes the nullifier for the input account
@@ -57,6 +57,7 @@ fn run_private_execution_of_transfer_program() {
     env_builder.write(&num_inputs).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();
@@ -69,10 +70,10 @@ fn run_private_execution_of_transfer_program() {
     // Sanity check
     receipt.verify(OUTER_ID).unwrap();
 
-    let output: [[u32; 8]; 3] = receipt.journal.decode().unwrap();
-    println!("nullifier: {:?}", output[0]);
-    println!("commitment_1: {:?}", output[1]);
-    println!("commitment_2: {:?}", output[2]);
+    let output: (Vec<Account>, Vec<[u32; 8]>, Vec<[u32; 8]>) = receipt.journal.decode().unwrap();
+    println!("public_outputs: {:?}", output.0);
+    println!("nullifiers: {:?}", output.1);
+    println!("commitments: {:?}", output.2);
 }
 
 fn prove_inner(
@@ -91,22 +92,22 @@ fn prove_inner(
 
     let receipt = prove_info.receipt;
 
-    let output: [Account; 4] = receipt.journal.decode().unwrap();
-    let [_, _, sender_post, receiver_post] = output;
+    let inputs_outputs: Vec<Account> = receipt.journal.decode().unwrap();
+    assert_eq!(inputs_outputs.len(), 4);
 
     println!(
         "sender_before: {:?}, sender_after: {:?}",
-        sender, sender_post
+        inputs_outputs[0], inputs_outputs[2]
     );
     println!(
         "receiver_before: {:?}, receiver_after: {:?}",
-        receiver, receiver_post
+        inputs_outputs[1], inputs_outputs[3]
     );
 
     // Sanity check
     receipt.verify(TRANSFER_ID).unwrap();
 
-    (receipt, vec![sender_post, receiver_post])
+    (receipt, inputs_outputs)
 }
 
 #[cfg(test)]

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

@@ -29,11 +29,21 @@ pub fn run_public_execution_of_transfer_program() {
     let env = env_builder.build().unwrap();
 
     let executor = default_executor();
-    let result: [Account; 4] = executor.execute(env, TRANSFER_ELF).unwrap().journal.decode().unwrap();
-    let [_, _, sender_post, receiver_post] = result;
+    let inputs_outputs: Vec<Account> = executor
+        .execute(env, TRANSFER_ELF)
+        .unwrap()
+        .journal
+        .decode()
+        .unwrap();
 
-    println!("sender_before: {:?}, sender_after: {:?}", sender, sender_post);
-    println!("receiver_before: {:?}, receiver_after: {:?}", receiver, receiver_post);
+    println!(
+        "sender_before: {:?}, sender_after: {:?}",
+        inputs_outputs[0], inputs_outputs[2]
+    );
+    println!(
+        "receiver_before: {:?}, receiver_after: {:?}",
+        inputs_outputs[1], inputs_outputs[3],
+    );
 }
 
 #[cfg(test)]

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

@@ -18,5 +18,5 @@ fn main() {
     sender_post.balance -= balance_to_move;
     receiver_post.balance += balance_to_move;
 
-    env::commit(&(sender, receiver, sender_post, receiver_post));
+    env::commit(&vec![sender, receiver, sender_post, receiver_post]);
 }