refactor: bundle private/public i/o into action struccts

lee/state_machine/src/validated_state_diff/mod.rs

@@ -4,7 +4,7 @@ use std::{
 };
 
 use lee_core::{
-    BlockId, Commitment, Nullifier, PrivacyPreservingCircuitOutput, Timestamp,
+    BlockId, Commitment, Nullifier, PrivacyPreservingCircuitOutput, PublicAction, Timestamp,
     account::{Account, AccountId, AccountWithMetadata},
     program::{
         ChainedCall, Claim, DEFAULT_PROGRAM_ID, ProgramId, compute_public_authorized_pdas,
@@ -321,10 +321,13 @@ impl ValidatedStateDiff {
     ) -> Result<Self, LeeError> {
         let message = &tx.message;
         let witness_set = &tx.witness_set;
+        let commitments = message.commitments();
+        let nullifiers = message.nullifiers();
+        let public_account_ids = message.public_account_ids();
 
         // 1. Commitments or nullifiers are non empty
         ensure!(
-            !message.new_commitments.is_empty() || !message.new_nullifiers.is_empty(),
+            !message.private_actions.is_empty(),
             LeeError::InvalidInput(
                 "Empty commitments and empty nullifiers found in message".into(),
             )
@@ -332,25 +335,19 @@ impl ValidatedStateDiff {
 
         // 2. Check there are no duplicate account_ids in the public_account_ids list.
         ensure!(
-            n_unique(&message.public_account_ids) == message.public_account_ids.len(),
+            n_unique(&public_account_ids) == public_account_ids.len(),
             LeeError::InvalidInput("Duplicate account_ids found in message".into())
         );
 
         // Check there are no duplicate nullifiers in the new_nullifiers list
         ensure!(
-            n_unique(
-                &message
-                    .new_nullifiers
-                    .iter()
-                    .map(|(n, _)| n)
-                    .collect::<Vec<_>>()
-            ) == message.new_nullifiers.len(),
+            n_unique(&nullifiers.iter().map(|(n, _)| n).collect::<Vec<_>>()) == nullifiers.len(),
             LeeError::InvalidInput("Duplicate nullifiers found in message".into())
         );
 
         // Check there are no duplicate commitments in the new_commitments list
         ensure!(
-            n_unique(&message.new_commitments) == message.new_commitments.len(),
+            n_unique(&commitments) == commitments.len(),
             LeeError::InvalidInput("Duplicate commitments found in message".into())
         );
 
@@ -387,8 +384,7 @@ impl ValidatedStateDiff {
         );
 
         // Build pre_states for proof verification
-        let public_pre_states: Vec<_> = message
-            .public_account_ids
+        let public_pre_states: Vec<_> = public_account_ids
             .iter()
             .map(|account_id| {
                 AccountWithMetadata::new(
@@ -407,28 +403,22 @@ impl ValidatedStateDiff {
         )?;
 
         // 5. Commitment freshness
-        state.check_commitments_are_new(&message.new_commitments)?;
+        state.check_commitments_are_new(&commitments)?;
 
         // 6. Nullifier uniqueness
-        state.check_nullifiers_are_valid(&message.new_nullifiers)?;
+        state.check_nullifiers_are_valid(&nullifiers)?;
 
-        let public_diff = message
-            .public_account_ids
+        let public_diff = public_account_ids
             .iter()
             .copied()
-            .zip(message.public_post_states.clone())
-            .collect();
-        let new_nullifiers = message
-            .new_nullifiers
-            .iter()
-            .copied()
-            .map(|(nullifier, _)| nullifier)
+            .zip(message.public_post_states())
             .collect();
+        let new_nullifiers = nullifiers.iter().map(|(nullifier, _)| *nullifier).collect();
 
         Ok(Self(StateDiff {
             signer_account_ids,
             public_diff,
-            new_commitments: message.new_commitments.clone(),
+            new_commitments: commitments,
             new_nullifiers,
             program: None,
         }))
@@ -509,11 +499,13 @@ fn check_privacy_preserving_circuit_proof_is_valid(
     message: &Message,
 ) -> Result<(), LeeError> {
     let output = PrivacyPreservingCircuitOutput {
-        public_pre_states: public_pre_states.to_vec(),
-        public_post_states: message.public_post_states.clone(),
-        encrypted_private_post_states: message.encrypted_private_post_states.clone(),
-        new_commitments: message.new_commitments.clone(),
-        new_nullifiers: message.new_nullifiers.clone(),
+        public_actions: public_pre_states
+            .iter()
+            .cloned()
+            .zip(message.public_post_states())
+            .map(|(pre, post)| PublicAction { pre, post })
+            .collect(),
+        private_actions: message.private_actions.clone(),
         block_validity_window: message.block_validity_window,
         timestamp_validity_window: message.timestamp_validity_window,
     };

lee/state_machine/src/validated_state_diff/tests.rs

@@ -188,12 +188,10 @@ fn privacy_malicious_programs_cannot_drain_public_victim() {
 
     // public_account_ids lists the Public entries from account_identities, in order.
     // The single ciphertext belongs to attacker's private account update.
-    let message = Message::try_from_circuit_output(
-        vec![victim_id, recipient_id],
+    let message = Message::from_circuit_output(
         vec![], // no public signers, no nonces
         circuit_output,
-    )
-    .unwrap();
+    );
 
     let witness_set = WitnessSet::for_message(&message, proof, &[]); // no signatures
     let tx = PrivacyPreservingTransaction::new(message, witness_set);
@@ -350,12 +348,10 @@ fn privacy_malicious_programs_cannot_drain_private_victim() {
 
     // public_account_ids lists the Public entries from account_identities, in order.
     // The single ciphertext belongs to attacker's private account update.
-    let message = Message::try_from_circuit_output(
-        vec![victim_id, recipient_id],
+    let message = Message::from_circuit_output(
         vec![], // no public signers, no nonces
         circuit_output,
-    )
-    .unwrap();
+    );
 
     let witness_set = WitnessSet::for_message(&message, proof, &[]); // no signatures
     let tx = PrivacyPreservingTransaction::new(message, witness_set);
@@ -480,8 +476,9 @@ fn malicious_programs_cannot_drain_victim_without_signature() {
 #[test]
 fn privacy_garbage_proof_is_rejected() {
     use lee_core::{
-        Commitment,
+        Commitment, EncryptedAccountData, Nullifier, PrivateAction,
         account::Account,
+        encryption::{Ciphertext, EphemeralPublicKey},
         program::{BlockValidityWindow, TimestampValidityWindow},
     };
 
@@ -503,12 +500,18 @@ fn privacy_garbage_proof_is_rejected() {
     ));
     let commitment = Commitment::new(&account_id, &Account::default());
     let message = Message {
-        public_account_ids: vec![],
+        public_actions: vec![],
         nonces: vec![],
-        public_post_states: vec![],
-        encrypted_private_post_states: vec![],
-        new_commitments: vec![commitment],
-        new_nullifiers: vec![],
+        private_actions: vec![PrivateAction {
+            nullifier: Nullifier::for_account_initialization(&account_id),
+            root: [0; 32],
+            commitment,
+            encrypted_post_state: EncryptedAccountData {
+                ciphertext: Ciphertext::from_inner(vec![]),
+                epk: EphemeralPublicKey(vec![]),
+                view_tag: 0,
+            },
+        }],
         block_validity_window: BlockValidityWindow::new_unbounded(),
         timestamp_validity_window: TimestampValidityWindow::new_unbounded(),
     };