Merge pull request #462 from logos-blockchain/moudy/feat-strong-type-circuit-input

integration_tests/tests/tps.rs

@@ -27,7 +27,7 @@ use nssa::{
     public_transaction as putx,
 };
 use nssa_core::{
-    MembershipProof, NullifierPublicKey,
+    InputAccountIdentity, MembershipProof, NullifierPublicKey,
     account::{AccountWithMetadata, Nonce, data::Data},
     encryption::ViewingPublicKey,
 };
@@ -251,10 +251,19 @@ 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, 0, sender_ss), (recipient_npk, 0, recipient_ss)],
-        vec![sender_nsk],
-        vec![Some(proof)],
+        vec![
+            InputAccountIdentity::PrivateAuthorizedUpdate {
+                ssk: sender_ss,
+                nsk: sender_nsk,
+                membership_proof: proof,
+                identifier: 0,
+            },
+            InputAccountIdentity::PrivateUnauthorized {
+                npk: recipient_npk,
+                ssk: recipient_ss,
+                identifier: 0,
+            },
+        ],
         &program.into(),
     )
     .unwrap();

nssa/core/src/circuit_io.rs

@@ -12,23 +12,92 @@ 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 and identifiers of private accounts.
-    pub private_account_keys: Vec<(NullifierPublicKey, Identifier, 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 account_identities: Vec<InputAccountIdentity>,
     /// 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 InputAccountIdentity {
+    /// 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: no membership proof. The `pre_state`
+    /// must be `Account::default()`. The `account_id` is derived as
+    /// `AccountId::from((&NullifierPublicKey::from(nsk), identifier))` and matched against
+    /// `pre_state.account_id`.
+    PrivateAuthorizedInit {
+        ssk: SharedSecretKey,
+        nsk: NullifierSecretKey,
+        identifier: Identifier,
+    },
+    /// Update of an authorized standalone private account: existing on-chain commitment, with
+    /// membership proof.
+    PrivateAuthorizedUpdate {
+        ssk: SharedSecretKey,
+        nsk: NullifierSecretKey,
+        membership_proof: MembershipProof,
+        identifier: Identifier,
+    },
+    /// Init of a standalone private account the caller does not own (e.g. a recipient who
+    /// doesn't yet exist on chain). No `nsk`, no membership proof.
+    PrivateUnauthorized {
+        npk: NullifierPublicKey,
+        ssk: SharedSecretKey,
+        identifier: Identifier,
+    },
+    /// Init of a private PDA, unauthorized. The npk-to-account_id binding is proven upstream
+    /// via `Claim::Pda(seed)` or a caller's `pda_seeds` match. Identifier is fixed by
+    /// convention to `PRIVATE_PDA_FIXED_IDENTIFIER` and not carried per-input.
+    PrivatePdaInit {
+        npk: NullifierPublicKey,
+        ssk: SharedSecretKey,
+    },
+    /// Update of an existing private PDA, 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. Identifier is fixed.
+    PrivatePdaUpdate {
+        ssk: SharedSecretKey,
+        nsk: NullifierSecretKey,
+        membership_proof: MembershipProof,
+    },
+}
+
+impl InputAccountIdentity {
+    #[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,9 @@
     reason = "We prefer to group methods by functionality rather than by type for encoding"
 )]
 
