add variable number of inputs/outputs with visibility

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

@@ -1,47 +1,59 @@
-#![cfg_attr(not(test), no_std)]
-
-use serde::{Serialize, Deserialize};
-use risc0_zkvm::{sha::{Impl, Sha256}, serde::to_vec};
+use risc0_zkvm::{
+    serde::to_vec,
+    sha::{Impl, Sha256},
+};
+use serde::{Deserialize, Serialize};
 
 #[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
 pub struct Account {
     pub address: [u32; 8],
     pub balance: u128,
-    pub nonce: [u32; 8]
+    pub nonce: [u32; 8],
 }
 
 impl Account {
     /// Creates a new account with address = hash(private_key) and balance = 0
     pub fn new_from_private_key(private_key: [u32; 8], nonce: [u32; 8]) -> Self {
         let address = hash(&private_key);
-        Self { address, balance: 0, nonce }
+        Self {
+            address,
+            balance: 0,
+            nonce,
+        }
     }
 
     pub fn new(address: [u32; 8], nonce: [u32; 8]) -> Self {
-        Self { address, balance: 0, nonce }
+        Self {
+            address,
+            balance: 0,
+            nonce,
+        }
     }
 
     /// Returns Hash(Account)
     pub fn commitment(&self) -> [u32; 8] {
         hash(&to_vec(&self).unwrap())
     }
-
 }
 
 pub fn hash(bytes: &[u32]) -> [u32; 8] {
-    Impl::hash_words(&bytes).as_words().try_into().unwrap()
+    Impl::hash_words(bytes).as_words().try_into().unwrap()
 }
 
+/// Dummy implementation
 pub fn is_in_commitment_tree(_commitment: [u32; 8], _tree_root: [u32; 8]) -> bool {
-    // Dummy implementation
     true
 }
 
 /// Returns Hash(Commitment || private_key)
-pub fn compute_nullifier(commitment: [u32; 8], private_key: [u32; 8]) -> [u32; 8] {
+pub fn compute_nullifier(commitment: &[u32; 8], private_key: &[u32; 8]) -> [u32; 8] {
     let mut bytes_to_hash = [0; 16];
-    bytes_to_hash[..8].copy_from_slice(&commitment);
-    bytes_to_hash[8..].copy_from_slice(&private_key);
+    bytes_to_hash[..8].copy_from_slice(commitment);
+    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/core/src/input.rs

@@ -0,0 +1,12 @@
+use serde::{Deserialize, Serialize};
+
+pub type PrivateKey = [u32; 8];
+
+#[derive(Serialize, Deserialize)]
+pub enum InputVisibiility {
+    // A public account
+    Public,
+    // A private account
+    Private(Option<PrivateKey>),
+}
+

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

@@ -1 +1,3 @@
 pub mod account;
+pub mod input;
+

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

@@ -1,57 +1,88 @@
-use risc0_zkvm::{guest::env, sha::{Impl, Sha256}, serde::to_vec};
-use toy_example_core::account::{Account, hash, compute_nullifier, is_in_commitment_tree};
+use risc0_zkvm::{
+    guest::env,
+    serde::to_vec,
+    sha::{Impl, Sha256},
+};
+use toy_example_core::{
+    account::{compute_nullifier, hash, is_in_commitment_tree, Account},
+    input::InputVisibiility,
+};
 
 /// Private execution logic.
 /// Circuit for proving correct execution of some program with program id
 /// equal to `program_id` (last input).
-/// 
+///
 /// Currently only supports private execution of a program with two input accounts, one
 /// of which must be a fresh new account (`account_2`) (for example a private transfer function).
-/// 
+///
 /// This circuit checks:
 /// - That accounts pre states and post states are consistent with the execution of the given `program_id`.
 /// - That `account_2` is fresh (meaning, for this toy example, that it has 0 balance).
 /// - That `program_id` execution didn't change addresses of the accounts.
-/// 
+///
 /// Outputs:
 /// - The nullifier for the only existing input account (account_1)
 /// - The commitments for the private accounts post states.
 fn main() {
-    // Read inputs
-    let account_1_private_key: [u32; 8] = env::read();
-    let account_1: Account = env::read();
-    let account_2: Account = env::read();
-    let account_1_post: Account = env::read();
-    let account_2_post: Account = env::read();
+    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);
+    }
+
+    // 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 commitment_tree_root: [u32; 8] = env::read();
     let program_id: [u32; 8] = env::read();
 
-    // Assert account_2 is a fresh account
-    assert_eq!(account_2.balance, 0);
+    let inputs = inputs_outputs.iter().take(num_inputs as usize);
+    let mut nullifiers = Vec::new();
+    for (visibility, input_account) in input_visibilities.iter().zip(inputs) {
+        match visibility {
+            InputVisibiility::Private(Some(private_key)) => {
+                // Prove ownership of input accounts by proving
+                // knowledge of the pre-image of their addresses.
+                assert_eq!(hash(private_key), input_account.address);
+                // Check the input account was created by a previous transaction
+                // by checking it belongs to the commitments tree.
+                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);
+                nullifiers.push(nullifier);
+            }
+            InputVisibiility::Private(None) => {
+                // Private accounts without a companion private key are
+                // enforced to have default values
+                assert_eq!(input_account.balance, 0);
+            }
+            // No checks on public accounts
+            InputVisibiility::Public => continue,
+        }
+    }
+    let outputs = inputs_outputs.iter().skip(num_inputs as usize);
+    let output_commitments: Vec<_> = outputs.map(|account| account.commitment()).collect();
 
