add tx types

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

@@ -51,7 +51,7 @@ fn main() {
         InputVisibiility::Private(None),
     ];
 
-    let receipt = nssa::execute_and_prove_privacy_execution::<TransferMultipleProgram>(
+    let (receipt, _) = nssa::invoke_privacy_execution::<TransferMultipleProgram>(
         &[sender, receiver_1, receiver_2],
         &vec![30, 40],
         &visibilities,

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

@@ -1,26 +1,170 @@
-use nssa::program::TransferProgram;
+use core::{
+    account::Account,
+    bytes_to_words, hash,
+    input::InputVisibiility,
+    types::{Address, Commitment, Key, Nullifier},
+};
 
-use crate::sequencer_mock::MockedSequencer;
+use nssa::program::TransferProgram;
+use risc0_zkvm::Receipt;
+
+use crate::sequencer_mock::{MockedSequencer, ACCOUNTS_PRIVATE_KEYS};
 
 mod sequencer_mock;
 
 fn main() {
     let mut sequencer = MockedSequencer::new();
     let addresses = sequencer.addresses();
+    println!("addresses: {:?}", addresses);
     println!("Initial balances");
     sequencer.print();
 
     // A public execution of the Transfer Program
-    let sender_addr = addresses[0];
-    let receiver_addr = addresses[1];
+    let sender_addr = addresses[1];
+    let receiver_addr = addresses[2];
     sequencer
         .invoke_public::<TransferProgram>(&[sender_addr, receiver_addr], 10)
         .unwrap();
-
     println!("Balances after transfer");
     sequencer.print();
 
+    // A shielded execution of the Transfer Program
+    let private_account_2 = send_shielded(&addresses[1], &addresses[2], 15, &mut sequencer);
+    println!("Balances after shielded execution");
+    sequencer.print();
 
     // A private execution of the Transfer Program
+    let private_account_1 = send_private(
+        &private_account_2,
+        &ACCOUNTS_PRIVATE_KEYS[1], // <-- this is shifted 🫠
+        &addresses[3],
+        8,
+        &mut sequencer,
+    );
+    println!("Balances after shielded execution");
+    sequencer.print();
 
+    // A deshielded execution of the Transfer Program
+    send_deshielded(
+        &private_account_1,
+        &ACCOUNTS_PRIVATE_KEYS[0],
+        &addresses[0],
+        1,
+        &mut sequencer,
+    );
+
+    println!("Balances after deshielded execution");
+    sequencer.print();
+}
+
+fn mint_fresh_account(address: Address) -> Account {
+    let nonce = [0; 8];
+    Account::new(address, nonce)
+}
+
+/// A shielded execution of the Transfer program
+fn send_shielded(
+    from_address: &Address,
+    to_address: &Address,
+    balance_to_move: u128,
+    sequencer: &mut MockedSequencer,
+) -> Account {
+    // All of this is executed locally by the sender
+    let sender_account = sequencer.get_account(&from_address).unwrap();
+    let commitment_tree_root = sequencer.get_commitment_tree_root();
+    let receiver_addr = to_address;
+    let mut receiver_account = mint_fresh_account(*receiver_addr);
+    let visibilities = vec![InputVisibiility::Public, InputVisibiility::Private(None)];
+    let (receipt, nonces) = nssa::invoke_privacy_execution::<TransferProgram>(
+        &[sender_account, receiver_account.clone()],
+        &balance_to_move,
+        &visibilities,
+        commitment_tree_root,
+    )
+    .unwrap();
+
+    // Assemble the private account
+    receiver_account.nonce = nonces[1];
+    receiver_account.balance = balance_to_move;
+    let output: (Vec<Account>, Vec<Nullifier>, Vec<Commitment>, [u32; 8]) =
+        receipt.journal.decode().unwrap();
+
+    // Send to te sequencer
+    sequencer
+        .invoke_privacy_execution(receipt, &output.0, &output.1, &output.2)
+        .unwrap();
+    receiver_account
+}
+
+/// A private execution of the Transfer program
+fn send_private(
+    from_account: &Account,
+    from_account_pk: &Key,
+    to_address: &Address,
+    balance_to_move: u128,
+    sequencer: &mut MockedSequencer,
+) -> Account {
+    // All of this is executed locally by the sender
+    let commitment_tree_root = sequencer.get_commitment_tree_root();
+    let receiver_addr = to_address;
+    let sender_commitment_auth_path =
+        sequencer.get_authentication_path_for(&from_account.commitment());
+    let mut receiver_account = mint_fresh_account(*receiver_addr);
+    let visibilities = vec![
+        InputVisibiility::Private(Some((from_account_pk.clone(), sender_commitment_auth_path))),
+        InputVisibiility::Private(None),
+    ];
+    let (receipt, nonces) = nssa::invoke_privacy_execution::<TransferProgram>(
+        &[from_account.clone(), receiver_account.clone()],
+        &balance_to_move,
+        &visibilities,
+        commitment_tree_root,
+    )
+    .unwrap();
+
+    // Assemble the private account
+    receiver_account.nonce = nonces[1];
+    receiver_account.balance = balance_to_move;
+    let output: (Vec<Account>, Vec<Nullifier>, Vec<Commitment>, [u32; 8]) =
+        receipt.journal.decode().unwrap();
+
+    // Send to te sequencer
+    sequencer
+        .invoke_privacy_execution(receipt, &output.0, &output.1, &output.2)
+        .unwrap();
+    receiver_account
+}
+
+fn send_deshielded(
+    from_account: &Account,
+    from_account_pk: &Key,
+    to_address: &Address,
+    balance_to_move: u128,
+    sequencer: &mut MockedSequencer,
+) {
+    // All of this is executed locally by the sender
+    let commitment_tree_root = sequencer.get_commitment_tree_root();
+    let receiver_addr = to_address;
+    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((from_account_pk.clone(), sender_commitment_auth_path))),
+        InputVisibiility::Public,
+    ];
+    let (receipt, nonces) = nssa::invoke_privacy_execution::<TransferProgram>(
+        &[from_account.clone(), to_account],
+        &balance_to_move,
+        &visibilities,
+        commitment_tree_root,
+    )
+    .unwrap();
+
+    let output: (Vec<Account>, Vec<Nullifier>, Vec<Commitment>, [u32; 8]) =
+        receipt.journal.decode().unwrap();
+
+    // Send to te sequencer
+    sequencer
+        .invoke_privacy_execution(receipt, &output.0, &output.1, &output.2)
+        .unwrap();
 }