-pub use circuit_io::{PrivacyPreservingCircuitInput, PrivacyPreservingCircuitOutput};
+pub use circuit_io::{
+    InputAccountIdentity, PrivacyPreservingCircuitInput, 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,7 @@ use std::collections::{HashMap, VecDeque};
 
 use borsh::{BorshDeserialize, BorshSerialize};
 use nssa_core::{
-    Identifier, MembershipProof, NullifierPublicKey, NullifierSecretKey,
-    PrivacyPreservingCircuitInput, PrivacyPreservingCircuitOutput, SharedSecretKey,
+    InputAccountIdentity, PrivacyPreservingCircuitInput, PrivacyPreservingCircuitOutput,
     account::AccountWithMetadata,
     program::{ChainedCall, InstructionData, ProgramId, ProgramOutput},
 };
@@ -63,14 +62,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, Identifier, SharedSecretKey)>,
-    private_account_nsks: Vec<NullifierSecretKey>,
-    private_account_membership_proofs: Vec<Option<MembershipProof>>,
+    account_identities: Vec<InputAccountIdentity>,
     program_with_dependencies: &ProgramWithDependencies,
 ) -> Result<(PrivacyPreservingCircuitOutput, Proof), NssaError> {
     let ProgramWithDependencies {
@@ -128,10 +123,7 @@ pub fn execute_and_prove(
 
     let circuit_input = PrivacyPreservingCircuitInput {
         program_outputs,
-        visibility_mask,
-        private_account_keys,
-        private_account_nsks,
-        private_account_membership_proofs,
+        account_identities,
         program_id: program_with_dependencies.program.id(),
     };
 
@@ -240,10 +232,14 @@ 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(), 0, shared_secret)],
-            vec![],
-            vec![None],
+            vec![
+                InputAccountIdentity::Public,
+                InputAccountIdentity::PrivateUnauthorized {
+                    npk: recipient_keys.npk(),
+                    ssk: shared_secret,
+                    identifier: 0,
+                },
+            ],
             &Program::authenticated_transfer_program().into(),
         )
         .unwrap();
@@ -333,13 +329,21 @@ 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(), 0, shared_secret_1),
-                (recipient_keys.npk(), 0, shared_secret_2),
+                InputAccountIdentity::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"),
+                    identifier: 0,
+                },
+                InputAccountIdentity::PrivateUnauthorized {
+                    npk: recipient_keys.npk(),
+                    ssk: shared_secret_2,
+                    identifier: 0,
+                },
             ],
-            vec![sender_keys.nsk],
-            vec![commitment_set.get_proof_for(&commitment_sender), None],
             &program.into(),
         )
         .unwrap();