-    // Prove ownership of account_1 account by proving
-    // knowledge of the pre-image of its address
-    assert_eq!(hash(&account_1_private_key), account_1.address);
-
-    // Compute account_1 account commitment and prove it belongs to commitments tree
-    let account_1_commitment = account_1.commitment();
-    assert!(is_in_commitment_tree(account_1_commitment, commitment_tree_root)); // <- Dummy implementation
-
-    // Compute nullifier of account_1 account
-    let account_1_nullifier = compute_nullifier(account_1_commitment, account_1_private_key);
-
-    // Compute accounts post states commitments
-    let account_1_post_commitment = account_1_post.commitment();
-    let account_2_post_commitment = account_2_post.commitment();
+    // 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))
+    {
+        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(&(account_1.clone(), account_2.clone(), account_1_post.clone(), account_2_post.clone())).unwrap()).unwrap();
+    env::verify(program_id, &to_vec(&inputs_outputs).unwrap()).unwrap();
 
-    // Assert `program_id` program didn't modify address fields
-    assert_eq!(account_1.address, account_1_post.address);
-    assert_eq!(account_2.address, account_2_post.address);
-
-    // Output nullifier and commitments of new private accounts
-    env::commit(&(account_1_nullifier, account_1_post_commitment, account_2_post_commitment));
+    // Output nullifier of consumed input accounts and commitments of new output private accounts
+    env::commit(&(nullifiers, output_commitments));
 }

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

@@ -1,17 +1,22 @@
-use transfer_methods::{
-    TRANSFER_ELF, TRANSFER_ID
-};
-use outer_methods::{
-    OUTER_ELF, OUTER_ID
-};
+use outer_methods::{OUTER_ELF, OUTER_ID};
 use risc0_zkvm::{default_prover, ExecutorEnv, Receipt};
