refactor

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

@@ -12,7 +12,7 @@ pub struct Account {
 
 impl Account {
     /// Creates a new account with address = hash(private_key) and balance = 0
-    pub fn new(private_key: [u32; 8], nonce: [u32; 8]) -> Self {
+    pub fn new_from_private_key(private_key: [u32; 8], nonce: [u32; 8]) -> Self {
         let address = hash(&private_key);
         Self { address, balance: 0, nonce }
     }

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

@@ -1,6 +1,20 @@
 use risc0_zkvm::{guest::env, sha::{Impl, Sha256}, serde::to_vec};
 use toy_example_core::{Account, hash, compute_nullifier, is_in_commitment_tree};
 
+/// 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 inputs, one
+/// of which must be a fresh new account (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 `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();
@@ -21,7 +35,7 @@ fn main() {
 
     // 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));
+    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);

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

@@ -1,122 +1,2 @@
-use std::u128;
-
-// These constants represent the RISC-V ELF and the image ID generated by risc0-build.
-// The ELF is used for proving and the ID is used for verification.
-use transfer_methods::{
-    TRANSFER_ELF, TRANSFER_ID
-};
-use outer_methods::{
-    OUTER_ELF, OUTER_ID
-};
-use risc0_zkvm::{default_executor, default_prover, ExecutorEnv, Receipt};
-use toy_example_core::Account;
-
-fn prove_inner(sender: Account, receiver: Account, balance_to_move: u128) -> (Receipt, Account, Account) {
-    let mut env_builder = ExecutorEnv::builder();
-    env_builder.write(&sender).unwrap();
-    env_builder.write(&receiver).unwrap();
-    env_builder.write(&balance_to_move).unwrap();
-    let env = env_builder.build().unwrap();
-
-    let prover = default_prover();
-    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);
-
-    // Sanity check
-    receipt
-        .verify(TRANSFER_ID)
-        .unwrap();
-
-    (receipt, sender_post, receiver_post)
-}
-
-pub fn run_private_execution_of_transfer_program() {
-    let commitment_tree_root = [0xdd, 0xee, 0xaa, 0xdd, 0xbb, 0xee, 0xee, 0xff];
-    let sender_private_key = [0; 8];
-    let mut sender = Account::new(sender_private_key, [1; 8]);
-    sender.balance = 150;
-
-    let receiver_private_key = [99; 8];
-    let receiver = Account::new(receiver_private_key, [1; 8]);
-    let balance_to_move: u128 = 3;
-
-    // Prove inner
-    let (inner_receipt, sender_post, receiver_post) = prove_inner(sender.clone(), receiver.clone(), balance_to_move);
-
-    // Prover outer
-    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(&balance_to_move).unwrap();
-    env_builder.write(&sender_post).unwrap();
-    env_builder.write(&receiver_post).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 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]);
-}
-
-pub fn run_public_execution_of_transfer_program() {
-    let sender_private_key = [0; 8];
-    let mut sender = Account::new(sender_private_key, [1; 8]);
-    sender.balance = 150;
-
-    let receiver_private_key = [99; 8];
-    let mut receiver = Account::new(receiver_private_key, [1; 8]);
-    receiver.balance = 900;
-
-    let balance_to_move: u128 = 3;
-
-    let mut env_builder = ExecutorEnv::builder();
-    env_builder.write(&sender).unwrap();
-    env_builder.write(&receiver).unwrap();
-    env_builder.write(&balance_to_move).unwrap();
-    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;
-
-    println!("sender_before: {:?}, sender_after: {:?}", sender, sender_post);
-    println!("receiver_before: {:?}, receiver_after: {:?}", receiver, receiver_post);
-}
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    #[test]
-    fn test_private() {
-        run_private_execution_of_transfer_program();
-    }
-
-    #[test]
-    fn test_public() {
-        run_public_execution_of_transfer_program();
-    }
-}
+mod private_execution;
+mod public_execution;

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

@@ -0,0 +1,88 @@
+use transfer_methods::{
+    TRANSFER_ELF, TRANSFER_ID
+};
+use outer_methods::{
+    OUTER_ELF, OUTER_ID
+};
+use risc0_zkvm::{default_prover, ExecutorEnv, Receipt};
+use toy_example_core::Account;
+
+fn prove_inner(sender: Account, receiver: Account, balance_to_move: u128) -> (Receipt, Account, Account) {
+    let mut env_builder = ExecutorEnv::builder();
+    env_builder.write(&sender).unwrap();
+    env_builder.write(&receiver).unwrap();
+    env_builder.write(&balance_to_move).unwrap();
+    let env = env_builder.build().unwrap();
+
+    let prover = default_prover();
+    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);
+
+    // Sanity check
+    receipt
+        .verify(TRANSFER_ID)
+        .unwrap();
+
+    (receipt, sender_post, receiver_post)
+}
+
+fn run_private_execution_of_transfer_program() {
+    let commitment_tree_root = [0xdd, 0xee, 0xaa, 0xdd, 0xbb, 0xee, 0xee, 0xff];
+    let sender_private_key = [0; 8];
+    let mut sender = Account::new_from_private_key(sender_private_key, [1; 8]);
+    sender.balance = 150;
+
+    let receiver_private_key = [99; 8];
+    let receiver = Account::new_from_private_key(receiver_private_key, [1; 8]);
+    let balance_to_move: u128 = 3;
+
+    // Prove inner
+    let (inner_receipt, sender_post, receiver_post) = prove_inner(sender.clone(), receiver.clone(), balance_to_move);
+
+    // Prover outer
+    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(&balance_to_move).unwrap();
+    env_builder.write(&sender_post).unwrap();
+    env_builder.write(&receiver_post).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 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]);
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_private() {
+        run_private_execution_of_transfer_program();
+    }
+}

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

@@ -0,0 +1,38 @@
+use risc0_zkvm::{default_executor, ExecutorEnv};
+use toy_example_core::Account;
+use transfer_methods::TRANSFER_ELF;
+
+pub fn run_public_execution_of_transfer_program() {
+    let sender_private_key = [0; 8];
+    let mut sender = Account::new_from_private_key(sender_private_key, [1; 8]);
+    sender.balance = 150;
+
+    let receiver_private_key = [99; 8];
+    let mut receiver = Account::new_from_private_key(receiver_private_key, [1; 8]);
+    receiver.balance = 900;
+
+    let balance_to_move: u128 = 3;
+
+    let mut env_builder = ExecutorEnv::builder();
+    env_builder.write(&sender).unwrap();
+    env_builder.write(&receiver).unwrap();
+    env_builder.write(&balance_to_move).unwrap();
+    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;
+
+    println!("sender_before: {:?}, sender_after: {:?}", sender, sender_post);
+    println!("receiver_before: {:?}, receiver_after: {:?}", receiver, receiver_post);
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_public() {
+        run_public_execution_of_transfer_program();
+    }
+}

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

@@ -1,16 +1,20 @@
 use risc0_zkvm::{guest::env, sha::{Impl, Sha256}, serde::to_vec};
 use toy_example_core::Account;
 
+
+/// A transfer of balance program.
+/// To be used both in public and private contexts
 fn main() {
     let sender: Account = env::read();
     let receiver: Account = env::read();
     let balance_to_move: u128 = env::read();
 
+    // Check sender has enough balance
     assert!(sender.balance >= balance_to_move);
 
+    // Create accounts post states, with updated balances
     let mut sender_post = sender.clone();
     let mut receiver_post = receiver.clone();
-
     sender_post.balance -= balance_to_move;
     receiver_post.balance += balance_to_move;