add comments

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

@@ -1,3 +1,4 @@
+/// For this POC we consider 32-bit commitments
 pub type Commitment = u32;
 pub type Nullifier = [u32; 8];
 pub type Address = [u32; 8];

risc0-selective-privacy-poc/examples/mocked_components/client/mod.rs

@@ -1,18 +1,17 @@
+use crate::mocked_components::sequencer::MockedSequencer;
 use core::{
     account::Account,
     types::{Address, Commitment, Key, Nullifier},
     visibility::InputVisibiility,
 };
-
 use nssa::program::TransferProgram;
 
-use crate::mocked_components::sequencer::MockedSequencer;
-
 pub mod transfer_deshielded;
 pub mod transfer_private;
 pub mod transfer_public;
 pub mod transfer_shielded;
 
+/// A client that creates and submits transfer transactions
 pub struct MockedClient {
     user_private_key: Key,
     private_accounts: Vec<Account>,
@@ -30,6 +29,8 @@ impl MockedClient {
         Account::address_for_key(&self.user_private_key)
     }
 
+    /// Runs the outer program and submits the proof to the sequencer.
+    /// Returns the output private accounts of the execution.
     pub fn prove_and_send_to_sequencer<P: nssa::Program>(
         input_accounts: &[Account],
         instruction_data: P::InstructionData,
@@ -37,19 +38,19 @@ impl MockedClient {
         commitment_tree_root: [u32; 8],
         sequencer: &mut MockedSequencer,
     ) -> Result<Vec<Account>, ()> {
-        let (receipt, private_outputs) = nssa::invoke_privacy_execution::<P>(
-            input_accounts,
-            instruction_data,
-            visibilities,
-            commitment_tree_root,
-        )
-        .unwrap();
-        // Send to te sequencer
+        // Execute and generate proof of the outer program
+        let (receipt, private_outputs) =
+            nssa::invoke_privacy_execution::<P>(input_accounts, instruction_data, visibilities, commitment_tree_root)
+                .unwrap();
+
+        // Send proof to the sequencer
         sequencer.process_privacy_execution(receipt)?;
 
+        // Return private outputs
         Ok(private_outputs)
     }
 
+    /// Returns a new account used in privacy executions that mint private accounts
     pub fn fresh_account_for_mint(address: Address) -> Account {
         Account::new(address, 0)
     }

risc0-selective-privacy-poc/examples/mocked_components/client/transfer_deshielded.rs

@@ -1,12 +1,14 @@
 use core::account::Account;
-use core::visibility::InputVisibiility;
 use core::types::{Address, Commitment, Key, Nullifier};
+use core::visibility::InputVisibiility;
 
 use nssa::program::TransferProgram;
 
 use super::{MockedClient, MockedSequencer};
 
 impl MockedClient {
+    /// A deshielded transaction of the Transfer program.
+    /// All of this is executed locally by the sender
     pub fn transfer_deshielded(
         &self,
         from_account: Account,
@@ -14,17 +16,22 @@ impl MockedClient {
         balance_to_move: u128,
         sequencer: &mut MockedSequencer,
     ) -> Result<Account, ()> {
-        // All of this is executed locally by the sender
+        // Fetch commitment tree root from the sequencer
         let commitment_tree_root = sequencer.get_commitment_tree_root();
-        let receiver_addr = to_address;
-        // let from_account = sequencer.get_account(&self.user_address()).ok_or(())?;
-        let sender_commitment_auth_path =
-            sequencer.get_authentication_path_for(&from_account.commitment());
+        // Compute authenticaton path for the input private account
+        let sender_commitment_auth_path = sequencer.get_authentication_path_for(&from_account.commitment());
+
+        // Fetch public account to deshield to
         let to_account = sequencer.get_account(&to_address).unwrap();
+
+        // Set input visibilities
+        // First entry is the private sender. Second entry is the public receiver
         let visibilities = vec![
             InputVisibiility::Private(Some((self.user_private_key, sender_commitment_auth_path))),
             InputVisibiility::Public,
         ];
+
+        // Execute privately (off-chain) and submit it to the sequencer
         let private_outputs = Self::prove_and_send_to_sequencer::<TransferProgram>(
             &[from_account, to_account],
             balance_to_move,
@@ -32,6 +39,9 @@ impl MockedClient {
             commitment_tree_root,
             sequencer,
         )?;
+
+        // There's only one private output account corresponding to the new private account of
+        // the sender, with the remaining balance.
         let [sender_private_account] = private_outputs.try_into().unwrap();
         Ok(sender_private_account)
     }

risc0-selective-privacy-poc/examples/mocked_components/client/transfer_private.rs

@@ -1,6 +1,6 @@
 use core::account::Account;
-use core::visibility::InputVisibiility;
 use core::types::{Address, Commitment, Key, Nullifier};
+use core::visibility::InputVisibiility;
 
 use nssa::program::TransferProgram;
 
@@ -8,6 +8,7 @@ use super::{MockedClient, MockedSequencer};
 
 impl MockedClient {
     /// A private execution of the Transfer program
+    // All of this is executed locally by the sender
     pub fn transfer_private(
         &self,
         owned_private_account: Account,
@@ -15,17 +16,21 @@ impl MockedClient {
         balance_to_move: u128,
         sequencer: &mut MockedSequencer,
     ) -> Result<[Account; 2], ()> {
-        // All of this is executed locally by the sender
+        // Fetch commitment tree root from the sequencer
         let commitment_tree_root = sequencer.get_commitment_tree_root();
-        let receiver_addr = to_address;
-        let sender_commitment_auth_path =
-            sequencer.get_authentication_path_for(&owned_private_account.commitment());
-        let mut receiver_account = Self::fresh_account_for_mint(*receiver_addr);
+        // Compute authenticaton path for the input private account
+        let sender_commitment_auth_path = sequencer.get_authentication_path_for(&owned_private_account.commitment());
+
+        // Create a new default private account for the recipient
+        let mut receiver_account = Self::fresh_account_for_mint(*to_address);
+
+        // Set visibilities. Both private accounts.
         let visibilities = vec![
             InputVisibiility::Private(Some((self.user_private_key, sender_commitment_auth_path))),
             InputVisibiility::Private(None),
         ];
 
+        // Execute privately (off-chain) and submit it to the sequencer
         let private_outputs = Self::prove_and_send_to_sequencer::<TransferProgram>(
             &[owned_private_account, receiver_account],
             balance_to_move,
@@ -34,6 +39,10 @@ impl MockedClient {
             sequencer,
         )?;
 
-        Ok(private_outputs.try_into().unwrap())
+        // There are two output private accounts of this execution.
+        // The first corresponds to the sender, with the remaining balance.
+        // The second corresponds to the newly minted private account for the recipient.
+        let [sender_private_account, receiver_private_account] = private_outputs.try_into().unwrap();
+        Ok([sender_private_account, receiver_private_account])
     }
 }

risc0-selective-privacy-poc/examples/mocked_components/client/transfer_public.rs

@@ -7,13 +7,11 @@ use crate::mocked_components::{client::MockedClient, sequencer::MockedSequencer}
 impl MockedClient {
     pub fn transfer_public(
         &self,
-        receiver_address: &Address,
+        to_address: &Address,
         amount_to_transfer: u128,
         sequencer: &mut MockedSequencer,
     ) -> Result<(), ()> {
-        sequencer.process_public_execution::<TransferProgram>(
-            &[self.user_address(), *receiver_address],
-            amount_to_transfer,
-        )
+        // Submit a public (on-chain) execution of the Transfer program to the sequencer
+        sequencer.process_public_execution::<TransferProgram>(&[self.user_address(), *to_address], amount_to_transfer)
     }
 }

risc0-selective-privacy-poc/examples/mocked_components/client/transfer_shielded.rs

@@ -8,26 +8,38 @@ use super::{MockedClient, MockedSequencer};
 
 impl MockedClient {
     /// A shielded execution of the Transfer program
+    // All of this is executed locally by the sender
     pub fn transfer_shielded(
         &self,
         to_address: &Address,
         balance_to_move: u128,
         sequencer: &mut MockedSequencer,
     ) -> Result<Account, ()> {
-        // All of this is executed locally by the sender
-        let sender_account = sequencer.get_account(&self.user_address()).ok_or(())?;
+        // Fetch commitment tree root from the sequencer
         let commitment_tree_root = sequencer.get_commitment_tree_root();
-        let receiver_addr = to_address;
-        let mut receiver_account = Self::fresh_account_for_mint(*receiver_addr);
+
+        // Fetch sender account from the sequencer
+        let from_account = sequencer.get_account(&self.user_address()).ok_or(())?;
+
+        // Create a new default private account for the receiver
+        let mut to_account = Self::fresh_account_for_mint(*to_address);
+
+        // Set input visibilities
+        // First is the public account of the sender. Second is the private account minted in this
+        // execution
         let visibilities = [InputVisibiility::Public, InputVisibiility::Private(None)];
 
+        // Execute privately (off-chain) and submit it to the sequencer
         let private_outputs = Self::prove_and_send_to_sequencer::<TransferProgram>(
-            &[sender_account, receiver_account],
+            &[from_account, to_account],
             balance_to_move,
             &visibilities,
             commitment_tree_root,
             sequencer,
         )?;
+
+        // There is only one private account as the output of this execution. It corresponds to the
+        // receiver.
         let [receiver_private_account] = private_outputs.try_into().unwrap();
         Ok(receiver_private_account)
     }

risc0-selective-privacy-poc/examples/mocked_components/sequencer/mod.rs

@@ -20,8 +20,11 @@ pub struct MockedSequencer {
     deployed_program_ids: HashSet<ProgramId>,
 }
 
+/// List of deployed programs
 const DEPLOYED_PROGRAM_IDS: [ProgramId; 3] = [TRANSFER_ID, TRANSFER_MULTIPLE_ID, PINATA_ID];
+/// The initial balance of the genesis accounts
 const INITIAL_BALANCE: u128 = 150;
+/// The address of the piñata program account
 const PINATA_ADDRESS: Address = [0xcafe; 8];
 
 impl MockedSequencer {
@@ -46,14 +49,17 @@ impl MockedSequencer {
         }
     }
 
+    /// Returns the current state of the account for the given address
     pub fn get_account(&self, address: &Address) -> Option<Account> {
         self.accounts.get(address).cloned()
     }
 
+    /// Returns the root of the commitment tree
     pub fn get_commitment_tree_root(&self) -> [u32; 8] {
         bytes_to_words(&self.commitment_tree.root())
     }
 
+    /// Computes the authentication path for the given commitment
     pub fn get_authentication_path_for(&self, commitment: &Commitment) -> AuthenticationPath {
         self.commitment_tree
             .get_authentication_path_for_value(*commitment)
@@ -64,11 +70,13 @@ impl MockedSequencer {
             .unwrap()
     }
 
+    /// Returns the list of all registered addresses
     pub fn addresses(&self) -> Vec<Address> {
         self.accounts.keys().cloned().collect()
     }
 }
 
+/// Pretty prints the chain's state
 pub fn print_accounts(sequencer: &MockedSequencer, private_accounts: &[&Account]) {
     println!("\n====================== ACCOUNT SNAPSHOT ======================\n");
 
@@ -107,10 +115,7 @@ pub fn print_accounts(sequencer: &MockedSequencer, private_accounts: &[&Account]
     println!("{:-<20}\n", "");
 
     println!(">> Private Accounts:");
-    println!(
-        "{:<20} | {:>10} | {:>10}",
-        "Address (first u32)", "Nonce", "Balance"
-    );
+    println!("{:<20} | {:>10} | {:>10}", "Address (first u32)", "Nonce", "Balance");
     println!("{:-<20}-+-{:-<10}-+-{:-<10}", "", "", "");
 
     for account in private_accounts.iter() {

risc0-selective-privacy-poc/examples/mocked_components/sequencer/process_privacy_execution.rs

@@ -8,19 +8,28 @@ use risc0_zkvm::Receipt;
 use super::MockedSequencer;
 
 impl MockedSequencer {
+    /// Processes a privacy execution request.
+    /// Verifies the proof of the privacy execution and updates the state of the chain.
     pub fn process_privacy_execution(&mut self, receipt: Receipt) -> Result<(), ()> {
-        let output: (Vec<Account>, Vec<Nullifier>, Vec<Commitment>, [u32; 8]) =
-            receipt.journal.decode().unwrap();
+        // Parse the output of the proof.
+        // This is the output of the "outer" program
+        let output: (Vec<Account>, Vec<Nullifier>, Vec<Commitment>, [u32; 8]) = receipt.journal.decode().unwrap();
         let (public_inputs_outputs, nullifiers, commitments, commitment_tree_root) = output;
 
+        // Reject in case the root used in the privacy execution is not the current root.
         if commitment_tree_root != self.get_commitment_tree_root() {
             return Err(());
         }
 
+        // Reject in case the number of accounts in the public_inputs_outputs is not even.
+        // This is because it is expected to contain the pre and post-states of public the accounts
+        // of the inner execution.
         if public_inputs_outputs.len() % 2 != 0 {
             return Err(());
         }
 
+        // Reject if the states of the public input accounts used in the inner execution do not
+        // coincide with the on-chain state.
         let num_input_public = public_inputs_outputs.len() >> 1;
         for account in public_inputs_outputs.iter().take(num_input_public) {
             let current_account = self.get_account(&account.address).ok_or(())?;
@@ -29,7 +38,7 @@ impl MockedSequencer {
             }
         }
 
-        // Check that nullifiers have not been added before
+        // Reject if the nullifiers of this privacy execution have already been published.
         if nullifiers
             .iter()
             .any(|nullifier| self.nullifier_set.contains(nullifier))
@@ -37,7 +46,7 @@ impl MockedSequencer {
             return Err(());
         }
 
-        // Check that commitments are new too
+        // Reject if the commitments have already been seen.
         if commitments
             .iter()
             .any(|commitment| self.commitment_tree.values().contains(commitment))
@@ -45,7 +54,13 @@ impl MockedSequencer {
             return Err(());
         }
 
-        // Verify consistency between public accounts, nullifiers and commitments
+        // Verify the proof of the privacy execution.
+        // This includes a proof of the following statements
+        // - Public inputs, public outputs, commitments and nullifiers are consistent with the
+        //   execution of some program.
+        // - The given nullifiers correctly correspond to commitments that currently belong to
+        //   the commitment tree.
+        // - The given commitments are correctly computed from valid accounts.
         nssa::verify_privacy_execution(
             receipt,
             &public_inputs_outputs,
@@ -54,20 +69,21 @@ impl MockedSequencer {
             &commitment_tree_root,
         )?;
 
-        // Update accounts
+        // At this point the privacy execution is considered valid.
+        //
+        // Update the state of the public accounts with the post-state of this privacy execution
         public_inputs_outputs
             .iter()
             .cloned()
             .skip(num_input_public)
             .for_each(|account_post_state| {
-                self.accounts
-                    .insert(account_post_state.address, account_post_state);
+                self.accounts.insert(account_post_state.address, account_post_state);
             });
 
-        // Add nullifiers
+        // Add all nullifiers to the nullifier set.
         self.nullifier_set.extend(nullifiers);
 
-        // Add commitments
+        // Add commitments to the commitment tree.
         for commitment in commitments.iter() {
             self.commitment_tree.add_value(*commitment);
         }

risc0-selective-privacy-poc/examples/mocked_components/sequencer/process_public_execution.rs

@@ -3,18 +3,19 @@ use core::{account::Account, types::Address};
 use super::MockedSequencer;
 
 impl MockedSequencer {
+    /// Processes a public execution request of the program `P`.
     pub fn process_public_execution<P: nssa::Program>(
         &mut self,
         input_account_addresses: &[Address],
         instruction_data: P::InstructionData,
     ) -> Result<(), ()> {
-        // Fetch accounts
+        // Fetch the current state of the input accounts.
         let input_accounts: Vec<Account> = input_account_addresses
             .iter()
             .map(|address| self.get_account(address).ok_or(()))
             .collect::<Result<_, _>>()?;
 
-        // Execute
+        // Execute the program
         let inputs_outputs = nssa::execute::<P>(&input_accounts, instruction_data)?;
 
         // Perform consistency checks
@@ -22,51 +23,58 @@ impl MockedSequencer {
             return Err(());
         }
 
-        // Update accounts
+        // 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);
+                self.accounts.insert(account_post_state.address, account_post_state);
             });
         Ok(())
     }
 
-    fn inputs_outputs_are_consistent(
-        &self,
-        input_accounts: &[Account],
-        inputs_outputs: &[Account],
-    ) -> bool {
+    /// Verifies that a program public execution didn't break the chain's rules.
+    /// `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 {
         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 {
             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 {
             return false;
         }
 
         for (account_pre, account_post) in input_accounts.iter().zip(accounts_post) {
+            // Fail if the program modified the addresses of the input accounts
             if account_pre.address != account_post.address {
                 return false;
             }
+            // Fail if the program modified the nonces of the input accounts
             if account_pre.nonce != account_post.nonce {
                 return false;
             }
-            // Redundant with previous checks, but better make it explicit.
+            // Fail if any of the output accounts is not yet registered.
+            // (redundant with previous checks, but better make it explicit)
             if !self.accounts.contains_key(&account_post.address) {
                 return false;
             }
         }
-        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 {
+
+        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();
+        // Fail if the execution didn't preserve the total supply.
+        if total_balance_pre != total_balance_post {
             return false;
         }
+
         return true;
     }
 }