add identifier logic and fix unit tests

integration_tests/tests/tps.rs

@@ -256,7 +256,7 @@ fn build_privacy_transaction() -> PrivacyPreservingTransaction {
         vec![1, 2],
         vec![(sender_npk, sender_ss), (recipient_npk, recipient_ss)],
         vec![sender_nsk],
-        vec![], // TODO: marvin fix
+        vec![],
         vec![Some(proof)],
         &program.into(),
     )

key_protocol/src/key_management/key_tree/keys_private.rs

@@ -1,5 +1,5 @@
 use k256::{Scalar, elliptic_curve::PrimeField as _};
-use nssa_core::{NullifierPublicKey, encryption::ViewingPublicKey};
+use nssa_core::{NullifierPublicKey, account::Identifier, encryption::ViewingPublicKey};
 use serde::{Deserialize, Serialize};
 
 use crate::key_management::{
@@ -111,7 +111,7 @@ impl KeyNode for ChildKeysPrivate {
     }
 
     fn account_id(&self) -> nssa::AccountId {
-        nssa::AccountId::private_account_id(&self.value.0.nullifier_public_key, None)
+        nssa::AccountId::private_account_id(&self.value.0.nullifier_public_key, Identifier(0_u128))
     }
 }
 

key_protocol/src/key_management/key_tree/keys_public.rs

@@ -99,7 +99,7 @@ impl KeyNode for ChildKeysPublic {
     }
 
     fn account_id(&self) -> nssa::AccountId {
-        nssa::AccountId::public_account_id(&self.cpk, None)
+        nssa::AccountId::public_account_id(&self.cpk)
     }
 }
 

key_protocol/src/key_management/mod.rs