@@ -402,10 +406,11 @@ mod tests {
         let result = execute_and_prove(
             vec![pre],
             instruction,
-            vec![2],
-            vec![(account_keys.npk(), 0, shared_secret)],
-            vec![],
-            vec![None],
+            vec![InputAccountIdentity::PrivateUnauthorized {
+                npk: account_keys.npk(),
+                ssk: shared_secret,
+                identifier: 0,
+            }],
             &program_with_deps,
         );
 

program_methods/guest/src/bin/privacy_preserving_circuit.rs

@@ -5,7 +5,7 @@ use std::{
 
 use nssa_core::{
     Commitment, CommitmentSetDigest, DUMMY_COMMITMENT_HASH, EncryptionScheme, Identifier,
-    MembershipProof, Nullifier, NullifierPublicKey, NullifierSecretKey,
+    InputAccountIdentity, MembershipProof, Nullifier, NullifierPublicKey, NullifierSecretKey,
     PrivacyPreservingCircuitInput, PrivacyPreservingCircuitOutput, SharedSecretKey,
     account::{Account, AccountId, AccountWithMetadata, Nonce},
     compute_digest_for_path,
@@ -17,7 +17,7 @@ use nssa_core::{
 };
 use risc0_zkvm::{guest::env, serde::to_vec};
 
-const PRIVATE_PDA_FIXED_IDENTIFIER: u128 = u128::MAX;
+const PRIVATE_PDA_FIXED_IDENTIFIER: Identifier = u128::MAX;
 
 /// State of the involved accounts before and after program execution.
 struct ExecutionState {
@@ -25,16 +25,16 @@ struct ExecutionState {
     post_states: HashMap<AccountId, Account>,
     block_validity_window: BlockValidityWindow,
     timestamp_validity_window: TimestampValidityWindow,
-    /// Positions (in `pre_states`) of mask-3 accounts whose supplied npk has been bound to
-    /// their `AccountId` via a proven `AccountId::for_private_pda(program_id, seed, npk)`
+    /// Positions (in `pre_states`) of private-PDA accounts whose supplied npk has been bound
+    /// to their `AccountId` via a proven `AccountId::for_private_pda(program_id, seed, npk)`
     /// check.
     /// Two proof paths populate this set: a `Claim::Pda(seed)` in a program's `post_state` on
     /// that `pre_state`, or a caller's `ChainedCall.pda_seeds` entry matching that `pre_state`
     /// under the private derivation. Binding is an idempotent property, not an event: the same
     /// position can legitimately be bound through both paths in the same tx (e.g. a program
     /// claims a private PDA and then delegates it to a callee), and the set uses `contains`,
-    /// not `assert!(insert)`. After the main loop, every mask-3 position must appear in this
-    /// set; otherwise the npk is unbound and the circuit rejects.
+    /// not `assert!(insert)`. After the main loop, every private-PDA position must appear in
+    /// this set; otherwise the npk is unbound and the circuit rejects.
     private_pda_bound_positions: HashSet<usize>,
     /// Across the whole transaction, each `(program_id, seed)` pair may resolve to at most one
     /// `AccountId`. A seed under a program can derive a family of accounts, one public PDA and
@@ -45,39 +45,29 @@ 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 `account_identities` to the npk that
+    /// variant supplies for that position. Populated once in `derive_from_outputs` by walking
+    /// `account_identities` 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, Identifier, SharedSecretKey)],
+        account_identities: &[InputAccountIdentity],
         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
+        // `account_identities`. The 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_identity) in account_identities.iter().enumerate() {
+            if let Some(npk) = account_identity.npk_if_private_pda() {
+                private_pda_npk_by_position.insert(pos, npk);
             }
         }
 
@@ -194,7 +184,7 @@ impl ExecutionState {
             }
 
             execution_state.validate_and_sync_states(
-                visibility_mask,
+                account_identities,
                 chained_call.program_id,
                 caller_program_id,
                 &chained_call.pda_seeds,
@@ -211,12 +201,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_identity) in account_identities.iter().enumerate() {
+            if account_identity.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"
@@ -251,7 +241,7 @@ impl ExecutionState {
     /// Validate program pre and post states and populate the execution state.
     fn validate_and_sync_states(
         &mut self,
-        visibility_mask: &[u8],
+        account_identities: &[InputAccountIdentity],
         program_id: ProgramId,
         caller_program_id: Option<ProgramId>,
         caller_pda_seeds: &[PdaSeed],
@@ -329,9 +319,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_identity = &account_identities[pre_state_position];
+                if account_identity.is_public() {
+                    match claim {
                         Claim::Authorized => {
                             // Note: no need to check authorized pdas because we have already
                             // checked consistency of authorization above.
@@ -353,40 +343,40 @@ 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_identity.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}"
+                                .expect(
+                                    "private PDA pre_state must have an npk in the position map",
                                 );
-                                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: 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;
@@ -488,10 +478,7 @@ fn resolve_authorization_and_record_bindings(
 
 fn compute_circuit_output(
     execution_state: ExecutionState,
-    visibility_mask: &[u8],
-    private_account_keys: &[(NullifierPublicKey, Identifier, SharedSecretKey)],
-    private_account_nsks: &[NullifierSecretKey],
-    private_account_membership_proofs: &[Option<MembershipProof>],
+    account_identities: &[InputAccountIdentity],
 ) -> PrivacyPreservingCircuitOutput {
     let mut output = PrivacyPreservingCircuitOutput {
         public_pre_states: Vec::new(),
@@ -505,290 +492,268 @@ fn compute_circuit_output(
 
     let states_iter = execution_state.into_states_iter();
     assert_eq!(
-        visibility_mask.len(),
+        account_identities.len(),
         states_iter.len(),
-        "Invalid visibility mask length"
+        "Invalid account_identities 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_identity, (pre_state, post_state)) in account_identities.iter().zip(states_iter) {
+        match account_identity {
+            InputAccountIdentity::Public => {
                 output.public_pre_states.push(pre_state);
                 output.public_post_states.push(post_state);
             }
-            1 | 2 => {
-                let Some((npk, identifier, shared_secret)) = private_keys_iter.next() else {
-                    panic!("Missing private account key");
-                };
+            InputAccountIdentity::PrivateAuthorizedInit {
+                ssk,
+                nsk,
+                identifier,
+            } => {
                 assert_ne!(
                     *identifier, PRIVATE_PDA_FIXED_IDENTIFIER,
                     "Identifier must be different from {PRIVATE_PDA_FIXED_IDENTIFIER}. This is reserved for private PDA."
                 );
+                let npk = NullifierPublicKey::from(nsk);
+                let account_id = AccountId::from((&npk, *identifier));
 
+                assert_eq!(account_id, 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(&account_id),
+                    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,
+                    &account_id,
+                    *identifier,
+                    ssk,
+                    new_nullifier,
+                    new_nonce,
+                );
+            }
+            InputAccountIdentity::PrivateAuthorizedUpdate {
+                ssk,
+                nsk,
+                membership_proof,
+                identifier,
+            } => {
+                assert_ne!(
+                    *identifier, PRIVATE_PDA_FIXED_IDENTIFIER,
+                    "Identifier must be different from {PRIVATE_PDA_FIXED_IDENTIFIER}. This is reserved for private PDA."
+                );
+                let npk = NullifierPublicKey::from(nsk);
+                let account_id = AccountId::from((&npk, *identifier));
+
+                assert_eq!(account_id, 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,
+                    &account_id,
+                    nsk,
+                );
+                let new_nonce = pre_state.account.nonce.private_account_nonce_increment(nsk);
+
+                emit_private_output(
+                    &mut output,
+                    &mut output_index,
+                    post_state,
+                    &account_id,
+                    *identifier,
+                    ssk,
+                    new_nullifier,
+                    new_nonce,
+                );
+            }
+            InputAccountIdentity::PrivateUnauthorized {
+                npk,
+                ssk,
+                identifier,
+            } => {
+                assert_ne!(
+                    *identifier, PRIVATE_PDA_FIXED_IDENTIFIER,
+                    "Identifier must be different from {PRIVATE_PDA_FIXED_IDENTIFIER}. This is reserved for private PDA."
+                );
                 let account_id = AccountId::from((npk, *identifier));
 
                 assert_eq!(account_id, 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,
-                        &account_id,
-                        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(&account_id);
-
-                    let new_nonce = Nonce::private_account_nonce_init(&account_id);
-
-                    ((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(&account_id, &post_with_updated_nonce);
-
-                // Encrypt and push post state
-                let encrypted_account = EncryptionScheme::encrypt(
-                    &post_with_updated_nonce,
-                    *identifier,
-                    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"));
-            }
-            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, identifier, shared_secret)) = private_keys_iter.next() else {
-                    panic!("Missing private account key");
-                };
-
                 assert_eq!(
-                    *identifier, PRIVATE_PDA_FIXED_IDENTIFIER,
-                    "Identifier for private PDAs must be {PRIVATE_PDA_FIXED_IDENTIFIER}."
+                    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 (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,
-                        &pre_state.account_id,
-                        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(&pre_state.account_id);
-                    let new_nonce = Nonce::private_account_nonce_init(&pre_state.account_id);
-                    ((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(&pre_state.account_id, &post_with_updated_nonce);
-
-                let encrypted_account = EncryptionScheme::encrypt(
-                    &post_with_updated_nonce,
-                    PRIVATE_PDA_FIXED_IDENTIFIER,
-                    shared_secret,
-                    &commitment_post,
-                    output_index,
+                let new_nullifier = (
+                    Nullifier::for_account_initialization(&account_id),
+                    DUMMY_COMMITMENT_HASH,
                 );
+                let new_nonce = Nonce::private_account_nonce_init(&account_id);
 
-                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"));
+                emit_private_output(
+                    &mut output,
+                    &mut output_index,
+                    post_state,
+                    &account_id,
+                    *identifier,
+                    ssk,
+                    new_nullifier,
+                    new_nonce,
+                );
+            }
+            InputAccountIdentity::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. The supplied npk on
+                // the variant has been recorded into `private_pda_npk_by_position` and used
+                // for the binding check; we use `pre_state.account_id` directly for nullifier
+                // and commitment derivation.
+                assert!(
+                    !pre_state.is_authorized,
+                    "PrivatePdaInit requires unauthorized pre_state"
+                );
+                assert_eq!(
+                    pre_state.account,
+                    Account::default(),
+                    "New private PDA must be default"
+                );
+
+                let new_nullifier = (
+                    Nullifier::for_account_initialization(&pre_state.account_id),
+                    DUMMY_COMMITMENT_HASH,
+                );
+                let new_nonce = Nonce::private_account_nonce_init(&pre_state.account_id);
+
+                let account_id = pre_state.account_id;
+                emit_private_output(
+                    &mut output,
+                    &mut output_index,
+                    post_state,
+                    &account_id,
+                    PRIVATE_PDA_FIXED_IDENTIFIER,
+                    ssk,
+                    new_nullifier,
+                    new_nonce,
+                );
+            }
+            InputAccountIdentity::PrivatePdaUpdate {
+                ssk,
+                nsk,
+                membership_proof,
+            } => {
+                // 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,
+                    &pre_state.account_id,
+                    nsk,
+                );
+                let new_nonce = pre_state.account.nonce.private_account_nonce_increment(nsk);
+
+                let account_id = pre_state.account_id;
+                emit_private_output(
+                    &mut output,
+                    &mut output_index,
+                    post_state,
+                    &account_id,
+                    PRIVATE_PDA_FIXED_IDENTIFIER,
+                    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>,
+#[expect(
+    clippy::too_many_arguments,
+    reason = "All seven inputs are distinct concerns from the variant arms; bundling would be artificial"
+)]
+fn emit_private_output(
+    output: &mut PrivacyPreservingCircuitOutput,
+    output_index: &mut u32,
+    post_state: Account,
+    account_id: &AccountId,
+    identifier: Identifier,
+    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(account_id, &post_with_updated_nonce);
+    let encrypted_account = EncryptionScheme::encrypt(
+        &post_with_updated_nonce,
+        identifier,
+        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,
     account_id: &AccountId,
     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(account_id);
-            (nullifier, DUMMY_COMMITMENT_HASH)
-        },
-        |membership_proof| {
-            // Compute commitment set digest associated with provided auth path
-            let commitment_pre = Commitment::new(account_id, 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(account_id, 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,
+        account_identities,
         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(&account_identities, 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, &account_identities);
 
     env::commit(&output);
 }

sequencer/core/src/lib.rs

@@ -1076,7 +1076,7 @@ mod tests {
             program::Program,
         };
         use nssa_core::{
-            SharedSecretKey,
+            InputAccountIdentity, SharedSecretKey,
             account::AccountWithMetadata,
             encryption::{EphemeralPublicKey, EphemeralSecretKey, ViewingPublicKey},
         };
@@ -1114,10 +1114,11 @@ mod tests {
                 (&npk, 0),
             )],
             Program::serialize_instruction(0_u128).unwrap(),
-            vec![1],
-            vec![(npk, 0, shared_secret)],
-            vec![nsk],
-            vec![None],
+            vec![InputAccountIdentity::PrivateAuthorizedInit {
+                ssk: shared_secret,
+                nsk,
+                identifier: 0,
+            }],
             &Program::authenticated_transfer_program().into(),
         )
         .unwrap();

wallet/src/lib.rs

@@ -413,13 +413,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.identifier, keys.ssk))
-                .collect::<Vec<_>>(),
-            acc_manager.private_account_auth(),
-            acc_manager.private_account_membership_proofs(),
+            acc_manager.account_identities(),
             &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::{
-    Identifier, MembershipProof, NullifierPublicKey, NullifierSecretKey, SharedSecretKey,
+    Identifier, InputAccountIdentity, MembershipProof, NullifierPublicKey, NullifierSecretKey,
+    SharedSecretKey,
     account::{AccountWithMetadata, Nonce},
     encryption::{EphemeralPublicKey, ViewingPublicKey},
 };
@@ -34,7 +35,7 @@ impl PrivacyPreservingAccount {
                 | Self::PrivateForeign {
                     npk: _,
                     vpk: _,
-                    identifier: _
+                    identifier: _,
                 }
         )
     }
@@ -42,7 +43,6 @@ impl PrivacyPreservingAccount {
 
 pub struct PrivateAccountKeys {
     pub npk: NullifierPublicKey,
-    pub identifier: Identifier,
     pub ssk: SharedSecretKey,
     pub vpk: ViewingPublicKey,
     pub epk: EphemeralPublicKey,
@@ -58,7 +58,6 @@ enum State {
 
 pub struct AccountManager {
     states: Vec<State>,
-    visibility_mask: Vec<u8>,
 }
 
 impl AccountManager {
@@ -66,11 +65,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)
@@ -80,13 +78,12 @@ 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,
@@ -95,6 +92,9 @@ impl AccountManager {
                 } => {
                     let acc = nssa_core::account::Account::default();
                     let auth_acc = AccountWithMetadata::new(acc, false, (&npk, identifier));
+                    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,
@@ -102,20 +102,18 @@ impl AccountManager {
                         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> {
@@ -128,10 +126,6 @@ impl AccountManager {
             .collect()
     }
 
-    pub fn visibility_mask(&self) -> &[u8] {
-        &self.visibility_mask
-    }
-
     pub fn public_account_nonces(&self) -> Vec<Nonce> {
         self.states
             .iter()
@@ -146,38 +140,46 @@ 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 {
-                        npk: pre.npk,
+    /// 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 account_identities(&self) -> Vec<InputAccountIdentity> {
+        self.states
+            .iter()
+            .map(|state| match state {
+                State::Public { .. } => InputAccountIdentity::Public,
+                State::Private(pre) => match (pre.nsk, pre.proof.clone()) {
+                    (Some(nsk), Some(membership_proof)) => {
+                        InputAccountIdentity::PrivateAuthorizedUpdate {
+                            ssk: pre.ssk,
+                            nsk,
+                            membership_proof,
+                            identifier: pre.identifier,
+                        }
+                    }
+                    (Some(nsk), None) => InputAccountIdentity::PrivateAuthorizedInit {
+                        ssk: pre.ssk,
+                        nsk,
                         identifier: pre.identifier,
-                        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,
+                    },
+                    (None, _) => InputAccountIdentity::PrivateUnauthorized {
+                        npk: pre.npk,
+                        ssk: pre.ssk,
+                        identifier: pre.identifier,
+                    },
+                },
             })
             .collect()
     }
@@ -210,6 +212,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 (`account_identities()`) 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(
@@ -237,6 +246,10 @@ async fn private_acc_preparation(
     // support from that in the wallet.
     let sender_pre = AccountWithMetadata::new(from_acc.clone(), true, (&from_npk, from_identifier));
 
+    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,
@@ -244,5 +257,7 @@ async fn private_acc_preparation(
         vpk: from_vpk,
         pre_state: sender_pre,
         proof,
+        ssk,
+        epk,
     })
 }