risc0-selective-privacy-poc/examples/sequencer_mock/mod.rs

@@ -1,7 +1,7 @@
 use core::{
     account::Account,
     bytes_to_words,
-    types::{Address, Commitment, Key, Nullifier, ProgramId},
+    types::{Address, AuthenticationPath, Commitment, Key, Nullifier, ProgramId},
 };
 use std::collections::{BTreeMap, HashSet};
 
@@ -17,8 +17,8 @@ pub struct MockedSequencer {
     deployed_program_ids: HashSet<ProgramId>,
 }
 
-const ACCOUNTS_PRIVATE_KEYS: [Key; 3] = [[1; 8], [2; 8], [3; 8]];
-const ACCOUNTS_INITIAL_BALANCES: [u128; 3] = [100, 1337, 0];
+pub const ACCOUNTS_PRIVATE_KEYS: [Key; 3] = [[1; 8], [2; 8], [3; 8]];
+const ACCOUNTS_INITIAL_BALANCES: [u128; 3] = [100, 1337, 37];
 const DEPLOYED_PROGRAM_IDS: [ProgramId; 3] = [TRANSFER_ID, TRANSFER_MULTIPLE_ID, PINATA_ID];
 
 impl MockedSequencer {
@@ -63,13 +63,13 @@ impl MockedSequencer {
         nullifiers: &[Nullifier],
         commitments: &[Commitment],
     ) -> Result<(), ()> {
-        let commitments_tree_root = bytes_to_words(&self.commitment_tree.root());
+        let commitments_tree_root = self.get_commitment_tree_root();
         if public_inputs_outputs.len() % 2 != 0 {
             return Err(());
         }
 
         let num_input_public = public_inputs_outputs.len() >> 1;
-        for account in public_inputs_outputs.iter() {
+        for account in public_inputs_outputs.iter().take(num_input_public) {
             let current_account = self.get_account(&account.address).ok_or(())?;
             if &current_account != account {
                 return Err(());
@@ -111,6 +111,14 @@ impl MockedSequencer {
                     .insert(account_post_state.address, account_post_state);
             });
 
+        // Add nullifiers
+        self.nullifier_set.extend(nullifiers);
+
+        // Add nullifiers
+        for commitment in commitments.iter() {
+            self.commitment_tree.add_value(*commitment);
+        }
+
         Ok(())
     }
 
@@ -192,8 +200,26 @@ impl MockedSequencer {
         }
         println!("{:-<20}-+-{:-<10}", "", "");
         println!("Commitments: {:?}", self.commitment_tree.values());
-        println!("Nullifiers: {:?}", self.nullifier_set);
+        let formatted: Vec<String> = self.nullifier_set
+            .iter()
+            .map(|arr| format!("0x{:x}", arr[0]))
+            .collect();
+        println!("Nullifiers: [{}]", formatted.join(", "));
         println!("");
         println!("");
     }
+
+    pub fn get_commitment_tree_root(&self) -> [u32; 8] {
+        bytes_to_words(&self.commitment_tree.root())
+    }
+
+    pub fn get_authentication_path_for(&self, commitment: &Commitment) -> AuthenticationPath {
+        self.commitment_tree
+            .get_authentication_path_for_value(*commitment)
+            .iter()
+            .map(bytes_to_words)
+            .collect::<Vec<_>>()
+            .try_into()
+            .unwrap()
+    }
 }

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

@@ -3,15 +3,14 @@ use core::{
     input::InputVisibiility,
     types::{Commitment, Nonce, Nullifier},
 };
+use program_methods::{OUTER_ELF, OUTER_ID};
 use rand::{rngs::OsRng, Rng};
 use risc0_zkvm::{default_executor, default_prover, ExecutorEnv, ExecutorEnvBuilder, Receipt};
-use program_methods::{OUTER_ELF, OUTER_ID};
 
 pub mod program;
 
 pub use program::Program;
 
-
 pub fn new_random_nonce() -> Nonce {
     let mut rng = OsRng;
     std::array::from_fn(|_| rng.gen())
@@ -67,12 +66,12 @@ pub fn execute<P: Program>(
     Ok(inputs_outputs)
 }
 
-pub fn execute_and_prove_privacy_execution<P: Program>(
+pub fn invoke_privacy_execution<P: Program>(
     inputs: &[Account],
     instruction_data: &P::InstructionData,
     visibilities: &[InputVisibiility],
     commitment_tree_root: [u32; 8],
-) -> Result<Receipt, ()> {
+) -> Result<(Receipt, Vec<Nonce>), ()> {
     // Prove inner program and get post state of the accounts
     let num_inputs = inputs.len();
     let (inner_receipt, inputs_outputs) = execute_and_prove_inner::<P>(inputs, instruction_data)?;
@@ -94,7 +93,7 @@ pub fn execute_and_prove_privacy_execution<P: Program>(
 
     let prover = default_prover();
     let prove_info = prover.prove(env, OUTER_ELF).unwrap();
-    Ok(prove_info.receipt)
+    Ok((prove_info.receipt, output_nonces))
 }
 
 pub fn verify_privacy_execution(