@@ -129,7 +129,7 @@ mod tests {
 
         let public_key = nssa_core::PublicKey::new_from_private_key(&pub_account_signing_key);
 
-        let account = nssa::AccountId::public_account_id(&public_key, None);
+        let account = nssa::AccountId::public_account_id(&public_key);
 
         println!("======Prerequisites======");
         println!();

key_protocol/src/key_protocol_core/mod.rs

@@ -2,6 +2,7 @@ use std::collections::BTreeMap;
 
 use anyhow::Result;
 use k256::AffinePoint;
+use nssa_core::account::Identifier;
 use serde::{Deserialize, Serialize};
 
 use crate::key_management::{
@@ -33,7 +34,6 @@ impl NSSAUserData {
         for (account_id, key) in accounts_keys_map {
             let expected_account_id = nssa::AccountId::public_account_id(
                 &nssa_core::PublicKey::new_from_private_key(key),
-                None,
             );
             if &expected_account_id != account_id {
                 println!("{expected_account_id}, {account_id}");
@@ -49,7 +49,7 @@ impl NSSAUserData {
         let mut check_res = true;
         for (account_id, (key, _)) in accounts_keys_map {
             let expected_account_id =
-                nssa::AccountId::private_account_id(&key.nullifier_public_key, None);
+                nssa::AccountId::private_account_id(&key.nullifier_public_key, Identifier(0_u128));
             if expected_account_id != *account_id {
                 println!("{expected_account_id}, {account_id}");
                 check_res = false;

nssa/Cargo.toml

@@ -8,7 +8,7 @@ license = { workspace = true }
 workspace = true
 
 [dependencies]
-nssa_core = { workspace = true, features = ["host"] }
+nssa_core = { workspace = true, features = ["host", "test-utils"] }
 
 anyhow.workspace = true
 thiserror.workspace = true

nssa/core/Cargo.toml

@@ -27,3 +27,4 @@ serde_json.workspace = true
 [features]
 default = ["dep:k256"]
 host = ["dep:k256"]
+test-utils = ["dep:k256"]

nssa/core/src/account.rs

@@ -10,12 +10,16 @@ use risc0_zkvm::sha::{Impl, Sha256 as _};
 use serde::{Deserialize, Serialize};
 use serde_with::{DeserializeFromStr, SerializeDisplay};
 
-use crate::{NullifierPublicKey, NullifierSecretKey, PublicKey, program::ProgramId};
+use crate::{
+    EphemeralPublicKey, NullifierPublicKey, NullifierSecretKey, PublicKey, program::ProgramId,
+};
 
 pub mod data;
 
 #[derive(Copy, Debug, Default, Clone, Eq, PartialEq)]
 pub struct Nonce(pub u128);
+#[derive(Copy, Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
+pub struct Identifier(pub u128);
 
 impl Nonce {
     pub const fn public_account_nonce_increment(&mut self) {
@@ -91,6 +95,23 @@ impl BorshDeserialize for Nonce {
 
 pub type Balance = u128;
 
+impl Identifier {
+    #[must_use]
+    pub fn private_identifier(epk: &EphemeralPublicKey, index: u8) -> Self {
+        const PRIVATE_ACCOUNT_ID_PREFIX: &[u8; 32] = b"/LEE/v0.3/AccountId/Identifier/\x00";
+
+        let mut bytes = Vec::<u8>::new();
+        bytes.extend_from_slice(PRIVATE_ACCOUNT_ID_PREFIX);
+        bytes.extend_from_slice(&epk.0);
+        bytes.extend_from_slice(&[index]);
+
+        let mut value = [0_u8; 16];
+        value.copy_from_slice(&Impl::hash_bytes(&bytes).as_bytes()[0..16]);
+
+        Self(u128::from_le_bytes(value))
+    }
+}
+
 /// Account to be used both in public and private contexts.
 #[derive(
     Default, Clone, Eq, PartialEq, Serialize, Deserialize, BorshSerialize, BorshDeserialize,
@@ -180,18 +201,14 @@ impl AccountId {
     }
 
     #[must_use]
-    pub fn private_account_id(value: &NullifierPublicKey, identifier: Option<u128>) -> Self {
+    pub fn private_account_id(value: &NullifierPublicKey, identifier: Identifier) -> Self {
         const PRIVATE_ACCOUNT_ID_PREFIX: &[u8; 32] =
             b"/LEE/v0.3/AccountId/Private/\x00\x00\x00\x00";
 
         let mut bytes = Vec::<u8>::new();
         bytes.extend_from_slice(PRIVATE_ACCOUNT_ID_PREFIX);
         bytes.extend_from_slice(&value.0);
-
+        bytes.extend_from_slice(&identifier.0.to_le_bytes());
-        match identifier {
-            None => {}
-            Some(identifier) => bytes.extend_from_slice(&identifier.to_le_bytes()),
-        }
 
         Self::new(
             Impl::hash_bytes(&bytes)
@@ -202,18 +219,14 @@ impl AccountId {
     }
 
     #[must_use]
-    pub fn public_account_id(value: &PublicKey, identifier: Option<u128>) -> Self {
+    pub fn public_account_id(value: &PublicKey) -> Self {
         const PUBLIC_ACCOUNT_ID_PREFIX: &[u8; 32] =
             b"/LEE/v0.3/AccountId/Public/\x00\x00\x00\x00\x00";
 
         let mut bytes = Vec::<u8>::new();
         bytes.extend_from_slice(PUBLIC_ACCOUNT_ID_PREFIX);
         bytes.extend_from_slice(value.value());
-
+        // bytes.extend_from_slice(&identifier.0.to_le_bytes());
-        match identifier {
-            None => {}
-            Some(identifier) => bytes.extend_from_slice(&identifier.to_le_bytes()),
-        }
 
         Self::new(
             Impl::hash_bytes(&bytes)
@@ -409,7 +422,9 @@ mod tests {
             126, 85, 106, 222, 127, 193, 125, 168, 62, 150, 129, 194, 135, 114,
         ]);
 
-        let account_id = AccountId::private_account_id(&npk, Some(13_u128));
+        let identifier = Identifier(13_u128);
+
+        let account_id = AccountId::private_account_id(&npk, identifier);
 
         assert_eq!(account_id, expected_account_id);
     }
@@ -419,12 +434,27 @@ mod tests {
         let pub_key = PublicKey::try_new([42_u8; 32]).expect("Expect valid Public Key");
 
         let expected_account_id = AccountId::new([
-            75, 60, 223, 47, 170, 89, 187, 173, 89, 16, 96, 18, 76, 101, 203, 128, 241, 4, 253, 18,
+            55, 223, 166, 27, 166, 126, 71, 128, 222, 225, 215, 176, 98, 21, 215, 13, 71, 74, 13,
-            61, 201, 37, 226, 199, 119, 9, 1, 239, 131, 221, 142,
+            72, 200, 175, 25, 19, 96, 160, 250, 230, 45, 15, 254, 134,
         ]);
 
-        let account_id = AccountId::public_account_id(&pub_key, Some(13_u128));
+        let account_id = AccountId::public_account_id(&pub_key);
 
         assert_eq!(account_id, expected_account_id);
     }
+
+    #[test]
+    fn identifier_from_ephemeral_public_key() {
+        let epk = EphemeralPublicKey::from_scalar([
+            185, 147, 32, 242, 145, 91, 123, 77, 42, 33, 134, 84, 12, 165, 117, 70, 158, 201, 95,
+            153, 14, 12, 92, 235, 128, 156, 194, 169, 68, 35, 165, 127,
+        ]);
+
+        let expected_identifier = Identifier(u128::from_le_bytes([
+            170, 216, 75, 182, 85, 117, 119, 230, 115, 121, 70, 204, 104, 96, 182, 122,
+        ]));
+        let identifier = Identifier::private_identifier(&epk, 13_u8);
+
+        assert_eq!(identifier, expected_identifier);
+    }
 }

nssa/core/src/circuit_io.rs

@@ -3,7 +3,7 @@ use serde::{Deserialize, Serialize};
 use crate::{
     Commitment, CommitmentSetDigest, MembershipProof, Nullifier, NullifierPublicKey,
     NullifierSecretKey, SharedSecretKey,
-    account::{Account, AccountWithMetadata},
+    account::{Account, AccountWithMetadata, Identifier},
     encryption::Ciphertext,
     program::{BlockValidityWindow, ProgramId, ProgramOutput, TimestampValidityWindow},
 };
@@ -22,8 +22,8 @@ pub struct PrivacyPreservingCircuitInput {
     pub private_account_keys: Vec<(NullifierPublicKey, SharedSecretKey)>,
     /// Nullifier secret keys for authorized private accounts.
     pub private_account_nsks: Vec<NullifierSecretKey>,
-    /// Identifiers used to generate `AccountId` (TODO: marvin double check logic).
+    /// Identifiers used to generate `AccountId`.
-    pub private_account_identifiers: Vec<u128>,
+    pub private_account_identifiers: Vec<Identifier>,
     /// Membership proofs for private accounts. Can be [`None`] for uninitialized accounts.
     pub private_account_membership_proofs: Vec<Option<MembershipProof>>,
     /// Program ID.
@@ -59,7 +59,7 @@ mod tests {
     use super::*;
     use crate::{
         Commitment, Nullifier, NullifierPublicKey,
-        account::{Account, AccountId, AccountWithMetadata, Nonce},
+        account::{Account, AccountId, AccountWithMetadata, Identifier, Nonce},
     };
 
     #[test]
@@ -95,7 +95,10 @@ mod tests {
             }],
             ciphertexts: vec![Ciphertext(vec![255, 255, 1, 1, 2, 2])],
             new_commitments: vec![Commitment::new(
-                &AccountId::private_account_id(&NullifierPublicKey::from(&[1_u8; 32]), None),
+                &AccountId::private_account_id(
+                    &NullifierPublicKey::from(&[1_u8; 32]),
+                    Identifier(0_u128),
+                ),
                 &Account::default(),
             )],
             new_nullifiers: vec![(
@@ -103,7 +106,7 @@ mod tests {
                     &Commitment::new(
                         &AccountId::private_account_id(
                             &NullifierPublicKey::from(&[2_u8; 32]),
-                            None,
+                            Identifier(0_u128),
                         ),
                         &Account::default(),
                     ),

nssa/core/src/encryption/mod.rs

@@ -5,11 +5,11 @@ use chacha20::{
 };
 use risc0_zkvm::sha::{Impl, Sha256 as _};
 use serde::{Deserialize, Serialize};
+pub use shared_key_derivation::EphemeralPublicKey;
 #[cfg(feature = "host")]
-pub use shared_key_derivation::{EphemeralPublicKey, EphemeralSecretKey, ViewingPublicKey};
+pub use shared_key_derivation::{EphemeralSecretKey, ViewingPublicKey};
 
 use crate::{Commitment, account::Account};
-#[cfg(feature = "host")]
 pub mod shared_key_derivation;
 
 pub type Scalar = [u8; 32];

nssa/core/src/lib.rs

@@ -8,7 +8,9 @@ pub use commitment::{
     Commitment, CommitmentSetDigest, DUMMY_COMMITMENT, DUMMY_COMMITMENT_HASH, MembershipProof,
     compute_digest_for_path,
 };
-pub use encryption::{EncryptionScheme, SharedSecretKey};
+pub use encryption::{
+    EncryptionScheme, SharedSecretKey, shared_key_derivation::EphemeralPublicKey,
+};
 pub use nullifier::{Nullifier, NullifierPublicKey, NullifierSecretKey};
 pub use signature::{PrivateKey, PublicKey, Signature};
 

nssa/core/src/nullifier.rs

@@ -79,6 +79,7 @@ impl Nullifier {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::account::Identifier;
 
     #[test]
     fn constructor_for_account_update() {
@@ -99,7 +100,9 @@ mod tests {
             255, 29, 105, 42, 186, 43, 11, 157, 168, 132, 225, 17, 163,
         ]);
 
-        let account_id = AccountId::private_account_id(&npk, Some(0_u128));
+        let identifier = Identifier(0_u128);
+
+        let account_id = AccountId::private_account_id(&npk, identifier);
 
         let expected_nullifier = Nullifier([
             63, 58, 51, 159, 15, 100, 240, 243, 60, 143, 151, 108, 116, 144, 101, 6, 134, 72, 198,

nssa/core/src/signature/mod.rs

@@ -72,6 +72,11 @@ impl Signature {
 
         pk.verify_raw(bytes, &sig).is_ok()
     }
+
+    #[cfg(feature = "test-utils")]
+    pub fn new_for_tests(value: [u8; 64]) -> Self {
+        Self { value }
+    }
 }
 
 #[cfg(test)]
@@ -82,12 +87,6 @@ mod tests {
 
     use crate::{Signature, signature::bip340_test_vectors};
 
-    impl Signature {
-        pub(crate) fn _new_for_tests(value: [u8; 64]) -> Self {
-            Self { value }
-        }
-    }
-
     #[test]
     fn signature_generation_from_bip340_test_vectors() {
         for (i, test_vector) in bip340_test_vectors::test_vectors().into_iter().enumerate() {

nssa/src/privacy_preserving_transaction/circuit.rs

@@ -4,7 +4,7 @@ use borsh::{BorshDeserialize, BorshSerialize};
 use nssa_core::{
     MembershipProof, NullifierPublicKey, NullifierSecretKey, PrivacyPreservingCircuitInput,
     PrivacyPreservingCircuitOutput, SharedSecretKey,
-    account::AccountWithMetadata,
+    account::{AccountWithMetadata, Identifier},
     program::{ChainedCall, InstructionData, ProgramId, ProgramOutput},
 };
 use risc0_zkvm::{ExecutorEnv, InnerReceipt, ProverOpts, Receipt, default_prover};
@@ -70,7 +70,7 @@ pub fn execute_and_prove(
     visibility_mask: Vec<u8>,
     private_account_keys: Vec<(NullifierPublicKey, SharedSecretKey)>,
     private_account_nsks: Vec<NullifierSecretKey>,
-    private_account_identifiers: Vec<u128>,
+    private_account_identifiers: Vec<Identifier>,
     private_account_membership_proofs: Vec<Option<MembershipProof>>,
     program_with_dependencies: &ProgramWithDependencies,
 ) -> Result<(PrivacyPreservingCircuitOutput, Proof), NssaError> {
@@ -177,7 +177,7 @@ mod tests {
 
     use nssa_core::{
         Commitment, DUMMY_COMMITMENT_HASH, EncryptionScheme, Nullifier, SharedSecretKey,
-        account::{Account, AccountId, AccountWithMetadata, Nonce, data::Data},
+        account::{Account, AccountId, AccountWithMetadata, Identifier, Nonce, data::Data},
     };
 
     use super::*;
@@ -208,7 +208,7 @@ mod tests {
         let recipient = AccountWithMetadata::new(
             Account::default(),
             false,
-            AccountId::private_account_id(&recipient_keys.npk(), None),
+            AccountId::private_account_id(&recipient_keys.npk(), Identifier(0_u128)),
         );
 
         let balance_to_move: u128 = 37;
@@ -238,7 +238,7 @@ mod tests {
             vec![0, 2],
             vec![(recipient_keys.npk(), shared_secret)],
             vec![],
-            vec![], // TODO check (Marvin)
+            vec![],
             vec![None],
             &Program::authenticated_transfer_program().into(),
         )
@@ -270,7 +270,7 @@ mod tests {
         let sender_keys = test_private_account_keys_1();
         let recipient_keys = test_private_account_keys_2();
         let recipient_id =
-            AccountId::private_account_id(&test_private_account_keys_2().npk(), None);
+            AccountId::private_account_id(&test_private_account_keys_2().npk(), Identifier(0_u128));
 
         let sender_nonce = Nonce(0xdead_beef);
         let sender_pre = AccountWithMetadata::new(
@@ -281,14 +281,14 @@ mod tests {
                 data: Data::default(),
             },
             true,
-            AccountId::private_account_id(&sender_keys.npk(), None),
+            AccountId::private_account_id(&sender_keys.npk(), Identifier(0_u128)),
         );
         let commitment_sender = Commitment::new(&sender_pre.account_id, &sender_pre.account);
 
         let recipient = AccountWithMetadata::new(
             Account::default(),
             false,
-            AccountId::private_account_id(&recipient_keys.npk(), None),
+            AccountId::private_account_id(&recipient_keys.npk(), Identifier(0_u128)),
         );
         let balance_to_move: u128 = 37;
 
@@ -340,7 +340,7 @@ mod tests {
                 (recipient_keys.npk(), shared_secret_2),
             ],
             vec![sender_keys.nsk],
-            vec![], // TODO check (Marvin)
+            vec![],
             vec![commitment_set.get_proof_for(&commitment_sender), None],
             &program.into(),
         )
@@ -378,7 +378,7 @@ mod tests {
         let pre = AccountWithMetadata::new(
             Account::default(),
             false,
-            AccountId::private_account_id(&account_keys.npk(), None),
+            AccountId::private_account_id(&account_keys.npk(), Identifier(0_u128)),
         );
 
         let validity_window_chain_caller = Program::validity_window_chain_caller();
@@ -407,7 +407,7 @@ mod tests {
             vec![2],
             vec![(account_keys.npk(), shared_secret)],
             vec![],
-            vec![], // TODO check (Marvin)
+            vec![],
             vec![None],
             &program_with_deps,
         );

nssa/src/privacy_preserving_transaction/message.rs

@@ -130,7 +130,7 @@ impl Message {
 pub mod tests {
     use nssa_core::{
         Commitment, EncryptionScheme, Nullifier, NullifierPublicKey, SharedSecretKey,
-        account::Account,
+        account::{Account, Identifier},
         encryption::{EphemeralPublicKey, ViewingPublicKey},
         program::{BlockValidityWindow, TimestampValidityWindow},
     };
@@ -152,8 +152,8 @@ pub mod tests {
         let npk1 = NullifierPublicKey::from(&nsk1);
         let npk2 = NullifierPublicKey::from(&nsk2);
 
-        let account_id1 = AccountId::private_account_id(&npk1, None);
+        let account_id1 = AccountId::private_account_id(&npk1, Identifier(0_u128));
-        let account_id2 = AccountId::private_account_id(&npk2, None);
+        let account_id2 = AccountId::private_account_id(&npk2, Identifier(0_u128));
         let public_account_ids = vec![account_id1, account_id2];
 
         let nonces = vec![1_u128.into(), 2_u128.into(), 3_u128.into()];
@@ -186,7 +186,7 @@ pub mod tests {
     fn encrypted_account_data_constructor() {
         let npk = NullifierPublicKey::from(&[1; 32]);
         let vpk = ViewingPublicKey::from_scalar([2; 32]);
-        let account_id = AccountId::private_account_id(&npk, None);
+        let account_id = AccountId::private_account_id(&npk, Identifier(0_u128));
         let account = Account::default();
         let commitment = Commitment::new(&account_id, &account);
         let esk = [3; 32];

nssa/src/privacy_preserving_transaction/transaction.rs

@@ -154,7 +154,7 @@ impl PrivacyPreservingTransaction {
         self.witness_set
             .signatures_and_public_keys()
             .iter()
-            .map(|(_, public_key)| AccountId::public_account_id(public_key, None))
+            .map(|(_, public_key)| AccountId::public_account_id(public_key))
             .collect()
     }
 
@@ -214,8 +214,8 @@ mod tests {
     fn keys_for_tests() -> (PrivateKey, PrivateKey, AccountId, AccountId) {
         let key1 = PrivateKey::try_new([1; 32]).unwrap();
         let key2 = PrivateKey::try_new([2; 32]).unwrap();
-        let addr1 = AccountId::public_account_id(&PublicKey::new_from_private_key(&key1), None);
+        let addr1 = AccountId::public_account_id(&PublicKey::new_from_private_key(&key1));
-        let addr2 = AccountId::public_account_id(&PublicKey::new_from_private_key(&key2), None);
+        let addr2 = AccountId::public_account_id(&PublicKey::new_from_private_key(&key2));
         (key1, key2, addr1, addr2)
     }
 

nssa/src/public_transaction/transaction.rs

@@ -45,7 +45,7 @@ impl PublicTransaction {
         self.witness_set
             .signatures_and_public_keys()
             .iter()
-            .map(|(_, public_key)| AccountId::public_account_id(public_key, None))
+            .map(|(_, public_key)| AccountId::public_account_id(public_key))
             .collect()
     }
 
@@ -282,7 +282,7 @@ impl PublicTransaction {
 
 #[cfg(test)]
 pub mod tests {
-    use nssa_core::{PrivateKey, PublicKey};
+    use nssa_core::{PrivateKey, PublicKey, Signature};
     use sha2::{Digest as _, digest::FixedOutput as _};
 
     use crate::{
@@ -295,8 +295,8 @@ pub mod tests {
     fn keys_for_tests() -> (PrivateKey, PrivateKey, AccountId, AccountId) {
         let key1 = PrivateKey::try_new([1; 32]).unwrap();
         let key2 = PrivateKey::try_new([2; 32]).unwrap();
-        let addr1 = AccountId::public_account_id(&PublicKey::new_from_private_key(&key1), None);
+        let addr1 = AccountId::public_account_id(&PublicKey::new_from_private_key(&key1));
-        let addr2 = AccountId::public_account_id(&PublicKey::new_from_private_key(&key2), None);
+        let addr2 = AccountId::public_account_id(&PublicKey::new_from_private_key(&key2));
         (key1, key2, addr1, addr2)
     }
 
@@ -422,27 +422,27 @@ pub mod tests {
         assert!(matches!(result, Err(NssaError::InvalidInput(_))));
     }
 
-    // TODO: fix Marvin (commented out since that test is local to nssa-core)
+    #[cfg(feature = "test-utils")]
-    // #[test]
+    #[test]
-    // fn all_signatures_must_be_valid() {
+    fn all_signatures_must_be_valid() {
-    // let (key1, key2, addr1, addr2) = keys_for_tests();
+        let (key1, key2, addr1, addr2) = keys_for_tests();
-    // let state = state_for_tests();
+        let state = state_for_tests();
-    // let nonces = vec![0_u128.into(), 0_u128.into()];
+        let nonces = vec![0_u128.into(), 0_u128.into()];
-    // let instruction = 1337;
+        let instruction = 1337;
-    // let message = Message::try_new(
+        let message = Message::try_new(
-    // Program::authenticated_transfer_program().id(),
+            Program::authenticated_transfer_program().id(),
-    // vec![addr1, addr2],
+            vec![addr1, addr2],
-    // nonces,
+            nonces,
-    // instruction,
+            instruction,
-    // )
+        )
-    // .unwrap();
+        .unwrap();
-    //
+
-    // let mut witness_set = WitnessSet::for_message(&message, &[&key1, &key2]);
+        let mut witness_set = WitnessSet::for_message(&message, &[&key1, &key2]);
-    // witness_set.signatures_and_public_keys[0].0 = Signature::new_for_tests([1; 64]);
+        witness_set.signatures_and_public_keys[0].0 = Signature::new_for_tests([1; 64]);
-    // let tx = PublicTransaction::new(message, witness_set);
+        let tx = PublicTransaction::new(message, witness_set);
-    // let result = tx.validate_and_produce_public_state_diff(&state, 1, 0);
+        let result = tx.validate_and_produce_public_state_diff(&state, 1, 0);
-    // assert!(matches!(result, Err(NssaError::InvalidInput(_))));
+        assert!(matches!(result, Err(NssaError::InvalidInput(_))));
-    // }
+    }
 
     #[test]
     fn nonces_must_match_the_state_current_nonces() {

nssa/src/public_transaction/witness_set.rs

@@ -66,8 +66,8 @@ mod tests {
         let key2 = PrivateKey::try_new([2; 32]).unwrap();
         let pubkey1 = PublicKey::new_from_private_key(&key1);
         let pubkey2 = PublicKey::new_from_private_key(&key2);
-        let addr1 = AccountId::public_account_id(&pubkey1, None);
+        let addr1 = AccountId::public_account_id(&pubkey1);
-        let addr2 = AccountId::public_account_id(&pubkey2, None);
+        let addr2 = AccountId::public_account_id(&pubkey2);
         let nonces = vec![1_u128.into(), 2_u128.into()];
         let instruction = vec![1, 2, 3, 4];
         let message = Message::try_new([0; 8], vec![addr1, addr2], nonces, instruction).unwrap();

program_methods/guest/src/bin/privacy_preserving_circuit.rs

@@ -7,7 +7,7 @@ use nssa_core::{
     Commitment, CommitmentSetDigest, DUMMY_COMMITMENT_HASH, EncryptionScheme, MembershipProof,
     Nullifier, NullifierPublicKey, NullifierSecretKey, PrivacyPreservingCircuitInput,
     PrivacyPreservingCircuitOutput, SharedSecretKey,
-    account::{Account, AccountId, AccountWithMetadata, Nonce},
+    account::{Account, AccountId, AccountWithMetadata, Identifier, Nonce},
     compute_digest_for_path,
     program::{
         AccountPostState, BlockValidityWindow, ChainedCall, Claim, DEFAULT_PROGRAM_ID,
@@ -295,7 +295,7 @@ fn compute_circuit_output(
     visibility_mask: &[u8],
     private_account_keys: &[(NullifierPublicKey, SharedSecretKey)],
     private_account_nsks: &[NullifierSecretKey],
-    private_account_identifiers: &[u128],
+    private_account_identifiers: &[Identifier],
     private_account_membership_proofs: &[Option<MembershipProof>],
 ) -> PrivacyPreservingCircuitOutput {
     let mut output = PrivacyPreservingCircuitOutput {
@@ -335,15 +335,12 @@ fn compute_circuit_output(
                     panic!("Missing private account key");
                 };
 
-                // TODO: (Marvin) double check
                 let Some(identifier) = private_identifiers_iter.next() else {
                     panic!("Missing private account identifier");
                 };
 
-                // TODO: (Marvin) identifier used here)
-                // Relevant here as this applies for both cases (authenicated and not authenicated).
                 assert_eq!(
-                    AccountId::private_account_id(npk, Some(*identifier)),
+                    AccountId::private_account_id(npk, *identifier),
                     pre_state.account_id,
                     "AccountId mismatch"
                 );
@@ -406,9 +403,7 @@ fn compute_circuit_output(
                         "Membership proof must be None for unauthorized accounts"
                     );
 
-                    // TODO: (Marvin) need to add a Vec<identifier> as input.
+                    let account_id = AccountId::private_account_id(npk, *identifier);
-                    // TODO: use here
-                    let account_id = AccountId::private_account_id(npk, Some(*identifier));
                     let nullifier = Nullifier::for_account_initialization(&account_id);
 
                     let new_nonce = Nonce::private_account_nonce_init(npk);
@@ -466,11 +461,11 @@ fn compute_nullifier_and_set_digest(
     pre_account: &Account,
     npk: &NullifierPublicKey,
     nsk: &NullifierSecretKey,
-    identifier: u128,
+    identifier: Identifier,
 ) -> (Nullifier, CommitmentSetDigest) {
     // TODO: consider rewriting the function to receive account id instead of npk.
     // NOTE: this does not use the identifier at all.
-    let account_id = AccountId::private_account_id(npk, Some(identifier));
+    let account_id = AccountId::private_account_id(npk, identifier);
     membership_proof_opt.as_ref().map_or_else(
         || {
             assert_eq!(

programs/amm/src/tests.rs

@@ -1313,24 +1313,21 @@ impl IdForExeTests {
     }
 
     fn user_token_a_id() -> AccountId {
-        AccountId::public_account_id(
+        AccountId::public_account_id(&PublicKey::new_from_private_key(
-            &PublicKey::new_from_private_key(&PrivateKeysForTests::user_token_a_key()),
+            &PrivateKeysForTests::user_token_a_key(),
-            None,
+        ))
-        )
     }
 
     fn user_token_b_id() -> AccountId {
-        AccountId::public_account_id(
+        AccountId::public_account_id(&PublicKey::new_from_private_key(
-            &PublicKey::new_from_private_key(&PrivateKeysForTests::user_token_b_key()),
+            &PrivateKeysForTests::user_token_b_key(),
-            None,
+        ))
-        )
     }
 
     fn user_token_lp_id() -> AccountId {
-        AccountId::public_account_id(
+        AccountId::public_account_id(&PublicKey::new_from_private_key(
-            &PublicKey::new_from_private_key(&PrivateKeysForTests::user_token_lp_key()),
+            &PrivateKeysForTests::user_token_lp_key(),
-            None,
+        ))
-        )
     }
 
     fn vault_a_id() -> AccountId {

storage/src/indexer/mod.rs

@@ -243,11 +243,11 @@ mod tests {
     }
 
     fn acc1() -> AccountId {
-        AccountId::public_account_id(&PublicKey::new_from_private_key(&acc1_sign_key()), None)
+        AccountId::public_account_id(&PublicKey::new_from_private_key(&acc1_sign_key()))
     }
 
     fn acc2() -> AccountId {
-        AccountId::public_account_id(&PublicKey::new_from_private_key(&acc2_sign_key()), None)
+        AccountId::public_account_id(&PublicKey::new_from_private_key(&acc2_sign_key()))
     }
 
     #[test]

testnet_initial_state/src/lib.rs

@@ -5,7 +5,8 @@ use key_protocol::key_management::{
 };
 use nssa::{Account, AccountId, Data, V03State};
 use nssa_core::{
-    NullifierPublicKey, PrivateKey, PublicKey, encryption::shared_key_derivation::Secp256k1Point,
+    NullifierPublicKey, PrivateKey, PublicKey, account::Identifier,
+    encryption::shared_key_derivation::Secp256k1Point,
 };
 use serde::{Deserialize, Serialize};
 
@@ -110,17 +111,15 @@ pub fn initial_pub_accounts_private_keys() -> Vec<PublicAccountPrivateInitialDat
 
     vec![
         PublicAccountPrivateInitialData {
-            account_id: AccountId::public_account_id(
+            account_id: AccountId::public_account_id(&PublicKey::new_from_private_key(
-                &PublicKey::new_from_private_key(&acc1_pub_sign_key),
+                &acc1_pub_sign_key,
-                None,
+            )),
-            ),
             pub_sign_key: acc1_pub_sign_key,
         },
         PublicAccountPrivateInitialData {
-            account_id: AccountId::public_account_id(
+            account_id: AccountId::public_account_id(&PublicKey::new_from_private_key(
-                &PublicKey::new_from_private_key(&acc2_pub_sign_key),
+                &acc2_pub_sign_key,
-                None,
+            )),
-            ),
             pub_sign_key: acc2_pub_sign_key,
         },
     ]
@@ -150,7 +149,10 @@ pub fn initial_priv_accounts_private_keys() -> Vec<PrivateAccountPrivateInitialD
 
     vec![
         PrivateAccountPrivateInitialData {
-            account_id: AccountId::private_account_id(&key_chain_1.nullifier_public_key, None),
+            account_id: AccountId::private_account_id(
+                &key_chain_1.nullifier_public_key,
+                Identifier(13_u128),
+            ),
             account: Account {
                 program_owner: DEFAULT_PROGRAM_OWNER,
                 balance: PRIV_ACC_A_INITIAL_BALANCE,
@@ -160,7 +162,10 @@ pub fn initial_priv_accounts_private_keys() -> Vec<PrivateAccountPrivateInitialD
             key_chain: key_chain_1,
         },
         PrivateAccountPrivateInitialData {
-            account_id: AccountId::private_account_id(&key_chain_2.nullifier_public_key, None),
+            account_id: AccountId::private_account_id(
+                &key_chain_2.nullifier_public_key,
+                Identifier(42_u128),
+            ),
             account: Account {
                 program_owner: DEFAULT_PROGRAM_OWNER,
                 balance: PRIV_ACC_B_INITIAL_BALANCE,
@@ -208,7 +213,7 @@ pub fn initial_state() -> V03State {
         .iter()
         .map(|init_comm_data| {
             let npk = &init_comm_data.npk;
-            let acc_id = &AccountId::private_account_id(npk, None);
+            let acc_id = &AccountId::private_account_id(npk, Identifier(0_u128));
 
             let mut acc = init_comm_data.account.clone();
 

wallet/src/lib.rs

@@ -396,7 +396,7 @@ impl WalletCore {
                 .map(|keys| (keys.npk.clone(), keys.ssk))
                 .collect::<Vec<_>>(),
             acc_manager.private_account_auth(),
-            vec![], // TODO check (Marvin)
+            vec![],
             acc_manager.private_account_membership_proofs(),
             &program.to_owned(),
         )
@@ -410,7 +410,7 @@ impl WalletCore {
                     .iter()
                     .map(|keys| {
                         (
-                            AccountId::private_account_id(&keys.npk.clone(), None),
+                            AccountId::private_account_id(&keys.npk.clone(), Identifier(0_u128)),
                             keys.vpk.clone(),
                             keys.epk.clone(),
                         )
@@ -505,7 +505,10 @@ impl WalletCore {
         let affected_accounts = private_account_key_chains
             .flat_map(|(acc_account_id, key_chain, index)| {
                 let view_tag = EncryptedAccountData::compute_view_tag(
-                    &AccountId::private_account_id(&key_chain.nullifier_public_key, None),
+                    &AccountId::private_account_id(
+                        &key_chain.nullifier_public_key,
+                        Identifier(0_u128),
+                    ),
                     &key_chain.viewing_public_key,
                 );