refactor

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

@@ -5,6 +5,7 @@ use crate::{
 use risc0_zkvm::{serde::to_vec, sha::Impl};
 use serde::{Deserialize, Serialize};
 
+/// Account to be used both in public and private contexts
 #[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
 pub struct Account {
     pub address: Address,
@@ -13,25 +14,26 @@ pub struct Account {
 }
 
 impl Account {
-    /// Creates a new account with address = hash(private_key) and balance = 0
-    pub fn new_from_private_key(private_key: Key, nonce: Nonce) -> Self {
-        let address = Self::address_for_key(&private_key);
-        Self::new(address, nonce)
+    pub fn new(address: Address, balance: u128) -> Self {
+        Self {
+            address,
+            balance,
+            nonce: [0; 8],
+        }
     }
 
+    /// Creates a new account with address = hash(private_key) and balance = 0
+    pub fn new_from_private_key(private_key: Key) -> Self {
+        let address = Self::address_for_key(&private_key);
+        Self::new(address, 0)
+    }
+
+    /// Computes the address corresponding to the given private key
     pub fn address_for_key(private_key: &Key) -> Address {
         hash(private_key)
     }
 
-    pub fn new(address: Address, nonce: Nonce) -> Self {
-        Self {
-            address,
-            balance: 0,
-            nonce,
-        }
-    }
-
-    /// Returns Hash(Account)[0] (only first word for this POC)
+    /// Returns (first 8 bytes of) SHA256(Account)
     pub fn commitment(&self) -> Commitment {
         hash(&to_vec(&self).unwrap())[0]
     }

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

@@ -1,6 +1,6 @@
 pub mod account;
-pub mod input;
 pub mod types;
+pub mod visibility;
 
 use crate::types::{AuthenticationPath, Commitment, Key, Nullifier};
 use risc0_zkvm::sha::{Impl, Sha256};

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

@@ -1,5 +1,5 @@
-use serde::{Deserialize, Serialize};
 use crate::types::{AuthenticationPath, Key};
+use serde::{Deserialize, Serialize};
 
 #[derive(Serialize, Deserialize)]
 pub enum InputVisibiility {

risc0-selective-privacy-poc/examples/happy_path.rs

@@ -1,14 +1,14 @@
 use core::{
     account::Account,
     bytes_to_words, hash,
-    input::InputVisibiility,
     types::{Address, Commitment, Key, Nullifier},
+    visibility::InputVisibiility,
 };
 
 use nssa::program::{PinataProgram, TransferProgram};
 use risc0_zkvm::Receipt;
 
-use crate::mocked_components::sequencer::MockedSequencer;
+use crate::mocked_components::sequencer::{print_accounts, MockedSequencer};
 use crate::mocked_components::{client::MockedClient, USER_CLIENTS};
 
 mod mocked_components;
@@ -16,44 +16,81 @@ mod mocked_components;
 fn main() {
     let mut sequencer = MockedSequencer::new();
     let addresses: [Address; 3] = USER_CLIENTS.map(|client| client.user_address());
-
-    println!("addresses: {:?}", addresses);
-    println!("🚀 Initial balances");
-    sequencer.print();
+    println!("📝 Initial balances");
+    print_accounts(&sequencer, &[]);
 
     // A public execution of the Transfer Program
-    USER_CLIENTS[0]
-        .transfer_public(&addresses[1], 10, &mut sequencer)
+    USER_CLIENTS[1]
+        .transfer_public(&[0xcafe; 8], 51, &mut sequencer)
         .unwrap();
-    println!("🚀 Balances after transfer");
-    sequencer.print();
+    println!("📝 Balances after transfer");
+    print_accounts(&sequencer, &[]);
 
     // A shielded execution of the Transfer Program
-    let private_account_1 = USER_CLIENTS[0]
+    let private_account_user_1 = USER_CLIENTS[0]
         .transfer_shielded(&addresses[1], 15, &mut sequencer)
         .unwrap();
-    println!("Balances after shielded execution");
-    sequencer.print();
+    println!("📝 Balances after shielded execution");
+    print_accounts(&sequencer, &[&private_account_user_1]);
 
     // A private execution of the Transfer Program
-    let [_, private_account_2] = USER_CLIENTS[1]
-        .transfer_private(private_account_1, &addresses[2], 8, &mut sequencer)
+    let [private_account_user_1, private_account_user_2] = USER_CLIENTS[1]
+        .transfer_private(private_account_user_1, &addresses[2], 8, &mut sequencer)
         .unwrap();
-    println!("🚀 Balances after shielded execution");
-    sequencer.print();
+    println!("📝 Balances after shielded execution");
+    print_accounts(
+        &sequencer,
+        &[&private_account_user_1, &private_account_user_2],
+    );
 
     // A deshielded execution of the Transfer Program
-    USER_CLIENTS[2]
-        .transfer_deshielded(private_account_2, &addresses[0], 1, &mut sequencer)
+    let private_acount_user_2 = USER_CLIENTS[2]
+        .transfer_deshielded(private_account_user_2, &addresses[0], 1, &mut sequencer)
         .unwrap();
-    println!("🚀 Balances after deshielded execution");
-    sequencer.print();
+    println!("📝 Balances after deshielded execution");
+    print_accounts(
+        &sequencer,
+        &[&private_account_user_1, &private_acount_user_2],
+    );
 
-    // A public execution of the Pinata program
+    // A public execution of the Piñata program
     let preimage = bytes_to_words(b"NSSA Selective privacy is great!").to_vec();
     sequencer
         .process_public_execution::<PinataProgram>(&[[0xcafe; 8], addresses[2]], preimage)
         .unwrap();
-    println!("🚀 Balances after public piñata execution");
-    sequencer.print();
+    println!("📝 Balances after public piñata execution");
+    print_accounts(
+        &sequencer,
+        &[&private_account_user_1, &private_acount_user_2],
+    );
+
+    // A deshielded execution of the Piñata program
+    let private_account_user_0 = {
+        // All of this is executed locally by the sender
+        let receiver_addr = USER_CLIENTS[1].user_address();
+        let pinata_account = sequencer.get_account(&[0xcafe; 8]).unwrap();
+        let mut receiver_account = MockedClient::fresh_account_for_mint(receiver_addr);
+        let visibilities = [InputVisibiility::Public, InputVisibiility::Private(None)];
+        let preimage = bytes_to_words(b"NSSA Selective privacy is great!").to_vec();
+
+        let private_outputs = MockedClient::prove_and_send_to_sequencer::<PinataProgram>(
+            &[pinata_account, receiver_account],
+            preimage,
+            &visibilities,
+            sequencer.get_commitment_tree_root(),
+            &mut sequencer,
+        )
+        .unwrap();
+        let [private_account_user_0] = private_outputs.try_into().unwrap();
+        private_account_user_0
+    };
+    println!("📝 Balances after private piñata execution");
+    print_accounts(
+        &sequencer,
+        &[
+            &private_account_user_1,
+            &private_acount_user_2,
+            &private_account_user_0,
+        ],
+    );
 }

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

@@ -1,7 +1,7 @@
 use core::{
     account::Account,
-    input::InputVisibiility,
     types::{Address, Commitment, Key, Nullifier},
+    visibility::InputVisibiility,
 };
 
 use nssa::program::TransferProgram;
@@ -15,19 +15,22 @@ pub mod transfer_shielded;
 
 pub struct MockedClient {
     user_private_key: Key,
+    private_accounts: Vec<Account>,
 }
 
 impl MockedClient {
     pub const fn new(user_private_key: Key) -> Self {
-        Self { user_private_key }
+        Self {
+            user_private_key,
+            private_accounts: Vec::new(),
+        }
     }
 
     pub fn user_address(&self) -> Address {
-        let address = Account::address_for_key(&self.user_private_key);
-        address 
+        Account::address_for_key(&self.user_private_key)
     }
 
-    fn prove_and_send_to_sequencer<P: nssa::Program>(
+    pub fn prove_and_send_to_sequencer<P: nssa::Program>(
         input_accounts: &[Account],
         instruction_data: P::InstructionData,
         visibilities: &[InputVisibiility],
@@ -48,7 +51,6 @@ impl MockedClient {
     }
 
     pub fn fresh_account_for_mint(address: Address) -> Account {
-        let nonce = [0; 8];
-        Account::new(address, nonce)
+        Account::new(address, 0)
     }
 }

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

@@ -1,5 +1,5 @@
 use core::account::Account;
-use core::input::InputVisibiility;
+use core::visibility::InputVisibiility;
 use core::types::{Address, Commitment, Key, Nullifier};
 
 use nssa::program::TransferProgram;
@@ -7,28 +7,32 @@ use nssa::program::TransferProgram;
 use super::{MockedClient, MockedSequencer};
 
 impl MockedClient {
-    /// A shielded execution of the Transfer program
-    pub fn transfer_shielded(
+    pub fn transfer_deshielded(
         &self,
+        from_account: Account,
         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(())?;
         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);
-        let visibilities = [InputVisibiility::Public, InputVisibiility::Private(None)];
-
+        // let from_account = sequencer.get_account(&self.user_address()).ok_or(())?;
+        let sender_commitment_auth_path =
+            sequencer.get_authentication_path_for(&from_account.commitment());
+        let to_account = sequencer.get_account(&to_address).unwrap();
+        let visibilities = vec![
+            InputVisibiility::Private(Some((self.user_private_key, sender_commitment_auth_path))),
+            InputVisibiility::Public,
+        ];
         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,
         )?;
-        let [receiver_private_account] = private_outputs.try_into().unwrap();
-        Ok(receiver_private_account)
+        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,5 +1,5 @@
 use core::account::Account;
-use core::input::InputVisibiility;
+use core::visibility::InputVisibiility;
 use core::types::{Address, Commitment, Key, Nullifier};
 
 use nssa::program::TransferProgram;
@@ -10,7 +10,7 @@ impl MockedClient {
     /// A private execution of the Transfer program
     pub fn transfer_private(
         &self,
-        from_account: Account,
+        owned_private_account: Account,
         to_address: &Address,
         balance_to_move: u128,
         sequencer: &mut MockedSequencer,
@@ -18,9 +18,8 @@ impl MockedClient {
         // All of this is executed locally by the sender
         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());
+            sequencer.get_authentication_path_for(&owned_private_account.commitment());
         let mut receiver_account = Self::fresh_account_for_mint(*receiver_addr);
         let visibilities = vec![
             InputVisibiility::Private(Some((self.user_private_key, sender_commitment_auth_path))),
@@ -28,7 +27,7 @@ impl MockedClient {
         ];
 
         let private_outputs = Self::prove_and_send_to_sequencer::<TransferProgram>(
-            &[from_account, receiver_account],
+            &[owned_private_account, receiver_account],
             balance_to_move,
             &visibilities,
             commitment_tree_root,

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

@@ -1,37 +1,34 @@
 use core::account::Account;
-use core::input::InputVisibiility;
 use core::types::{Address, Commitment, Key, Nullifier};
+use core::visibility::InputVisibiility;
 
 use nssa::program::TransferProgram;
 
 use super::{MockedClient, MockedSequencer};
 
 impl MockedClient {
-    pub fn transfer_deshielded(
+    /// A shielded execution of the Transfer program
+    pub fn transfer_shielded(
         &self,
-        from_account: Account,
         to_address: &Address,
         balance_to_move: u128,
         sequencer: &mut MockedSequencer,
-    ) -> Result<(), ()> {
+    ) -> Result<Account, ()> {
         // All of this is executed locally by the sender
+        let sender_account = sequencer.get_account(&self.user_address()).ok_or(())?;
         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());
-        let to_account = sequencer.get_account(&to_address).unwrap();
-        let visibilities = vec![
-            InputVisibiility::Private(Some((self.user_private_key, sender_commitment_auth_path))),
-            InputVisibiility::Public,
-        ];
+        let mut receiver_account = Self::fresh_account_for_mint(*receiver_addr);
+        let visibilities = [InputVisibiility::Public, InputVisibiility::Private(None)];
+
         let private_outputs = Self::prove_and_send_to_sequencer::<TransferProgram>(
-            &[from_account, to_account],
+            &[sender_account, receiver_account],
             balance_to_move,
             &visibilities,
             commitment_tree_root,
             sequencer,
         )?;
-        Ok(())
+        let [receiver_private_account] = private_outputs.try_into().unwrap();
+        Ok(receiver_private_account)
     }
 }

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

@@ -20,28 +20,20 @@ pub struct MockedSequencer {
     deployed_program_ids: HashSet<ProgramId>,
 }
 
-const ACCOUNTS_INITIAL_BALANCES: [u128; 3] = [100, 1337, 37];
 const DEPLOYED_PROGRAM_IDS: [ProgramId; 3] = [TRANSFER_ID, TRANSFER_MULTIPLE_ID, PINATA_ID];
+const INITIAL_BALANCE: u128 = 150;
+const PINATA_ADDRESS: Address = [0xcafe; 8];
 
 impl MockedSequencer {
     pub fn new() -> Self {
         let mut accounts: BTreeMap<Address, Account> = USER_CLIENTS
             .iter()
             .map(|client| client.user_address())
-            .zip(ACCOUNTS_INITIAL_BALANCES)
-            .map(|(address, initial_balance)| {
-                let mut this = Account::new(address, [0; 8]);
-                this.balance = initial_balance;
-                this
-            })
+            .map(|address| Account::new(address, INITIAL_BALANCE))
             .map(|account| (account.address, account))
             .collect();
 
-        let pinata_account = {
-            let mut this = Account::new([0xcafe; 8], [0; 8]);
-            this.balance = 100;
-            this
-        };
+        let pinata_account = Account::new(PINATA_ADDRESS, INITIAL_BALANCE);
         accounts.insert(pinata_account.address, pinata_account);
 
         let commitment_tree = SparseMerkleTree::new_empty();
@@ -75,23 +67,60 @@ impl MockedSequencer {
     pub fn addresses(&self) -> Vec<Address> {
         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());
-        let formatted: Vec<String> = self
-            .nullifier_set
-            .iter()
-            .map(|arr| format!("0x{:x}", arr[0]))
-            .collect();
-        println!("Nullifiers: [{}]", formatted.join(", "));
-        println!("");
-        println!("");
-    }
+}
+
+pub fn print_accounts(sequencer: &MockedSequencer, private_accounts: &[&Account]) {
+    println!("\n====================== ACCOUNT SNAPSHOT ======================\n");
+
+    println!(">> Public Accounts:");
+    println!("{:<20} | {:>10} |", "Address (first u32)", "Balance");
+    println!("{:-<20}-+-{:-<10}", "", "");
+
+    for account in sequencer.accounts.values() {
+        println!("0x{:<20x} | {:>10} |", account.address[0], account.balance);
+    }
+
+    println!("{:-<20}-+-{:-<10}\n", "", "");
+
+    println!(">> Commitments:");
+    println!("{:-<20}", "");
+
+    for commitment in sequencer.commitment_tree.values().iter() {
+        println!("{:<20x}", commitment);
+    }
+
+    println!("{:-<20}\n", "");
+
+    let formatted: Vec<String> = sequencer
+        .nullifier_set
+        .iter()
+        .map(|nullifier| format!("0x{:x}", nullifier[0]))
+        .collect();
+
+    println!(">> Nullifiers (first u32):");
+    println!("{:-<20}", "");
+
+    for entry in formatted {
+        println!("{:<20}", entry);
+    }
+
+    println!("{:-<20}\n", "");
+
+    println!(">> Private Accounts:");
+    println!(
+        "{:<20} | {:>10} | {:>10}",
+        "Address (first u32)", "Nonce", "Balance"
+    );
+    println!("{:-<20}-+-{:-<10}-+-{:-<10}", "", "", "");
+
+    for account in private_accounts.iter() {
+        println!(
+            "{:<20x} | {:>10x}| {:>10} | ",
+            account.address[0], account.nonce[0], account.balance,
+        );
+    }
+
+    println!("{:-<20}-+-{:-<10}-+-{:-<10}", "", "", "");
+
+    println!("\n=============================================================\n");
 }

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

@@ -1,7 +1,7 @@
 use core::{
     account::Account,
     bytes_to_words,
-    input::InputVisibiility,
+    visibility::InputVisibiility,
     types::{Address, AuthenticationPath, Commitment, Nullifier},
 };
 use nssa::program::TransferMultipleProgram;
@@ -22,7 +22,7 @@ fn main() {
     let sender_private_key = [1, 2, 3, 4, 4, 3, 2, 1];
     let sender = {
         // Creating it now but it's supposed to be already created by other previous transactions.
-        let mut account = Account::new_from_private_key(sender_private_key, [1; 8]);
+        let mut account = Account::new_from_private_key(sender_private_key);
         account.balance = 150;
         account
     };

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

@@ -10,24 +10,13 @@ use nssa::program::TransferMultipleProgram;
 /// the initiating transaction includes the sender's signature.
 pub fn main() {
     // Account fetched from the chain state with 150 in its balance.
-    let sender = {
-        let mut account = Account::new([5; 8], [98; 8]);
-        account.balance = 150;
-        account
-    };
+    let sender = Account::new([5; 8], 150);
 
     // Account fetched from the chain state with 900 in its balance.
-    let receiver_1 = {
-        let mut account = Account::new([6; 8], [99; 8]);
-        account.balance = 900;
-        account
-    };
+    let receiver_1 = Account::new([6; 8], 900);
 
-    let receiver_2 = {
-        let mut account = Account::new([6; 8], [99; 8]);
-        account.balance = 500;
-        account
-    };
+    // Account fetched from the chain state with 500 in its balance.
+    let receiver_2 = Account::new([6; 8], 500);
 
     let balance_to_move = vec![10, 20];
 

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

@@ -1,9 +1,8 @@
 use core::{
     account::Account,
-    compute_nullifier, hash,
-    input::InputVisibiility,
-    is_in_tree,
+    compute_nullifier, hash, is_in_tree,
     types::{Nonce, ProgramId},
+    visibility::InputVisibiility,
 };
 use risc0_zkvm::{guest::env, serde::to_vec};
 

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

@@ -1,7 +1,7 @@
 use core::{
     account::Account,
-    input::InputVisibiility,
     types::{AuthenticationPath, Commitment, Key, Nonce, Nullifier},
+    visibility::InputVisibiility,
 };
 use program_methods::{OUTER_ELF, OUTER_ID};
 use rand::{rngs::OsRng, Rng};