refactor: strong-type PrivacyPreservingCircuitInput with per-account enum

integration_tests/tests/tps.rs

@@ -27,7 +27,7 @@ use nssa::{
     public_transaction as putx,
 };
 use nssa_core::{
-    MembershipProof, NullifierPublicKey,
+    MembershipProof, NullifierPublicKey, PrivacyPreservingCircuitInputAccount,
     account::{AccountWithMetadata, Nonce, data::Data},
     encryption::ViewingPublicKey,
 };
@@ -248,10 +248,17 @@ fn build_privacy_transaction() -> PrivacyPreservingTransaction {
     let (output, proof) = circuit::execute_and_prove(
         vec![sender_pre, recipient_pre],
         Program::serialize_instruction(balance_to_move).unwrap(),
-        vec![1, 2],
-        vec![(sender_npk, sender_ss), (recipient_npk, recipient_ss)],
-        vec![sender_nsk],
-        vec![Some(proof)],
+        vec![
+            PrivacyPreservingCircuitInputAccount::PrivateAuthorizedUpdate {
+                ssk: sender_ss,
+                nsk: sender_nsk,
+                membership_proof: proof,
+            },
+            PrivacyPreservingCircuitInputAccount::PrivateUnauthorized {
+                npk: recipient_npk,
+                ssk: recipient_ss,
+            },
+        ],
         &program.into(),
     )
     .unwrap();

nssa/core/src/circuit_io.rs

@@ -12,23 +12,83 @@ use crate::{
 pub struct PrivacyPreservingCircuitInput {
     /// Outputs of the program execution.
     pub program_outputs: Vec<ProgramOutput>,
-    /// Visibility mask for accounts.
-    ///
-    /// - `0` - public account
-    /// - `1` - private account with authentication
-    /// - `2` - private account without authentication
-    /// - `3` - private PDA account
-    pub visibility_mask: Vec<u8>,
-    /// Public keys of private accounts.
-    pub private_account_keys: Vec<(NullifierPublicKey, SharedSecretKey)>,
-    /// Nullifier secret keys for authorized private accounts.
-    pub private_account_nsks: Vec<NullifierSecretKey>,
-    /// Membership proofs for private accounts. Can be [`None`] for uninitialized accounts.
-    pub private_account_membership_proofs: Vec<Option<MembershipProof>>,
+    /// One entry per `pre_state`, in the same order as the program's `pre_states`.
+    /// Length must equal the number of `pre_states` derived from `program_outputs`.
+    /// The guest's `private_pda_npk_by_position` and `private_pda_bound_positions`
+    /// rely on this position alignment.
+    pub accounts: Vec<PrivacyPreservingCircuitInputAccount>,
     /// Program ID.
     pub program_id: ProgramId,
 }
 
+/// Per-account input to the privacy-preserving circuit. Each variant carries exactly the fields
+/// the guest needs for that account's code path.
+#[derive(Serialize, Deserialize, Clone)]
+pub enum PrivacyPreservingCircuitInputAccount {
+    /// Public account. The guest reads pre/post state from `program_outputs` and emits no
+    /// commitment, ciphertext, or nullifier.
+    Public,
+    /// Init of an authorized standalone private account (mask 1, no membership proof). The
+    /// `pre_state` must be `Account::default()`. `npk` is derived from `nsk` and matched
+    /// against `pre_state.account_id` via `AccountId::from(npk)`.
+    PrivateAuthorizedInit {
+        ssk: SharedSecretKey,
+        nsk: NullifierSecretKey,
+    },
+    /// Update of an authorized standalone private account (mask 1, with membership proof).
+    PrivateAuthorizedUpdate {
+        ssk: SharedSecretKey,
+        nsk: NullifierSecretKey,
+        membership_proof: MembershipProof,
+    },
+    /// Unauthorized init of a standalone private account (mask 2). Used for recipients who
+    /// don't yet exist on chain. No `nsk`, no membership proof.
+    PrivateUnauthorized {
+        npk: NullifierPublicKey,
+        ssk: SharedSecretKey,
+    },
+    /// Init of a private PDA (mask 3, unauthorized). The npk-to-account_id binding is proven
+    /// upstream via `Claim::Pda(seed)` or a caller's `pda_seeds` match.
+    PrivatePdaInit {
+        npk: NullifierPublicKey,
+        ssk: SharedSecretKey,
+    },
+    /// Update of an existing private PDA (mask 3, authorized, with membership proof). `npk` is
+    /// derived from `nsk`. Authorization is established upstream by a caller `pda_seeds` match
+    /// or a previously-seen authorization in a chained call.
+    PrivatePdaUpdate {
+        ssk: SharedSecretKey,
+        nsk: NullifierSecretKey,
+        membership_proof: MembershipProof,
+    },
+}
+
+impl PrivacyPreservingCircuitInputAccount {
+    #[must_use]
+    pub const fn is_public(&self) -> bool {
+        matches!(self, Self::Public)
+    }
+
+    #[must_use]
+    pub const fn is_private_pda(&self) -> bool {
+        matches!(self, Self::PrivatePdaInit { .. } | Self::PrivatePdaUpdate { .. })
+    }
+
+    /// For private PDA variants, return the nullifier public key. `Init` carries it directly;
+    /// `Update` derives it from `nsk`. For non-PDA variants returns `None`.
+    #[must_use]
+    pub fn npk_if_private_pda(&self) -> Option<NullifierPublicKey> {
+        match self {
+            Self::PrivatePdaInit { npk, .. } => Some(*npk),
+            Self::PrivatePdaUpdate { nsk, .. } => Some(NullifierPublicKey::from(nsk)),
+            Self::Public
+            | Self::PrivateAuthorizedInit { .. }
+            | Self::PrivateAuthorizedUpdate { .. }
+            | Self::PrivateUnauthorized { .. } => None,
+        }
+    }
+}
+
 #[derive(Serialize, Deserialize)]
 #[cfg_attr(any(feature = "host", test), derive(Debug, PartialEq, Eq))]
 pub struct PrivacyPreservingCircuitOutput {

nssa/core/src/lib.rs

@@ -3,7 +3,10 @@
     reason = "We prefer to group methods by functionality rather than by type for encoding"
 )]
 
-pub use circuit_io::{PrivacyPreservingCircuitInput, PrivacyPreservingCircuitOutput};
+pub use circuit_io::{
+    PrivacyPreservingCircuitInput, PrivacyPreservingCircuitInputAccount,
+    PrivacyPreservingCircuitOutput,
+};
 pub use commitment::{
     Commitment, CommitmentSetDigest, DUMMY_COMMITMENT, DUMMY_COMMITMENT_HASH, MembershipProof,
     compute_digest_for_path,

nssa/src/privacy_preserving_transaction/circuit.rs

@@ -2,8 +2,8 @@ use std::collections::{HashMap, VecDeque};
 
 use borsh::{BorshDeserialize, BorshSerialize};
 use nssa_core::{
-    MembershipProof, NullifierPublicKey, NullifierSecretKey, PrivacyPreservingCircuitInput,
-    PrivacyPreservingCircuitOutput, SharedSecretKey,
+    PrivacyPreservingCircuitInput, PrivacyPreservingCircuitInputAccount,
+    PrivacyPreservingCircuitOutput,
     account::AccountWithMetadata,
     program::{ChainedCall, InstructionData, ProgramId, ProgramOutput},
 };
@@ -63,14 +63,10 @@ impl From<Program> for ProgramWithDependencies {
 
 /// Generates a proof of the execution of a NSSA program inside the privacy preserving execution
 /// circuit.
-/// TODO: too many parameters.
 pub fn execute_and_prove(
     pre_states: Vec<AccountWithMetadata>,
     instruction_data: InstructionData,
-    visibility_mask: Vec<u8>,
-    private_account_keys: Vec<(NullifierPublicKey, SharedSecretKey)>,
-    private_account_nsks: Vec<NullifierSecretKey>,
-    private_account_membership_proofs: Vec<Option<MembershipProof>>,
+    accounts: Vec<PrivacyPreservingCircuitInputAccount>,
     program_with_dependencies: &ProgramWithDependencies,
 ) -> Result<(PrivacyPreservingCircuitOutput, Proof), NssaError> {
     let ProgramWithDependencies {
@@ -128,10 +124,7 @@ pub fn execute_and_prove(
 
     let circuit_input = PrivacyPreservingCircuitInput {
         program_outputs,
-        visibility_mask,
-        private_account_keys,
-        private_account_nsks,
-        private_account_membership_proofs,
+        accounts,
         program_id: program_with_dependencies.program.id(),
     };
 
@@ -243,10 +236,13 @@ mod tests {
         let (output, proof) = execute_and_prove(
             vec![sender, recipient],
             Program::serialize_instruction(balance_to_move).unwrap(),
-            vec![0, 2],
-            vec![(recipient_keys.npk(), shared_secret)],
-            vec![],
-            vec![None],
+            vec![
+                PrivacyPreservingCircuitInputAccount::Public,
+                PrivacyPreservingCircuitInputAccount::PrivateUnauthorized {
+                    npk: recipient_keys.npk(),
+                    ssk: shared_secret,
+                },
+            ],
             &Program::authenticated_transfer_program().into(),
         )
         .unwrap();
@@ -339,13 +335,19 @@ mod tests {
         let (output, proof) = execute_and_prove(
             vec![sender_pre, recipient],
             Program::serialize_instruction(balance_to_move).unwrap(),
-            vec![1, 2],
             vec![
-                (sender_keys.npk(), shared_secret_1),
-                (recipient_keys.npk(), shared_secret_2),
+                PrivacyPreservingCircuitInputAccount::PrivateAuthorizedUpdate {
+                    ssk: shared_secret_1,
+                    nsk: sender_keys.nsk,
+                    membership_proof: commitment_set
+                        .get_proof_for(&commitment_sender)
+                        .expect("sender's commitment must be in the set"),
+                },
+                PrivacyPreservingCircuitInputAccount::PrivateUnauthorized {
+                    npk: recipient_keys.npk(),
+                    ssk: shared_secret_2,
+                },
             ],
-            vec![sender_keys.nsk],
-            vec![commitment_set.get_proof_for(&commitment_sender), None],
             &program.into(),
         )
         .unwrap();
@@ -408,10 +410,10 @@ mod tests {
         let result = execute_and_prove(
             vec![pre],
             instruction,
-            vec![2],
-            vec![(account_keys.npk(), shared_secret)],
-            vec![],
-            vec![None],
+            vec![PrivacyPreservingCircuitInputAccount::PrivateUnauthorized {
+                npk: account_keys.npk(),
+                ssk: shared_secret,
+            }],
             &program_with_deps,
         );
 

program_methods/guest/src/bin/privacy_preserving_circuit.rs

@@ -6,7 +6,7 @@ use std::{
 use nssa_core::{
     Commitment, CommitmentSetDigest, DUMMY_COMMITMENT_HASH, EncryptionScheme, MembershipProof,
     Nullifier, NullifierPublicKey, NullifierSecretKey, PrivacyPreservingCircuitInput,
-    PrivacyPreservingCircuitOutput, SharedSecretKey,
+    PrivacyPreservingCircuitInputAccount, PrivacyPreservingCircuitOutput, SharedSecretKey,
     account::{Account, AccountId, AccountWithMetadata, Nonce},
     compute_digest_for_path,
     program::{
@@ -43,39 +43,28 @@ struct ExecutionState {
     /// `AccountId` entry or as an equality check against the existing one, making the rule: one
     /// `(program, seed)` → one account per tx.
     pda_family_binding: HashMap<(ProgramId, PdaSeed), AccountId>,
-    /// Map from a mask-3 `pre_state`'s position in `visibility_mask` to the npk supplied for
-    /// that position in `private_account_keys`. Built once in `derive_from_outputs` by walking
-    /// `visibility_mask` in lock-step with `private_account_keys`, used later by the claim and
-    /// caller-seeds authorization paths.
+    /// Map from a private-PDA `pre_state`'s position in `accounts` to the npk that variant
+    /// supplies for that position. Populated once in `derive_from_outputs` by walking
+    /// `accounts` and consulting `npk_if_private_pda`. Used later by the claim and
+    /// caller-seeds authorization paths to verify
+    /// `AccountId::for_private_pda(program_id, seed, npk) == pre_state.account_id`.
     private_pda_npk_by_position: HashMap<usize, NullifierPublicKey>,
 }
 
 impl ExecutionState {
     /// Validate program outputs and derive the overall execution state.
     pub fn derive_from_outputs(
-        visibility_mask: &[u8],
-        private_account_keys: &[(NullifierPublicKey, SharedSecretKey)],
+        accounts: &[PrivacyPreservingCircuitInputAccount],
         program_id: ProgramId,
         program_outputs: Vec<ProgramOutput>,
     ) -> Self {
-        // Build position → npk map for mask-3 pre_states. `private_account_keys` is consumed in
-        // pre_state order across all masks 1/2/3, so walk `visibility_mask` in lock-step. The
-        // downstream `compute_circuit_output` also consumes the same iterator and its trailing
-        // assertions catch an over-supply of keys; under-supply surfaces here.
+        // Build position → npk map for private-PDA pre_states, indexed by position in `accounts`.
+        // The `accounts` vec is documented as 1:1 with the program's pre_state order, so position
+        // here matches `pre_state_position` used downstream in `validate_and_sync_states`.
         let mut private_pda_npk_by_position: HashMap<usize, NullifierPublicKey> = HashMap::new();
-        {
-            let mut keys_iter = private_account_keys.iter();
-            for (pos, &mask) in visibility_mask.iter().enumerate() {
-                if matches!(mask, 1..=3) {
-                    let (npk, _) = keys_iter.next().unwrap_or_else(|| {
-                        panic!(
-                            "private_account_keys shorter than visibility_mask demands: no key for masked position {pos} (mask {mask})"
-                        )
-                    });
-                    if mask == 3 {
-                        private_pda_npk_by_position.insert(pos, *npk);
-                    }
-                }
+        for (pos, account) in accounts.iter().enumerate() {
+            if let Some(npk) = account.npk_if_private_pda() {
+                private_pda_npk_by_position.insert(pos, npk);
             }
         }
 
@@ -192,7 +181,7 @@ impl ExecutionState {
             }
 
             execution_state.validate_and_sync_states(
-                visibility_mask,
+                accounts,
                 chained_call.program_id,
                 caller_program_id,
                 &chained_call.pda_seeds,
@@ -209,12 +198,12 @@ impl ExecutionState {
             "Inner call without a chained call found",
         );
 
-        // Every mask-3 pre_state must have had its npk bound to its account_id, either via a
-        // `Claim::Pda(seed)` in some program's post_state or via a caller's `pda_seeds` matching
-        // the private derivation. An unbound mask-3 pre_state has no cryptographic link between
-        // the supplied npk and the account_id, and must be rejected.
-        for (pos, &mask) in visibility_mask.iter().enumerate() {
-            if mask == 3 {
+        // Every private-PDA pre_state must have had its npk bound to its account_id, either via
+        // a `Claim::Pda(seed)` in some program's post_state or via a caller's `pda_seeds`
+        // matching the private derivation. An unbound private-PDA pre_state has no
+        // cryptographic link between the supplied npk and the account_id, and must be rejected.
+        for (pos, account) in accounts.iter().enumerate() {
+            if account.is_private_pda() {
                 assert!(
                     execution_state.private_pda_bound_positions.contains(&pos),
                     "private PDA pre_state at position {pos} has no proven (seed, npk) binding via Claim::Pda or caller pda_seeds"
@@ -249,7 +238,7 @@ impl ExecutionState {
     /// Validate program pre and post states and populate the execution state.
     fn validate_and_sync_states(
         &mut self,
-        visibility_mask: &[u8],
+        accounts: &[PrivacyPreservingCircuitInputAccount],
         program_id: ProgramId,
         caller_program_id: Option<ProgramId>,
         caller_pda_seeds: &[PdaSeed],
@@ -327,9 +316,9 @@ impl ExecutionState {
                     .position(|acc| acc.account_id == pre_account_id)
                     .expect("Pre state must exist at this point");
 
-                let mask = visibility_mask[pre_state_position];
-                match mask {
-                    0 => match claim {
+                let account = &accounts[pre_state_position];
+                if account.is_public() {
+                    match claim {
                         Claim::Authorized => {
                             // Note: no need to check authorized pdas because we have already
                             // checked consistency of authorization above.
@@ -351,40 +340,38 @@ impl ExecutionState {
                                 pre_account_id,
                             );
                         }
-                    },
-                    3 => {
-                        match claim {
-                            Claim::Authorized => {
-                                assert!(
-                                    pre_is_authorized,
-                                    "Cannot claim unauthorized private PDA {pre_account_id}"
-                                );
-                            }
-                            Claim::Pda(seed) => {
-                                let npk = self
+                    }
+                } else if account.is_private_pda() {
+                    match claim {
+                        Claim::Authorized => {
+                            assert!(
+                                pre_is_authorized,
+                                "Cannot claim unauthorized private PDA {pre_account_id}"
+                            );
+                        }
+                        Claim::Pda(seed) => {
+                            let npk = self
                                 .private_pda_npk_by_position
                                 .get(&pre_state_position)
                                 .expect("private PDA pre_state must have an npk in the position map");
-                                let pda = AccountId::for_private_pda(&program_id, &seed, npk);
-                                assert_eq!(
-                                    pre_account_id, pda,
-                                    "Invalid private PDA claim for account {pre_account_id}"
-                                );
-                                self.private_pda_bound_positions.insert(pre_state_position);
-                                assert_family_binding(
-                                    &mut self.pda_family_binding,
-                                    program_id,
-                                    seed,
-                                    pre_account_id,
-                                );
-                            }
+                            let pda = AccountId::for_private_pda(&program_id, &seed, npk);
+                            assert_eq!(
+                                pre_account_id, pda,
+                                "Invalid private PDA claim for account {pre_account_id}"
+                            );
+                            self.private_pda_bound_positions.insert(pre_state_position);
+                            assert_family_binding(
+                                &mut self.pda_family_binding,
+                                program_id,
+                                seed,
+                                pre_account_id,
+                            );
                         }
                     }
-                    _ => {
-                        // Mask 1/2: standard private accounts don't enforce the claim semantics.
-                        // Unauthorized private claiming is intentionally allowed since operating
-                        // these accounts requires the npk/nsk keypair anyway.
-                    }
+                } else {
+                    // Standalone private accounts (mask 1/2): don't enforce the claim semantics.
+                    // Unauthorized private claiming is intentionally allowed since operating
+                    // these accounts requires the npk/nsk keypair anyway.
                 }
 
                 post.account_mut().program_owner = program_id;
@@ -486,10 +473,7 @@ fn resolve_authorization_and_record_bindings(
 
 fn compute_circuit_output(
     execution_state: ExecutionState,
-    visibility_mask: &[u8],
-    private_account_keys: &[(NullifierPublicKey, SharedSecretKey)],
-    private_account_nsks: &[NullifierSecretKey],
-    private_account_membership_proofs: &[Option<MembershipProof>],
+    accounts: &[PrivacyPreservingCircuitInputAccount],
 ) -> PrivacyPreservingCircuitOutput {
     let mut output = PrivacyPreservingCircuitOutput {
         public_pre_states: Vec::new(),
@@ -503,280 +487,241 @@ fn compute_circuit_output(
 
     let states_iter = execution_state.into_states_iter();
     assert_eq!(
-        visibility_mask.len(),
+        accounts.len(),
         states_iter.len(),
-        "Invalid visibility mask length"
+        "Invalid accounts length"
     );
 
-    let mut private_keys_iter = private_account_keys.iter();
-    let mut private_nsks_iter = private_account_nsks.iter();
-    let mut private_membership_proofs_iter = private_account_membership_proofs.iter();
-
     let mut output_index = 0;
-    for (account_visibility_mask, (pre_state, post_state)) in
-        visibility_mask.iter().copied().zip(states_iter)
-    {
-        match account_visibility_mask {
-            0 => {
-                // Public account
+    for (account, (pre_state, post_state)) in accounts.iter().zip(states_iter) {
+        match account {
+            PrivacyPreservingCircuitInputAccount::Public => {
                 output.public_pre_states.push(pre_state);
                 output.public_post_states.push(post_state);
             }
-            1 | 2 => {
-                let Some((npk, shared_secret)) = private_keys_iter.next() else {
-                    panic!("Missing private account key");
-                };
+            PrivacyPreservingCircuitInputAccount::PrivateAuthorizedInit { ssk, nsk } => {
+                let npk = NullifierPublicKey::from(nsk);
 
+                assert_eq!(
+                    AccountId::from(&npk),
+                    pre_state.account_id,
+                    "AccountId mismatch"
+                );
+                assert!(
+                    pre_state.is_authorized,
+                    "Pre-state not authorized for authenticated private account"
+                );
+                assert_eq!(
+                    pre_state.account,
+                    Account::default(),
+                    "Found new private account with non default values"
+                );
+
+                let new_nullifier = (
+                    Nullifier::for_account_initialization(&npk),
+                    DUMMY_COMMITMENT_HASH,
+                );
+                let new_nonce = pre_state.account.nonce.private_account_nonce_increment(nsk);
+
+                emit_private_output(
+                    &mut output,
+                    &mut output_index,
+                    post_state,
+                    &npk,
+                    ssk,
+                    new_nullifier,
+                    new_nonce,
+                );
+            }
+            PrivacyPreservingCircuitInputAccount::PrivateAuthorizedUpdate {
+                ssk,
+                nsk,
+                membership_proof,
+            } => {
+                let npk = NullifierPublicKey::from(nsk);
+
+                assert_eq!(
+                    AccountId::from(&npk),
+                    pre_state.account_id,
+                    "AccountId mismatch"
+                );
+                assert!(
+                    pre_state.is_authorized,
+                    "Pre-state not authorized for authenticated private account"
+                );
+
+                let new_nullifier = compute_update_nullifier_and_set_digest(
+                    membership_proof,
+                    &pre_state.account,
+                    &npk,
+                    nsk,
+                );
+                let new_nonce = pre_state.account.nonce.private_account_nonce_increment(nsk);
+
+                emit_private_output(
+                    &mut output,
+                    &mut output_index,
+                    post_state,
+                    &npk,
+                    ssk,
+                    new_nullifier,
+                    new_nonce,
+                );
+            }
+            PrivacyPreservingCircuitInputAccount::PrivateUnauthorized { npk, ssk } => {
                 assert_eq!(
                     AccountId::from(npk),
                     pre_state.account_id,
                     "AccountId mismatch"
                 );
-
-                let (new_nullifier, new_nonce) = if account_visibility_mask == 1 {
-                    // Private account with authentication
-
-                    let Some(nsk) = private_nsks_iter.next() else {
-                        panic!("Missing private account nullifier secret key");
-                    };
-
-                    // Verify the nullifier public key
-                    assert_eq!(
-                        npk,
-                        &NullifierPublicKey::from(nsk),
-                        "Nullifier public key mismatch"
-                    );
-
-                    // Check pre_state authorization
-                    assert!(
-                        pre_state.is_authorized,
-                        "Pre-state not authorized for authenticated private account"
-                    );
-
-                    let Some(membership_proof_opt) = private_membership_proofs_iter.next() else {
-                        panic!("Missing membership proof");
-                    };
-
-                    let new_nullifier = compute_nullifier_and_set_digest(
-                        membership_proof_opt.as_ref(),
-                        &pre_state.account,
-                        npk,
-                        nsk,
-                    );
-
-                    let new_nonce = pre_state.account.nonce.private_account_nonce_increment(nsk);
-
-                    (new_nullifier, new_nonce)
-                } else {
-                    // Private account without authentication
-
-                    assert_eq!(
-                        pre_state.account,
-                        Account::default(),
-                        "Found new private account with non default values",
-                    );
-
-                    assert!(
-                        !pre_state.is_authorized,
-                        "Found new private account marked as authorized."
-                    );
-
-                    let Some(membership_proof_opt) = private_membership_proofs_iter.next() else {
-                        panic!("Missing membership proof");
-                    };
-
-                    assert!(
-                        membership_proof_opt.is_none(),
-                        "Membership proof must be None for unauthorized accounts"
-                    );
-
-                    let nullifier = Nullifier::for_account_initialization(npk);
-
-                    let new_nonce = Nonce::private_account_nonce_init(npk);
-
-                    ((nullifier, DUMMY_COMMITMENT_HASH), new_nonce)
-                };
-                output.new_nullifiers.push(new_nullifier);
-
-                // Update post-state with new nonce
-                let mut post_with_updated_nonce = post_state;
-                post_with_updated_nonce.nonce = new_nonce;
-
-                // Compute commitment
-                let commitment_post = Commitment::new(npk, &post_with_updated_nonce);
-
-                // Encrypt and push post state
-                let encrypted_account = EncryptionScheme::encrypt(
-                    &post_with_updated_nonce,
-                    shared_secret,
-                    &commitment_post,
-                    output_index,
+                assert_eq!(
+                    pre_state.account,
+                    Account::default(),
+                    "Found new private account with non default values",
+                );
+                assert!(
+                    !pre_state.is_authorized,
+                    "Found new private account marked as authorized."
                 );
 
-                output.new_commitments.push(commitment_post);
-                output.ciphertexts.push(encrypted_account);
-                output_index = output_index
-                    .checked_add(1)
-                    .unwrap_or_else(|| panic!("Too many private accounts, output index overflow"));
+                let new_nullifier = (
+                    Nullifier::for_account_initialization(npk),
+                    DUMMY_COMMITMENT_HASH,
+                );
+                let new_nonce = Nonce::private_account_nonce_init(npk);
+
+                emit_private_output(
+                    &mut output,
+                    &mut output_index,
+                    post_state,
+                    npk,
+                    ssk,
+                    new_nullifier,
+                    new_nonce,
+                );
             }
-            3 => {
-                // Private PDA account. The supplied npk has already been bound to
-                // `pre_state.account_id` upstream in `validate_and_sync_states`, either via a
-                // `Claim::Pda(seed)` match or via a caller `pda_seeds` match, both of which
-                // assert `AccountId::for_private_pda(owner, seed, npk) == account_id`. The
-                // post-loop assertion in `derive_from_outputs` (see the
-                // `private_pda_bound_positions` check) guarantees that every mask-3
-                // position has been through at least one such binding, so this
-                // branch can safely use the wallet npk without re-verifying.
-                let Some((npk, shared_secret)) = private_keys_iter.next() else {
-                    panic!("Missing private account key");
-                };
-
-                let (new_nullifier, new_nonce) = if pre_state.is_authorized {
-                    // Existing private PDA with authentication (like mask 1)
-                    let Some(nsk) = private_nsks_iter.next() else {
-                        panic!("Missing private account nullifier secret key");
-                    };
-                    assert_eq!(
-                        npk,
-                        &NullifierPublicKey::from(nsk),
-                        "Nullifier public key mismatch"
-                    );
-
-                    let Some(membership_proof_opt) = private_membership_proofs_iter.next() else {
-                        panic!("Missing membership proof");
-                    };
-
-                    let new_nullifier = compute_nullifier_and_set_digest(
-                        membership_proof_opt.as_ref(),
-                        &pre_state.account,
-                        npk,
-                        nsk,
-                    );
-                    let new_nonce = pre_state.account.nonce.private_account_nonce_increment(nsk);
-                    (new_nullifier, new_nonce)
-                } else {
-                    // New private PDA (like mask 2). The default + unauthorized requirement
-                    // here rules out use cases like a fully-private multisig, which would need
-                    // a non-default, non-authorized private PDA input account.
-                    // TODO(private-pdas-pr-2/3): relax this once the wallet can supply a
-                    // `(seed, owner)` side input so the npk-to-account_id binding can be
-                    // re-verified for an existing private PDA without a `Claim::Pda` or caller
-                    // `pda_seeds` match.
-                    assert_eq!(
-                        pre_state.account,
-                        Account::default(),
-                        "New private PDA must be default"
-                    );
-
-                    let Some(membership_proof_opt) = private_membership_proofs_iter.next() else {
-                        panic!("Missing membership proof");
-                    };
-                    assert!(
-                        membership_proof_opt.is_none(),
-                        "Membership proof must be None for new accounts"
-                    );
-
-                    let nullifier = Nullifier::for_account_initialization(npk);
-                    let new_nonce = Nonce::private_account_nonce_init(npk);
-                    ((nullifier, DUMMY_COMMITMENT_HASH), new_nonce)
-                };
-                output.new_nullifiers.push(new_nullifier);
-
-                let mut post_with_updated_nonce = post_state;
-                post_with_updated_nonce.nonce = new_nonce;
-
-                let commitment_post = Commitment::new(npk, &post_with_updated_nonce);
-
-                let encrypted_account = EncryptionScheme::encrypt(
-                    &post_with_updated_nonce,
-                    shared_secret,
-                    &commitment_post,
-                    output_index,
+            PrivacyPreservingCircuitInputAccount::PrivatePdaInit { npk, ssk } => {
+                // The npk-to-account_id binding is established upstream in
+                // `validate_and_sync_states` via `Claim::Pda(seed)` or a caller `pda_seeds`
+                // match. Here we only enforce the init pre-conditions.
+                assert!(
+                    !pre_state.is_authorized,
+                    "PrivatePdaInit requires unauthorized pre_state"
+                );
+                assert_eq!(
+                    pre_state.account,
+                    Account::default(),
+                    "New private PDA must be default"
                 );
 
-                output.new_commitments.push(commitment_post);
-                output.ciphertexts.push(encrypted_account);
-                output_index = output_index
-                    .checked_add(1)
-                    .unwrap_or_else(|| panic!("Too many private accounts, output index overflow"));
+                let new_nullifier = (
+                    Nullifier::for_account_initialization(npk),
+                    DUMMY_COMMITMENT_HASH,
+                );
+                let new_nonce = Nonce::private_account_nonce_init(npk);
+
+                emit_private_output(
+                    &mut output,
+                    &mut output_index,
+                    post_state,
+                    npk,
+                    ssk,
+                    new_nullifier,
+                    new_nonce,
+                );
+            }
+            PrivacyPreservingCircuitInputAccount::PrivatePdaUpdate {
+                ssk,
+                nsk,
+                membership_proof,
+            } => {
+                let npk = NullifierPublicKey::from(nsk);
+
+                // The npk binding is established upstream. Authorization must already be set;
+                // an unauthorized PrivatePdaUpdate would mean the prover supplied an nsk for an
+                // unbound PDA, which the upstream binding check would have rejected anyway,
+                // but we assert here to fail fast and document the precondition.
+                assert!(
+                    pre_state.is_authorized,
+                    "PrivatePdaUpdate requires authorized pre_state"
+                );
+
+                let new_nullifier = compute_update_nullifier_and_set_digest(
+                    membership_proof,
+                    &pre_state.account,
+                    &npk,
+                    nsk,
+                );
+                let new_nonce = pre_state.account.nonce.private_account_nonce_increment(nsk);
+
+                emit_private_output(
+                    &mut output,
+                    &mut output_index,
+                    post_state,
+                    &npk,
+                    ssk,
+                    new_nullifier,
+                    new_nonce,
+                );
             }
-            _ => panic!("Invalid visibility mask value"),
         }
     }
 
-    assert!(
-        private_keys_iter.next().is_none(),
-        "Too many private account keys"
-    );
-
-    assert!(
-        private_nsks_iter.next().is_none(),
-        "Too many private account nullifier secret keys"
-    );
-
-    assert!(
-        private_membership_proofs_iter.next().is_none(),
-        "Too many private account membership proofs"
-    );
-
     output
 }
 
-fn compute_nullifier_and_set_digest(
-    membership_proof_opt: Option<&MembershipProof>,
+fn emit_private_output(
+    output: &mut PrivacyPreservingCircuitOutput,
+    output_index: &mut u32,
+    post_state: Account,
+    npk: &NullifierPublicKey,
+    shared_secret: &SharedSecretKey,
+    new_nullifier: (Nullifier, CommitmentSetDigest),
+    new_nonce: Nonce,
+) {
+    output.new_nullifiers.push(new_nullifier);
+
+    let mut post_with_updated_nonce = post_state;
+    post_with_updated_nonce.nonce = new_nonce;
+
+    let commitment_post = Commitment::new(npk, &post_with_updated_nonce);
+    let encrypted_account = EncryptionScheme::encrypt(
+        &post_with_updated_nonce,
+        shared_secret,
+        &commitment_post,
+        *output_index,
+    );
+
+    output.new_commitments.push(commitment_post);
+    output.ciphertexts.push(encrypted_account);
+    *output_index = output_index
+        .checked_add(1)
+        .unwrap_or_else(|| panic!("Too many private accounts, output index overflow"));
+}
+
+fn compute_update_nullifier_and_set_digest(
+    membership_proof: &MembershipProof,
     pre_account: &Account,
     npk: &NullifierPublicKey,
     nsk: &NullifierSecretKey,
 ) -> (Nullifier, CommitmentSetDigest) {
-    membership_proof_opt.as_ref().map_or_else(
-        || {
-            assert_eq!(
-                *pre_account,
-                Account::default(),
-                "Found new private account with non default values"
-            );
-
-            // Compute initialization nullifier
-            let nullifier = Nullifier::for_account_initialization(npk);
-            (nullifier, DUMMY_COMMITMENT_HASH)
-        },
-        |membership_proof| {
-            // Compute commitment set digest associated with provided auth path
-            let commitment_pre = Commitment::new(npk, pre_account);
-            let set_digest = compute_digest_for_path(&commitment_pre, membership_proof);
-
-            // Compute update nullifier
-            let nullifier = Nullifier::for_account_update(&commitment_pre, nsk);
-            (nullifier, set_digest)
-        },
-    )
+    let commitment_pre = Commitment::new(npk, pre_account);
+    let set_digest = compute_digest_for_path(&commitment_pre, membership_proof);
+    let nullifier = Nullifier::for_account_update(&commitment_pre, nsk);
+    (nullifier, set_digest)
 }
 
 fn main() {
     let PrivacyPreservingCircuitInput {
         program_outputs,
-        visibility_mask,
-        private_account_keys,
-        private_account_nsks,
-        private_account_membership_proofs,
+        accounts,
         program_id,
     } = env::read();
 
-    let execution_state = ExecutionState::derive_from_outputs(
-        &visibility_mask,
-        &private_account_keys,
-        program_id,
-        program_outputs,
-    );
+    let execution_state =
+        ExecutionState::derive_from_outputs(&accounts, program_id, program_outputs);
 
-    let output = compute_circuit_output(
-        execution_state,
-        &visibility_mask,
-        &private_account_keys,
-        &private_account_nsks,
-        &private_account_membership_proofs,
-    );
+    let output = compute_circuit_output(execution_state, &accounts);
 
     env::commit(&output);
 }

sequencer/core/src/lib.rs

@@ -1075,7 +1075,7 @@ mod tests {
             program::Program,
         };
         use nssa_core::{
-            SharedSecretKey,
+            PrivacyPreservingCircuitInputAccount, SharedSecretKey,
             account::AccountWithMetadata,
             encryption::{EphemeralPublicKey, EphemeralSecretKey, ViewingPublicKey},
         };
@@ -1109,10 +1109,10 @@ mod tests {
         let (output, proof) = execute_and_prove(
             vec![AccountWithMetadata::new(Account::default(), true, &npk)],
             Program::serialize_instruction(0_u128).unwrap(),
-            vec![1],
-            vec![(npk, shared_secret)],
-            vec![nsk],
-            vec![None],
+            vec![PrivacyPreservingCircuitInputAccount::PrivateAuthorizedInit {
+                ssk: shared_secret,
+                nsk,
+            }],
             &Program::authenticated_transfer_program().into(),
         )
         .unwrap();

wallet/src/lib.rs

@@ -390,13 +390,7 @@ impl WalletCore {
         let (output, proof) = nssa::privacy_preserving_transaction::circuit::execute_and_prove(
             pre_states,
             instruction_data,
-            acc_manager.visibility_mask().to_vec(),
-            private_account_keys
-                .iter()
-                .map(|keys| (keys.npk, keys.ssk))
-                .collect::<Vec<_>>(),
-            acc_manager.private_account_auth(),
-            acc_manager.private_account_membership_proofs(),
+            acc_manager.accounts(),
             &program.to_owned(),
         )
         .unwrap();

wallet/src/privacy_preserving_tx.rs

@@ -2,7 +2,8 @@ use anyhow::Result;
 use key_protocol::key_management::ephemeral_key_holder::EphemeralKeyHolder;
 use nssa::{AccountId, PrivateKey};
 use nssa_core::{
-    MembershipProof, NullifierPublicKey, NullifierSecretKey, SharedSecretKey,
+    MembershipProof, NullifierPublicKey, NullifierSecretKey, PrivacyPreservingCircuitInputAccount,
+    SharedSecretKey,
     account::{AccountWithMetadata, Nonce},
     encryption::{EphemeralPublicKey, ViewingPublicKey},
 };
@@ -51,7 +52,6 @@ enum State {
 
 pub struct AccountManager {
     states: Vec<State>,
-    visibility_mask: Vec<u8>,
 }
 
 impl AccountManager {
@@ -59,11 +59,10 @@ impl AccountManager {
         wallet: &WalletCore,
         accounts: Vec<PrivacyPreservingAccount>,
     ) -> Result<Self, ExecutionFailureKind> {
-        let mut pre_states = Vec::with_capacity(accounts.len());
-        let mut visibility_mask = Vec::with_capacity(accounts.len());
+        let mut states = Vec::with_capacity(accounts.len());
 
         for account in accounts {
-            let (state, mask) = match account {
+            let state = match account {
                 PrivacyPreservingAccount::Public(account_id) => {
                     let acc = wallet
                         .get_account_public(account_id)
@@ -73,37 +72,37 @@ impl AccountManager {
                     let sk = wallet.get_account_public_signing_key(account_id).cloned();
                     let account = AccountWithMetadata::new(acc.clone(), sk.is_some(), account_id);
 
-                    (State::Public { account, sk }, 0)
+                    State::Public { account, sk }
                 }
                 PrivacyPreservingAccount::PrivateOwned(account_id) => {
                     let pre = private_acc_preparation(wallet, account_id).await?;
-                    let mask = if pre.pre_state.is_authorized { 1 } else { 2 };
 
-                    (State::Private(pre), mask)
+                    State::Private(pre)
                 }
                 PrivacyPreservingAccount::PrivateForeign { npk, vpk } => {
                     let acc = nssa_core::account::Account::default();
                     let auth_acc = AccountWithMetadata::new(acc, false, &npk);
+                    let eph_holder = EphemeralKeyHolder::new(&npk);
+                    let ssk = eph_holder.calculate_shared_secret_sender(&vpk);
+                    let epk = eph_holder.generate_ephemeral_public_key();
                     let pre = AccountPreparedData {
                         nsk: None,
                         npk,
                         vpk,
                         pre_state: auth_acc,
                         proof: None,
+                        ssk,
+                        epk,
                     };
 
-                    (State::Private(pre), 2)
+                    State::Private(pre)
                 }
             };
 
-            pre_states.push(state);
-            visibility_mask.push(mask);
+            states.push(state);
         }
 
-        Ok(Self {
-            states: pre_states,
-            visibility_mask,
-        })
+        Ok(Self { states })
     }
 
     pub fn pre_states(&self) -> Vec<AccountWithMetadata> {
@@ -116,10 +115,6 @@ impl AccountManager {
             .collect()
     }
 
-    pub fn visibility_mask(&self) -> &[u8] {
-        &self.visibility_mask
-    }
-
     pub fn public_account_nonces(&self) -> Vec<Nonce> {
         self.states
             .iter()
@@ -134,37 +129,45 @@ impl AccountManager {
         self.states
             .iter()
             .filter_map(|state| match state {
-                State::Private(pre) => {
-                    let eph_holder = EphemeralKeyHolder::new(&pre.npk);
+                State::Private(pre) => Some(PrivateAccountKeys {
+                    npk: pre.npk,
+                    ssk: pre.ssk,
+                    vpk: pre.vpk.clone(),
+                    epk: pre.epk.clone(),
+                }),
+                State::Public { .. } => None,
+            })
+            .collect()
+    }
 
-                    Some(PrivateAccountKeys {
+    /// Build the per-account input vec for the privacy-preserving circuit. Each variant carries
+    /// exactly the fields the circuit's code path for that account needs, with the ephemeral
+    /// keys (`ssk`) drawn from the cached values that `private_account_keys` and the message
+    /// construction also use, so all three views agree on the same ephemeral key.
+    pub fn accounts(&self) -> Vec<PrivacyPreservingCircuitInputAccount> {
+        self.states
+            .iter()
+            .map(|state| match state {
+                State::Public { .. } => PrivacyPreservingCircuitInputAccount::Public,
+                State::Private(pre) => match (pre.nsk, pre.proof.clone()) {
+                    (Some(nsk), Some(membership_proof)) => {
+                        PrivacyPreservingCircuitInputAccount::PrivateAuthorizedUpdate {
+                            ssk: pre.ssk,
+                            nsk,
+                            membership_proof,
+                        }
+                    }
+                    (Some(nsk), None) => {
+                        PrivacyPreservingCircuitInputAccount::PrivateAuthorizedInit {
+                            ssk: pre.ssk,
+                            nsk,
+                        }
+                    }
+                    (None, _) => PrivacyPreservingCircuitInputAccount::PrivateUnauthorized {
                         npk: pre.npk,
-                        ssk: eph_holder.calculate_shared_secret_sender(&pre.vpk),
-                        vpk: pre.vpk.clone(),
-                        epk: eph_holder.generate_ephemeral_public_key(),
-                    })
-                }
-                State::Public { .. } => None,
-            })
-            .collect()
-    }
-
-    pub fn private_account_auth(&self) -> Vec<NullifierSecretKey> {
-        self.states
-            .iter()
-            .filter_map(|state| match state {
-                State::Private(pre) => pre.nsk,
-                State::Public { .. } => None,
-            })
-            .collect()
-    }
-
-    pub fn private_account_membership_proofs(&self) -> Vec<Option<MembershipProof>> {
-        self.states
-            .iter()
-            .filter_map(|state| match state {
-                State::Private(pre) => Some(pre.proof.clone()),
-                State::Public { .. } => None,
+                        ssk: pre.ssk,
+                    },
+                },
             })
             .collect()
     }
@@ -196,6 +199,13 @@ struct AccountPreparedData {
     vpk: ViewingPublicKey,
     pre_state: AccountWithMetadata,
     proof: Option<MembershipProof>,
+    /// Cached shared-secret key derived once at `AccountManager::new`. Reused for both the
+    /// circuit input variant (`accounts()`) and the message ephemeral-key tuples
+    /// (`private_account_keys()`), so all consumers see the same key. The corresponding
+    /// `EphemeralKeyHolder` uses `OsRng` and would produce a different value on a second call.
+    ssk: SharedSecretKey,
+    /// Cached ephemeral public key, paired with `ssk`.
+    epk: EphemeralPublicKey,
 }
 
 async fn private_acc_preparation(
@@ -226,11 +236,17 @@ async fn private_acc_preparation(
     // support from that in the wallet.
     let sender_pre = AccountWithMetadata::new(from_acc.clone(), true, &from_npk);
 
+    let eph_holder = EphemeralKeyHolder::new(&from_npk);
+    let ssk = eph_holder.calculate_shared_secret_sender(&from_vpk);
+    let epk = eph_holder.generate_ephemeral_public_key();
+
     Ok(AccountPreparedData {
         nsk: Some(nsk),
         npk: from_npk,
         vpk: from_vpk,
         pre_state: sender_pre,
         proof,
+        ssk,
+        epk,
     })
 }