-use toy_example_core::account::Account;
+use toy_example_core::{
+    account::{new_random_nonce, Account},
+    input::InputVisibiility,
+};
+use transfer_methods::{TRANSFER_ELF, TRANSFER_ID};
+
+const COMMITMENT_TREE_ROOT: [u32; 8] = [0xdd, 0xee, 0xaa, 0xdd, 0xbb, 0xee, 0xee, 0xff];
+
+fn mint_fresh_account(address: [u32; 8]) -> Account {
+    let nonce = new_random_nonce();
+    Account::new(address, nonce)
+}
 
 /// A private execution of the transfer function.
 /// This actually "burns" a sender private account and "mints" two new private accounts:
 /// one for the recipient with the transferred balance, and another owned by the sender with the remaining balance.
 fn run_private_execution_of_transfer_program() {
-    let commitment_tree_root = [0xdd, 0xee, 0xaa, 0xdd, 0xbb, 0xee, 0xee, 0xff];
     // This is supposed to be an existing private account (UTXO) with balance equal to 150.
     // And it is supposed to be a private account of the user running this private execution (hence the access to the private key)
     let sender_private_key = [0; 8];
@@ -21,47 +26,60 @@ fn run_private_execution_of_transfer_program() {
         account.balance = 150;
         account
     };
+
     let balance_to_move: u128 = 3;
 
     // This is the new private account (UTXO) being minted by this private execution.
     // (The `receiver_address` would be <Npk> in UTXO's terminology)
-    let receiver_address = [99; 8]; 
-    let receiver = Account::new(receiver_address, [1; 8]);
+    let receiver_address = [99; 8];
+    let receiver = mint_fresh_account(receiver_address);
 
     // Prove inner program and get post state of the accounts
-    let (inner_receipt, sender_post, receiver_post) = prove_inner(sender.clone(), receiver.clone(), balance_to_move);
+    let (inner_receipt, 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;
 
     // Prove outer program.
     // This computes the nullifier for the input account
     // and commitments for the accounts post states.
     let mut env_builder = ExecutorEnv::builder();
     env_builder.add_assumption(inner_receipt);
-    env_builder.write(&sender_private_key).unwrap();
-    env_builder.write(&sender).unwrap();
-    env_builder.write(&receiver) .unwrap();
-    env_builder.write(&sender_post).unwrap();
-    env_builder.write(&receiver_post).unwrap();
-    env_builder.write(&commitment_tree_root).unwrap();
+    env_builder.write(&num_inputs).unwrap();
+    env_builder.write(&inputs_outputs).unwrap();
+    env_builder.write(&visibilities).unwrap();
+    env_builder.write(&COMMITMENT_TREE_ROOT).unwrap();
     env_builder.write(&TRANSFER_ID).unwrap();
     let env = env_builder.build().unwrap();
 
     let prover = default_prover();
-    let prove_info = prover
-        .prove(env, OUTER_ELF)
-        .unwrap();
+    let prove_info = prover.prove(env, OUTER_ELF).unwrap();
 
     let receipt = prove_info.receipt;
-    
+
     // 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]);
 }
 
-fn prove_inner(sender: Account, receiver: Account, balance_to_move: u128) -> (Receipt, Account, Account) {
+fn prove_inner(
+    sender: &Account,
+    receiver: &Account,
+    balance_to_move: u128,
+) -> (Receipt, Vec<Account>) {
     let mut env_builder = ExecutorEnv::builder();
     env_builder.write(&sender).unwrap();
     env_builder.write(&receiver).unwrap();
@@ -69,24 +87,26 @@ fn prove_inner(sender: Account, receiver: Account, balance_to_move: u128) -> (Re
     let env = env_builder.build().unwrap();
 
     let prover = default_prover();
-    let prove_info = prover
-        .prove(env, TRANSFER_ELF)
-        .unwrap();
+    let prove_info = prover.prove(env, TRANSFER_ELF).unwrap();
 
     let receipt = prove_info.receipt;
 
     let output: [Account; 4] = receipt.journal.decode().unwrap();
     let [_, _, sender_post, receiver_post] = output;
 
-    println!("sender_before: {:?}, sender_after: {:?}", sender, sender_post);
-    println!("receiver_before: {:?}, receiver_after: {:?}", receiver, receiver_post);
+    println!(
+        "sender_before: {:?}, sender_after: {:?}",
+        sender, sender_post
+    );
+    println!(
+        "receiver_before: {:?}, receiver_after: {:?}",
+        receiver, receiver_post
+    );
 
     // Sanity check
-    receipt
-        .verify(TRANSFER_ID)
-        .unwrap();
+    receipt.verify(TRANSFER_ID).unwrap();
 
-    (receipt, sender_post, receiver_post)
+    (receipt, vec![sender_post, receiver_post])
 }
 
 #[cfg(test)]