lssa/nssa/src/state.rs

use std::collections::{BTreeSet, HashMap, HashSet};

use borsh::{BorshDeserialize, BorshSerialize};
use clock_core::ClockAccountData;
pub use clock_core::{
    CLOCK_01_PROGRAM_ACCOUNT_ID, CLOCK_10_PROGRAM_ACCOUNT_ID, CLOCK_50_PROGRAM_ACCOUNT_ID,
    CLOCK_PROGRAM_ACCOUNT_IDS,
};
use nssa_core::{
    BlockId, Commitment, CommitmentSetDigest, DUMMY_COMMITMENT, MembershipProof, Nullifier,
    Timestamp,
    account::{Account, AccountId, Nonce},
    program::ProgramId,
};

use crate::{
    error::NssaError,
    merkle_tree::MerkleTree,
    privacy_preserving_transaction::PrivacyPreservingTransaction,
    program::Program,
    program_deployment_transaction::ProgramDeploymentTransaction,
    public_transaction::PublicTransaction,
    validated_state_diff::{StateDiff, ValidatedStateDiff},
};

pub const MAX_NUMBER_CHAINED_CALLS: usize = 10;

#[derive(Clone, BorshSerialize, BorshDeserialize)]
#[cfg_attr(test, derive(Debug, PartialEq, Eq))]
pub struct CommitmentSet {
    merkle_tree: MerkleTree,
    commitments: HashMap<Commitment, usize>,
    root_history: HashSet<CommitmentSetDigest>,
}

impl CommitmentSet {
    pub(crate) fn digest(&self) -> CommitmentSetDigest {
        self.merkle_tree.root()
    }

    /// Queries the `CommitmentSet` for a membership proof of commitment.
    pub fn get_proof_for(&self, commitment: &Commitment) -> Option<MembershipProof> {
        let index = *self.commitments.get(commitment)?;

        self.merkle_tree
            .get_authentication_path_for(index)
            .map(|path| (index, path))
    }

    /// Inserts a list of commitments to the `CommitmentSet`.
    pub(crate) fn extend(&mut self, commitments: &[Commitment]) {
        for commitment in commitments.iter().cloned() {
            let index = self.merkle_tree.insert(commitment.to_byte_array());
            self.commitments.insert(commitment, index);
        }
        self.root_history.insert(self.digest());
    }

    fn contains(&self, commitment: &Commitment) -> bool {
        self.commitments.contains_key(commitment)
    }

    /// Initializes an empty `CommitmentSet` with a given capacity.
    /// If the capacity is not a `power_of_two`, then capacity is taken
    /// to be the next `power_of_two`.
    pub(crate) fn with_capacity(capacity: usize) -> Self {
        Self {
            merkle_tree: MerkleTree::with_capacity(capacity),
            commitments: HashMap::new(),
            root_history: HashSet::new(),
        }
    }
}

#[cfg_attr(test, derive(Debug, PartialEq, Eq))]
#[derive(Clone)]
struct NullifierSet(BTreeSet<Nullifier>);

impl NullifierSet {
    const fn new() -> Self {
        Self(BTreeSet::new())
    }

    fn extend(&mut self, new_nullifiers: &[Nullifier]) {
        self.0.extend(new_nullifiers);
    }

    fn contains(&self, nullifier: &Nullifier) -> bool {
        self.0.contains(nullifier)
    }
}

impl BorshSerialize for NullifierSet {
    fn serialize<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
        self.0.iter().collect::<Vec<_>>().serialize(writer)
    }
}

impl BorshDeserialize for NullifierSet {
    fn deserialize_reader<R: std::io::Read>(reader: &mut R) -> std::io::Result<Self> {
        let vec = Vec::<Nullifier>::deserialize_reader(reader)?;

        let mut set = BTreeSet::new();
        for n in vec {
            if !set.insert(n) {
                return Err(std::io::Error::new(
                    std::io::ErrorKind::InvalidData,
                    "duplicate nullifier in NullifierSet",
                ));
            }
        }

        Ok(Self(set))
    }
}

#[derive(Clone, BorshSerialize, BorshDeserialize)]
#[cfg_attr(test, derive(Debug, PartialEq, Eq))]
pub struct V03State {
    public_state: HashMap<AccountId, Account>,
    private_state: (CommitmentSet, NullifierSet),
    programs: HashMap<ProgramId, Program>,
}

impl V03State {
    #[must_use]
    pub fn new_with_genesis_accounts(
        initial_data: &[(AccountId, u128)],
        initial_commitments: &[nssa_core::Commitment],
        genesis_timestamp: nssa_core::Timestamp,
    ) -> Self {
        let authenticated_transfer_program = Program::authenticated_transfer_program();
        let public_state = initial_data
            .iter()
            .copied()
            .map(|(account_id, balance)| {
                let account = Account {
                    balance,
                    program_owner: authenticated_transfer_program.id(),
                    ..Account::default()
                };
                (account_id, account)
            })
            .collect();

        let mut private_state = CommitmentSet::with_capacity(32);
        private_state.extend(&[DUMMY_COMMITMENT]);
        private_state.extend(initial_commitments);

        let mut this = Self {
            public_state,
            private_state: (private_state, NullifierSet::new()),
            programs: HashMap::new(),
        };

        this.insert_program(Program::clock());
        this.insert_clock_accounts(genesis_timestamp);

        this.insert_program(Program::authenticated_transfer_program());
        this.insert_program(Program::token());
        this.insert_program(Program::amm());
        this.insert_program(Program::ata());

        this
    }

    fn insert_clock_accounts(&mut self, genesis_timestamp: nssa_core::Timestamp) {
        let data = ClockAccountData {
            block_id: 0,
            timestamp: genesis_timestamp,
        }
        .to_bytes();
        let clock_program_id = Program::clock().id();
        for account_id in CLOCK_PROGRAM_ACCOUNT_IDS {
            self.public_state.insert(
                account_id,
                Account {
                    program_owner: clock_program_id,
                    data: data
                        .clone()
                        .try_into()
                        .expect("Clock account data should fit within accounts data"),
                    ..Account::default()
                },
            );
        }
    }

    pub(crate) fn insert_program(&mut self, program: Program) {
        self.programs.insert(program.id(), program);
    }

    pub fn apply_state_diff(&mut self, diff: ValidatedStateDiff) {
        let StateDiff {
            signer_account_ids,
            public_diff,
            new_commitments,
            new_nullifiers,
            program,
        } = diff.into_state_diff();
        #[expect(
            clippy::iter_over_hash_type,
            reason = "Iteration order doesn't matter here"
        )]
        for (account_id, account) in public_diff {
            *self.get_account_by_id_mut(account_id) = account;
        }
        for account_id in signer_account_ids {
            self.get_account_by_id_mut(account_id)
                .nonce
                .public_account_nonce_increment();
        }
        self.private_state.0.extend(&new_commitments);
        self.private_state.1.extend(&new_nullifiers);
        if let Some(program) = program {
            self.insert_program(program);
        }
    }

    pub fn transition_from_public_transaction(
        &mut self,
        tx: &PublicTransaction,
        block_id: BlockId,
        timestamp: Timestamp,
    ) -> Result<(), NssaError> {
        let diff = ValidatedStateDiff::from_public_transaction(tx, self, block_id, timestamp)?;
        self.apply_state_diff(diff);
        Ok(())
    }

    pub fn transition_from_privacy_preserving_transaction(
        &mut self,
        tx: &PrivacyPreservingTransaction,
        block_id: BlockId,
        timestamp: Timestamp,
    ) -> Result<(), NssaError> {
        let diff =
            ValidatedStateDiff::from_privacy_preserving_transaction(tx, self, block_id, timestamp)?;
        self.apply_state_diff(diff);
        Ok(())
    }

    pub fn transition_from_program_deployment_transaction(
        &mut self,
        tx: &ProgramDeploymentTransaction,
    ) -> Result<(), NssaError> {
        let diff = ValidatedStateDiff::from_program_deployment_transaction(tx, self)?;
        self.apply_state_diff(diff);
        Ok(())
    }

    fn get_account_by_id_mut(&mut self, account_id: AccountId) -> &mut Account {
        self.public_state.entry(account_id).or_default()
    }

    #[must_use]
    pub fn get_account_by_id(&self, account_id: AccountId) -> Account {
        self.public_state
            .get(&account_id)
            .cloned()
            .unwrap_or_else(Account::default)
    }

    #[must_use]
    pub fn get_proof_for_commitment(&self, commitment: &Commitment) -> Option<MembershipProof> {
        self.private_state.0.get_proof_for(commitment)
    }

    pub(crate) const fn programs(&self) -> &HashMap<ProgramId, Program> {
        &self.programs
    }

    #[must_use]
    pub fn commitment_set_digest(&self) -> CommitmentSetDigest {
        self.private_state.0.digest()
    }

    pub(crate) fn check_commitments_are_new(
        &self,
        new_commitments: &[Commitment],
    ) -> Result<(), NssaError> {
        for commitment in new_commitments {
            if self.private_state.0.contains(commitment) {
                return Err(NssaError::InvalidInput(
                    "Commitment already seen".to_owned(),
                ));
            }
        }
        Ok(())
    }

    pub(crate) fn check_nullifiers_are_valid(
        &self,
        new_nullifiers: &[(Nullifier, CommitmentSetDigest)],
    ) -> Result<(), NssaError> {
        for (nullifier, digest) in new_nullifiers {
            if self.private_state.1.contains(nullifier) {
                return Err(NssaError::InvalidInput("Nullifier already seen".to_owned()));
            }
            if !self.private_state.0.root_history.contains(digest) {
                return Err(NssaError::InvalidInput(
                    "Unrecognized commitment set digest".to_owned(),
                ));
            }
        }
        Ok(())
    }
}

// TODO: Testnet only. Refactor to prevent compilation on mainnet.
impl V03State {
    pub fn add_pinata_program(&mut self, account_id: AccountId) {
        self.insert_program(Program::pinata());

        self.public_state.insert(
            account_id,
            Account {
                program_owner: Program::pinata().id(),
                balance: 1_500_000,
                // Difficulty: 3
                data: vec![3; 33].try_into().expect("should fit"),
                nonce: Nonce::default(),
            },
        );
    }

    pub fn add_pinata_token_program(&mut self, account_id: AccountId) {
        self.insert_program(Program::pinata_token());

        self.public_state.insert(
            account_id,
            Account {
                program_owner: Program::pinata_token().id(),
                // Difficulty: 3
                data: vec![3; 33].try_into().expect("should fit"),
                ..Account::default()
            },
        );
    }
}

#[cfg(any(test, feature = "test-utils"))]
impl V03State {
    pub fn force_insert_account(&mut self, account_id: AccountId, account: Account) {
        self.public_state.insert(account_id, account);
    }
}

#[cfg(test)]
pub mod tests {
    #![expect(
        clippy::arithmetic_side_effects,
        clippy::shadow_unrelated,
        reason = "We don't care about it in tests"
    )]

    use std::collections::HashMap;

    use nssa_core::account::Identifier;
    #[allow(unused_imports)]
    use nssa_core::{
        BlockId, Commitment, Nullifier, NullifierPublicKey, NullifierSecretKey, PrivateKey,
        PublicKey, SharedSecretKey, Timestamp,
        account::{Account, AccountId, AccountWithMetadata, Nonce, data::Data},
        encryption::{EphemeralPublicKey, Scalar, ViewingPublicKey},
        program::{BlockValidityWindow, PdaSeed, ProgramId, TimestampValidityWindow},
    };

    use crate::{
        PublicTransaction, V03State,
        error::NssaError,
        execute_and_prove,
        privacy_preserving_transaction::{
            PrivacyPreservingTransaction,
            circuit::{self, ProgramWithDependencies},
            message::Message,
            witness_set::WitnessSet,
        },
        program::Program,
        public_transaction,
        signature::PrivateKey,
        state::{
            CLOCK_01_PROGRAM_ACCOUNT_ID, CLOCK_10_PROGRAM_ACCOUNT_ID, CLOCK_50_PROGRAM_ACCOUNT_ID,
            CLOCK_PROGRAM_ACCOUNT_IDS, MAX_NUMBER_CHAINED_CALLS,
        },
    };

    impl V03State {
        /// Include test programs in the builtin programs map.
        #[must_use]
        pub fn with_test_programs(mut self) -> Self {
            self.insert_program(Program::nonce_changer_program());
            self.insert_program(Program::extra_output_program());
            self.insert_program(Program::missing_output_program());
            self.insert_program(Program::program_owner_changer());
            self.insert_program(Program::simple_balance_transfer());
            self.insert_program(Program::data_changer());
            self.insert_program(Program::minter());
            self.insert_program(Program::burner());
            self.insert_program(Program::chain_caller());
            self.insert_program(Program::amm());
            self.insert_program(Program::claimer());
            self.insert_program(Program::changer_claimer());
            self.insert_program(Program::validity_window());
            self.insert_program(Program::flash_swap_initiator());
            self.insert_program(Program::flash_swap_callback());
            self.insert_program(Program::malicious_self_program_id());
            self.insert_program(Program::malicious_caller_program_id());
            self.insert_program(Program::time_locked_transfer());
            self.insert_program(Program::pinata_cooldown());
            self
        }

        #[must_use]
        pub fn with_non_default_accounts_but_default_program_owners(mut self) -> Self {
            let account_with_default_values_except_balance = Account {
                balance: 100,
                ..Account::default()
            };
            let account_with_default_values_except_nonce = Account {
                nonce: Nonce(37),
                ..Account::default()
            };
            let account_with_default_values_except_data = Account {
                data: vec![0xca, 0xfe].try_into().unwrap(),
                ..Account::default()
            };
            self.force_insert_account(
                AccountId::new([255; 32]),
                account_with_default_values_except_balance,
            );
            self.force_insert_account(
                AccountId::new([254; 32]),
                account_with_default_values_except_nonce,
            );
            self.force_insert_account(
                AccountId::new([253; 32]),
                account_with_default_values_except_data,
            );
            self
        }

        #[must_use]
        pub fn with_account_owned_by_burner_program(mut self) -> Self {
            let account = Account {
                program_owner: Program::burner().id(),
                balance: 100,
                ..Default::default()
            };
            self.force_insert_account(AccountId::new([252; 32]), account);
            self
        }

        #[must_use]
        pub fn with_private_account(mut self, keys: &TestPrivateKeys, account: &Account) -> Self {
            let account_id = &AccountId::private_account_id(&keys.npk(), keys.identifier);
            let commitment = Commitment::new(account_id, account);
            self.private_state.0.extend(&[commitment]);
            self
        }
    }

    pub struct TestPublicKeys {
        pub signing_key: PrivateKey,
    }

    impl TestPublicKeys {
        pub fn account_id(&self) -> AccountId {
            AccountId::public_account_id(&PublicKey::new_from_private_key(&self.signing_key))
        }
    }

    pub struct TestPrivateKeys {
        pub nsk: NullifierSecretKey,
        pub vsk: Scalar,
        pub identifier: Identifier,
    }

    impl TestPrivateKeys {
        pub fn npk(&self) -> NullifierPublicKey {
            NullifierPublicKey::from(&self.nsk)
        }

        pub fn vpk(&self) -> ViewingPublicKey {
            ViewingPublicKey::from_scalar(self.vsk)
        }
    }

    // ── Flash Swap types (mirrors of guest types for host-side serialisation) ──

    #[derive(serde::Serialize, serde::Deserialize)]
    struct CallbackInstruction {
        return_funds: bool,
        token_program_id: ProgramId,
        amount: u128,
    }

    #[derive(serde::Serialize, serde::Deserialize)]
    enum FlashSwapInstruction {
        Initiate {
            token_program_id: ProgramId,
            callback_program_id: ProgramId,
            amount_out: u128,
            callback_instruction_data: Vec<u32>,
        },
        InvariantCheck {
            min_vault_balance: u128,
        },
    }

    fn transfer_transaction(
        from: AccountId,
        from_key: &PrivateKey,
        from_nonce: u128,
        to: AccountId,
        to_key: &PrivateKey,
        to_nonce: u128,
        balance: u128,
    ) -> PublicTransaction {
        let account_ids = vec![from, to];
        let nonces = vec![Nonce(from_nonce), Nonce(to_nonce)];
        let program_id = Program::authenticated_transfer_program().id();
        let message =
            public_transaction::Message::try_new(program_id, account_ids, nonces, balance).unwrap();
        let witness_set =
            public_transaction::WitnessSet::for_message(&message, &[from_key, to_key]);
        PublicTransaction::new(message, witness_set)
    }

    fn build_flash_swap_tx(
        initiator: &Program,
        vault_id: AccountId,
        receiver_id: AccountId,
        instruction: FlashSwapInstruction,
    ) -> PublicTransaction {
        let message = public_transaction::Message::try_new(
            initiator.id(),
            vec![vault_id, receiver_id],
            vec![], // no signers — vault is PDA-authorised
            instruction,
        )
        .unwrap();
        let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
        PublicTransaction::new(message, witness_set)
    }

    #[test]
    fn new_with_genesis() {
        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));
        let addr2 = AccountId::public_account_id(&PublicKey::new_from_private_key(&key2));
        let initial_data = [(addr1, 100_u128), (addr2, 151_u128)];
        let authenticated_transfers_program = Program::authenticated_transfer_program();
        let clock_program = Program::clock();
        let expected_public_state = {
            let mut this = HashMap::new();
            this.insert(
                addr1,
                Account {
                    balance: 100,
                    program_owner: authenticated_transfers_program.id(),
                    ..Account::default()
                },
            );
            this.insert(
                addr2,
                Account {
                    balance: 151,
                    program_owner: authenticated_transfers_program.id(),
                    ..Account::default()
                },
            );
            for account_id in CLOCK_PROGRAM_ACCOUNT_IDS {
                this.insert(
                    account_id,
                    Account {
                        program_owner: clock_program.id(),
                        data: [0_u8; 16].to_vec().try_into().unwrap(),
                        ..Account::default()
                    },
                );
            }
            this
        };
        let expected_builtin_programs = {
            let mut this = HashMap::new();
            this.insert(
                authenticated_transfers_program.id(),
                authenticated_transfers_program,
            );
            this.insert(clock_program.id(), clock_program);
            this.insert(Program::token().id(), Program::token());
            this.insert(Program::amm().id(), Program::amm());
            this.insert(Program::ata().id(), Program::ata());
            this
        };

        let state = V03State::new_with_genesis_accounts(&initial_data, &[], 0);

        assert_eq!(state.public_state, expected_public_state);
        assert_eq!(state.programs, expected_builtin_programs);
    }

    #[test]
    fn insert_program() {
        let mut state = V03State::new_with_genesis_accounts(&[], &[], 0);
        let program_to_insert = Program::simple_balance_transfer();
        let program_id = program_to_insert.id();
        assert!(!state.programs.contains_key(&program_id));

        state.insert_program(program_to_insert);

        assert!(state.programs.contains_key(&program_id));
    }

    #[test]
    fn get_account_by_account_id_non_default_account() {
        let key = PrivateKey::try_new([1; 32]).unwrap();
        let account_id = AccountId::public_account_id(&PublicKey::new_from_private_key(&key));
        let initial_data = [(account_id, 100_u128)];
        let state = V03State::new_with_genesis_accounts(&initial_data, &[], 0);
        let expected_account = &state.public_state[&account_id];

        let account = state.get_account_by_id(account_id);

        assert_eq!(&account, expected_account);
    }

    #[test]
    fn get_account_by_account_id_default_account() {
        let addr2 = AccountId::new([0; 32]);
        let state = V03State::new_with_genesis_accounts(&[], &[], 0);
        let expected_account = Account::default();

        let account = state.get_account_by_id(addr2);

        assert_eq!(account, expected_account);
    }

    #[test]
    fn builtin_programs_getter() {
        let state = V03State::new_with_genesis_accounts(&[], &[], 0);

        let builtin_programs = state.programs();

        assert_eq!(builtin_programs, &state.programs);
    }

    #[test]
    fn transition_from_authenticated_transfer_program_invocation_default_account_destination() {
        let key = PrivateKey::try_new([1; 32]).unwrap();
        let account_id = AccountId::public_account_id(&PublicKey::new_from_private_key(&key));
        let initial_data = [(account_id, 100)];
        let mut state = V03State::new_with_genesis_accounts(&initial_data, &[], 0);
        let from = account_id;
        let to_key = PrivateKey::try_new([2; 32]).unwrap();
        let to = AccountId::public_account_id(&PublicKey::new_from_private_key(&to_key));
        assert_eq!(state.get_account_by_id(to), Account::default());
        let balance_to_move = 5;

        let tx = transfer_transaction(from, &key, 0, to, &to_key, 0, balance_to_move);
        state.transition_from_public_transaction(&tx, 1, 0).unwrap();

        assert_eq!(state.get_account_by_id(from).balance, 95);
        assert_eq!(state.get_account_by_id(to).balance, 5);
        assert_eq!(state.get_account_by_id(from).nonce, Nonce(1));
        assert_eq!(state.get_account_by_id(to).nonce, Nonce(1));
    }

    #[test]
    fn transition_from_authenticated_transfer_program_invocation_insuficient_balance() {
        let key = PrivateKey::try_new([1; 32]).unwrap();
        let account_id = AccountId::public_account_id(&PublicKey::new_from_private_key(&key));
        let initial_data = [(account_id, 100)];
        let mut state = V03State::new_with_genesis_accounts(&initial_data, &[], 0);
        let from = account_id;
        let from_key = key;
        let to_key = PrivateKey::try_new([2; 32]).unwrap();
        let to = AccountId::public_account_id(&PublicKey::new_from_private_key(&to_key));
        let balance_to_move = 101;
        assert!(state.get_account_by_id(from).balance < balance_to_move);

        let tx = transfer_transaction(from, &from_key, 0, to, &to_key, 0, balance_to_move);
        let result = state.transition_from_public_transaction(&tx, 1, 0);

        assert!(matches!(result, Err(NssaError::ProgramExecutionFailed(_))));
        assert_eq!(state.get_account_by_id(from).balance, 100);
        assert_eq!(state.get_account_by_id(to).balance, 0);
        assert_eq!(state.get_account_by_id(from).nonce, Nonce(0));
        assert_eq!(state.get_account_by_id(to).nonce, Nonce(0));
    }

    #[test]
    fn transition_from_authenticated_transfer_program_invocation_non_default_account_destination() {
        let key1 = PrivateKey::try_new([1; 32]).unwrap();
        let key2 = PrivateKey::try_new([2; 32]).unwrap();
        let account_id1 = AccountId::public_account_id(&PublicKey::new_from_private_key(&key1));
        let account_id2 = AccountId::public_account_id(&PublicKey::new_from_private_key(&key2));
        let initial_data = [(account_id1, 100), (account_id2, 200)];
        let mut state = V03State::new_with_genesis_accounts(&initial_data, &[], 0);
        let from = account_id2;
        let from_key = key2;
        let to = account_id1;
        let to_key = key1;
        assert_ne!(state.get_account_by_id(to), Account::default());
        let balance_to_move = 8;

        let tx = transfer_transaction(from, &from_key, 0, to, &to_key, 0, balance_to_move);
        state.transition_from_public_transaction(&tx, 1, 0).unwrap();

        assert_eq!(state.get_account_by_id(from).balance, 192);
        assert_eq!(state.get_account_by_id(to).balance, 108);
        assert_eq!(state.get_account_by_id(from).nonce, Nonce(1));
        assert_eq!(state.get_account_by_id(to).nonce, Nonce(1));
    }

    #[test]
    fn transition_from_sequence_of_authenticated_transfer_program_invocations() {
        let key1 = PrivateKey::try_new([8; 32]).unwrap();
        let account_id1 = AccountId::public_account_id(&PublicKey::new_from_private_key(&key1));
        let key2 = PrivateKey::try_new([2; 32]).unwrap();
        let account_id2 = AccountId::public_account_id(&PublicKey::new_from_private_key(&key2));
        let initial_data = [(account_id1, 100)];
        let mut state = V03State::new_with_genesis_accounts(&initial_data, &[], 0);
        let key3 = PrivateKey::try_new([3; 32]).unwrap();
        let account_id3 = AccountId::public_account_id(&PublicKey::new_from_private_key(&key3));
        let balance_to_move = 5;

        let tx = transfer_transaction(
            account_id1,
            &key1,
            0,
            account_id2,
            &key2,
            0,
            balance_to_move,
        );
        state.transition_from_public_transaction(&tx, 1, 0).unwrap();
        let balance_to_move = 3;
        let tx = transfer_transaction(
            account_id2,
            &key2,
            1,
            account_id3,
            &key3,
            0,
            balance_to_move,
        );
        state.transition_from_public_transaction(&tx, 1, 0).unwrap();

        assert_eq!(state.get_account_by_id(account_id1).balance, 95);
        assert_eq!(state.get_account_by_id(account_id2).balance, 2);
        assert_eq!(state.get_account_by_id(account_id3).balance, 3);
        assert_eq!(state.get_account_by_id(account_id1).nonce, Nonce(1));
        assert_eq!(state.get_account_by_id(account_id2).nonce, Nonce(2));
        assert_eq!(state.get_account_by_id(account_id3).nonce, Nonce(1));
    }

    fn clock_transaction(timestamp: nssa_core::Timestamp) -> PublicTransaction {
        let message = public_transaction::Message::try_new(
            Program::clock().id(),
            CLOCK_PROGRAM_ACCOUNT_IDS.to_vec(),
            vec![],
            timestamp,
        )
        .unwrap();
        PublicTransaction::new(
            message,
            public_transaction::WitnessSet::from_raw_parts(vec![]),
        )
    }

    fn clock_account_data(state: &V03State, account_id: AccountId) -> (u64, nssa_core::Timestamp) {
        let data = state.get_account_by_id(account_id).data.into_inner();
        let parsed = clock_core::ClockAccountData::from_bytes(&data);
        (parsed.block_id, parsed.timestamp)
    }

    #[test]
    fn clock_genesis_state_has_zero_block_id_and_genesis_timestamp() {
        let genesis_timestamp = 1_000_000_u64;
        let state = V03State::new_with_genesis_accounts(&[], &[], genesis_timestamp);

        let (block_id, timestamp) = clock_account_data(&state, CLOCK_01_PROGRAM_ACCOUNT_ID);

        assert_eq!(block_id, 0);
        assert_eq!(timestamp, genesis_timestamp);
    }

    #[test]
    fn clock_invocation_increments_block_id() {
        let mut state = V03State::new_with_genesis_accounts(&[], &[], 0);

        let tx = clock_transaction(1234);
        state.transition_from_public_transaction(&tx, 0, 0).unwrap();

        let (block_id, _) = clock_account_data(&state, CLOCK_01_PROGRAM_ACCOUNT_ID);
        assert_eq!(block_id, 1);
    }

    #[test]
    fn clock_invocation_stores_timestamp_from_instruction() {
        let mut state = V03State::new_with_genesis_accounts(&[], &[], 0);
        let block_timestamp = 1_700_000_000_000_u64;

        let tx = clock_transaction(block_timestamp);
        state.transition_from_public_transaction(&tx, 0, 0).unwrap();

        let (_, timestamp) = clock_account_data(&state, CLOCK_01_PROGRAM_ACCOUNT_ID);
        assert_eq!(timestamp, block_timestamp);
    }

    #[test]
    fn clock_invocation_sequence_correctly_increments_block_id() {
        let mut state = V03State::new_with_genesis_accounts(&[], &[], 0);

        for expected_block_id in 1_u64..=5 {
            let tx = clock_transaction(expected_block_id * 1000);
            state.transition_from_public_transaction(&tx, 0, 0).unwrap();

            let (block_id, timestamp) = clock_account_data(&state, CLOCK_01_PROGRAM_ACCOUNT_ID);
            assert_eq!(block_id, expected_block_id);
            assert_eq!(timestamp, expected_block_id * 1000);
        }
    }

    #[test]
    fn clock_10_account_not_updated_when_block_id_not_multiple_of_10() {
        let genesis_timestamp = 0_u64;
        let mut state = V03State::new_with_genesis_accounts(&[], &[], genesis_timestamp);

        // Run 9 clock ticks (block_ids 1..=9), none of which are multiples of 10.
        for tick in 1_u64..=9 {
            let tx = clock_transaction(tick * 1000);
            state.transition_from_public_transaction(&tx, 0, 0).unwrap();
        }

        let (block_id_10, timestamp_10) = clock_account_data(&state, CLOCK_10_PROGRAM_ACCOUNT_ID);
        // The 10-block account should still reflect genesis state.
        assert_eq!(block_id_10, 0);
        assert_eq!(timestamp_10, genesis_timestamp);
    }

    #[test]
    fn clock_10_account_updated_when_block_id_is_multiple_of_10() {
        let mut state = V03State::new_with_genesis_accounts(&[], &[], 0);

        // Run 10 clock ticks so block_id reaches 10.
        for tick in 1_u64..=10 {
            let tx = clock_transaction(tick * 1000);
            state.transition_from_public_transaction(&tx, 0, 0).unwrap();
        }

        let (block_id_1, timestamp_1) = clock_account_data(&state, CLOCK_01_PROGRAM_ACCOUNT_ID);
        let (block_id_10, timestamp_10) = clock_account_data(&state, CLOCK_10_PROGRAM_ACCOUNT_ID);
        assert_eq!(block_id_1, 10);
        assert_eq!(block_id_10, 10);
        assert_eq!(timestamp_10, timestamp_1);
    }

    #[test]
    fn clock_50_account_only_updated_at_multiples_of_50() {
        let mut state = V03State::new_with_genesis_accounts(&[], &[], 0);

        // After 49 ticks the 50-block account should be unchanged.
        for tick in 1_u64..=49 {
            let tx = clock_transaction(tick * 1000);
            state.transition_from_public_transaction(&tx, 0, 0).unwrap();
        }
        let (block_id_50, _) = clock_account_data(&state, CLOCK_50_PROGRAM_ACCOUNT_ID);
        assert_eq!(block_id_50, 0);

        // Tick 50 — now the 50-block account should update.
        let tx = clock_transaction(50 * 1000);
        state.transition_from_public_transaction(&tx, 0, 0).unwrap();
        let (block_id_50, timestamp_50) = clock_account_data(&state, CLOCK_50_PROGRAM_ACCOUNT_ID);
        assert_eq!(block_id_50, 50);
        assert_eq!(timestamp_50, 50 * 1000);
    }

    #[test]
    fn all_three_clock_accounts_updated_at_multiple_of_50() {
        let mut state = V03State::new_with_genesis_accounts(&[], &[], 0);

        // Advance to block 50 (a multiple of both 10 and 50).
        for tick in 1_u64..=50 {
            let tx = clock_transaction(tick * 1000);
            state.transition_from_public_transaction(&tx, 0, 0).unwrap();
        }

        let (block_id_1, ts_1) = clock_account_data(&state, CLOCK_01_PROGRAM_ACCOUNT_ID);
        let (block_id_10, ts_10) = clock_account_data(&state, CLOCK_10_PROGRAM_ACCOUNT_ID);
        let (block_id_50, ts_50) = clock_account_data(&state, CLOCK_50_PROGRAM_ACCOUNT_ID);

        assert_eq!(block_id_1, 50);
        assert_eq!(block_id_10, 50);
        assert_eq!(block_id_50, 50);
        assert_eq!(ts_1, ts_10);
        assert_eq!(ts_1, ts_50);
    }

    #[test]
    fn program_should_fail_if_modifies_nonces() {
        let initial_data = [(AccountId::new([1; 32]), 100)];
        let mut state =
            V03State::new_with_genesis_accounts(&initial_data, &[], 0).with_test_programs();
        let account_ids = vec![AccountId::new([1; 32])];
        let program_id = Program::nonce_changer_program().id();
        let message =
            public_transaction::Message::try_new(program_id, account_ids, vec![], ()).unwrap();
        let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
        let tx = PublicTransaction::new(message, witness_set);

        let result = state.transition_from_public_transaction(&tx, 1, 0);

        assert!(matches!(result, Err(NssaError::InvalidProgramBehavior)));
    }

    #[test]
    fn program_should_fail_if_output_accounts_exceed_inputs() {
        let initial_data = [(AccountId::new([1; 32]), 100)];
        let mut state =
            V03State::new_with_genesis_accounts(&initial_data, &[], 0).with_test_programs();
        let account_ids = vec![AccountId::new([1; 32])];
        let program_id = Program::extra_output_program().id();
        let message =
            public_transaction::Message::try_new(program_id, account_ids, vec![], ()).unwrap();
        let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
        let tx = PublicTransaction::new(message, witness_set);

        let result = state.transition_from_public_transaction(&tx, 1, 0);

        assert!(matches!(result, Err(NssaError::InvalidProgramBehavior)));
    }

    #[test]
    fn program_should_fail_with_missing_output_accounts() {
        let initial_data = [(AccountId::new([1; 32]), 100)];
        let mut state =
            V03State::new_with_genesis_accounts(&initial_data, &[], 0).with_test_programs();
        let account_ids = vec![AccountId::new([1; 32]), AccountId::new([2; 32])];
        let program_id = Program::missing_output_program().id();
        let message =
            public_transaction::Message::try_new(program_id, account_ids, vec![], ()).unwrap();
        let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
        let tx = PublicTransaction::new(message, witness_set);

        let result = state.transition_from_public_transaction(&tx, 1, 0);

        assert!(matches!(result, Err(NssaError::InvalidProgramBehavior)));
    }

    #[test]
    fn program_should_fail_if_modifies_program_owner_with_only_non_default_program_owner() {
        let initial_data = [(AccountId::new([1; 32]), 0)];
        let mut state =
            V03State::new_with_genesis_accounts(&initial_data, &[], 0).with_test_programs();
        let account_id = AccountId::new([1; 32]);
        let account = state.get_account_by_id(account_id);
        // Assert the target account only differs from the default account in the program owner
        // field
        assert_ne!(account.program_owner, Account::default().program_owner);
        assert_eq!(account.balance, Account::default().balance);
        assert_eq!(account.nonce, Account::default().nonce);
        assert_eq!(account.data, Account::default().data);
        let program_id = Program::program_owner_changer().id();
        let message =
            public_transaction::Message::try_new(program_id, vec![account_id], vec![], ()).unwrap();
        let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
        let tx = PublicTransaction::new(message, witness_set);

        let result = state.transition_from_public_transaction(&tx, 1, 0);

        assert!(matches!(result, Err(NssaError::InvalidProgramBehavior)));
    }

    #[test]
    fn program_should_fail_if_modifies_program_owner_with_only_non_default_balance() {
        let initial_data = [];
        let mut state = V03State::new_with_genesis_accounts(&initial_data, &[], 0)
            .with_test_programs()
            .with_non_default_accounts_but_default_program_owners();
        let account_id = AccountId::new([255; 32]);
        let account = state.get_account_by_id(account_id);
        // Assert the target account only differs from the default account in balance field
        assert_eq!(account.program_owner, Account::default().program_owner);
        assert_ne!(account.balance, Account::default().balance);
        assert_eq!(account.nonce, Account::default().nonce);
        assert_eq!(account.data, Account::default().data);
        let program_id = Program::program_owner_changer().id();
        let message =
            public_transaction::Message::try_new(program_id, vec![account_id], vec![], ()).unwrap();
        let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
        let tx = PublicTransaction::new(message, witness_set);

        let result = state.transition_from_public_transaction(&tx, 1, 0);

        assert!(matches!(result, Err(NssaError::InvalidProgramBehavior)));
    }

    #[test]
    fn program_should_fail_if_modifies_program_owner_with_only_non_default_nonce() {
        let initial_data = [];
        let mut state = V03State::new_with_genesis_accounts(&initial_data, &[], 0)
            .with_test_programs()
            .with_non_default_accounts_but_default_program_owners();
        let account_id = AccountId::new([254; 32]);
        let account = state.get_account_by_id(account_id);
        // Assert the target account only differs from the default account in nonce field
        assert_eq!(account.program_owner, Account::default().program_owner);
        assert_eq!(account.balance, Account::default().balance);
        assert_ne!(account.nonce, Account::default().nonce);
        assert_eq!(account.data, Account::default().data);
        let program_id = Program::program_owner_changer().id();
        let message =
            public_transaction::Message::try_new(program_id, vec![account_id], vec![], ()).unwrap();
        let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
        let tx = PublicTransaction::new(message, witness_set);

        let result = state.transition_from_public_transaction(&tx, 1, 0);

        assert!(matches!(result, Err(NssaError::InvalidProgramBehavior)));
    }

    #[test]
    fn program_should_fail_if_modifies_program_owner_with_only_non_default_data() {
        let initial_data = [];
        let mut state = V03State::new_with_genesis_accounts(&initial_data, &[], 0)
            .with_test_programs()
            .with_non_default_accounts_but_default_program_owners();
        let account_id = AccountId::new([253; 32]);
        let account = state.get_account_by_id(account_id);
        // Assert the target account only differs from the default account in data field
        assert_eq!(account.program_owner, Account::default().program_owner);
        assert_eq!(account.balance, Account::default().balance);
        assert_eq!(account.nonce, Account::default().nonce);
        assert_ne!(account.data, Account::default().data);
        let program_id = Program::program_owner_changer().id();
        let message =
            public_transaction::Message::try_new(program_id, vec![account_id], vec![], ()).unwrap();
        let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
        let tx = PublicTransaction::new(message, witness_set);

        let result = state.transition_from_public_transaction(&tx, 1, 0);

        assert!(matches!(result, Err(NssaError::InvalidProgramBehavior)));
    }

    #[test]
    fn program_should_fail_if_transfers_balance_from_non_owned_account() {
        let initial_data = [(AccountId::new([1; 32]), 100)];
        let mut state =
            V03State::new_with_genesis_accounts(&initial_data, &[], 0).with_test_programs();
        let sender_account_id = AccountId::new([1; 32]);
        let receiver_account_id = AccountId::new([2; 32]);
        let balance_to_move: u128 = 1;
        let program_id = Program::simple_balance_transfer().id();
        assert_ne!(
            state.get_account_by_id(sender_account_id).program_owner,
            program_id
        );
        let message = public_transaction::Message::try_new(
            program_id,
            vec![sender_account_id, receiver_account_id],
            vec![],
            balance_to_move,
        )
        .unwrap();
        let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
        let tx = PublicTransaction::new(message, witness_set);

        let result = state.transition_from_public_transaction(&tx, 1, 0);

        assert!(matches!(result, Err(NssaError::InvalidProgramBehavior)));
    }

    #[test]
    fn program_should_fail_if_modifies_data_of_non_owned_account() {
        let initial_data = [];
        let mut state = V03State::new_with_genesis_accounts(&initial_data, &[], 0)
            .with_test_programs()
            .with_non_default_accounts_but_default_program_owners();
        let account_id = AccountId::new([255; 32]);
        let program_id = Program::data_changer().id();

        assert_ne!(state.get_account_by_id(account_id), Account::default());
        assert_ne!(
            state.get_account_by_id(account_id).program_owner,
            program_id
        );
        let message =
            public_transaction::Message::try_new(program_id, vec![account_id], vec![], vec![0])
                .unwrap();
        let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
        let tx = PublicTransaction::new(message, witness_set);

        let result = state.transition_from_public_transaction(&tx, 1, 0);

        assert!(matches!(result, Err(NssaError::InvalidProgramBehavior)));
    }

    #[test]
    fn program_should_fail_if_does_not_preserve_total_balance_by_minting() {
        let initial_data = [];
        let mut state =
            V03State::new_with_genesis_accounts(&initial_data, &[], 0).with_test_programs();
        let account_id = AccountId::new([1; 32]);
        let program_id = Program::minter().id();

        let message =
            public_transaction::Message::try_new(program_id, vec![account_id], vec![], ()).unwrap();
        let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
        let tx = PublicTransaction::new(message, witness_set);

        let result = state.transition_from_public_transaction(&tx, 1, 0);

        assert!(matches!(result, Err(NssaError::InvalidProgramBehavior)));
    }

    #[test]
    fn program_should_fail_if_does_not_preserve_total_balance_by_burning() {
        let initial_data = [];
        let mut state = V03State::new_with_genesis_accounts(&initial_data, &[], 0)
            .with_test_programs()
            .with_account_owned_by_burner_program();
        let program_id = Program::burner().id();
        let account_id = AccountId::new([252; 32]);
        assert_eq!(
            state.get_account_by_id(account_id).program_owner,
            program_id
        );
        let balance_to_burn: u128 = 1;
        assert!(state.get_account_by_id(account_id).balance > balance_to_burn);

        let message = public_transaction::Message::try_new(
            program_id,
            vec![account_id],
            vec![],
            balance_to_burn,
        )
        .unwrap();
        let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
        let tx = PublicTransaction::new(message, witness_set);
        let result = state.transition_from_public_transaction(&tx, 1, 0);

        assert!(matches!(result, Err(NssaError::InvalidProgramBehavior)));
    }

    fn test_public_account_keys_1() -> TestPublicKeys {
        TestPublicKeys {
            signing_key: PrivateKey::try_new([37; 32]).unwrap(),
        }
    }

    pub fn test_private_account_keys_1() -> TestPrivateKeys {
        TestPrivateKeys {
            nsk: [13; 32],
            vsk: [31; 32],
            identifier: Identifier(12_u128),
        }
    }

    pub fn test_private_account_keys_2() -> TestPrivateKeys {
        TestPrivateKeys {
            nsk: [38; 32],
            vsk: [83; 32],
            identifier: Identifier(42_u128),
        }
    }

    fn shielded_balance_transfer_for_tests(
        sender_keys: &TestPublicKeys,
        recipient_keys: &TestPrivateKeys,
        balance_to_move: u128,
        state: &V03State,
    ) -> PrivacyPreservingTransaction {
        let sender = AccountWithMetadata::new(
            state.get_account_by_id(sender_keys.account_id()),
            true,
            sender_keys.account_id(),
        );

        let sender_nonce = sender.account.nonce;

        let recipient_id =
            AccountId::private_account_id(&recipient_keys.npk(), recipient_keys.identifier);
        let recipient = AccountWithMetadata::new(Account::default(), false, recipient_id);

        let esk = [3; 32];
        let shared_secret = SharedSecretKey::new(&esk, &recipient_keys.vpk());
        let epk = EphemeralPublicKey::from_scalar(esk);

        let (output, proof) = circuit::execute_and_prove(
            vec![sender, recipient],
            Program::serialize_instruction(balance_to_move).unwrap(),
            vec![0, 2],
            vec![(recipient_keys.npk(), shared_secret)],
            vec![],
            vec![recipient_keys.identifier],
            vec![None],
            &Program::authenticated_transfer_program().into(),
        )
        .unwrap();

        let message = Message::try_from_circuit_output(
            vec![sender_keys.account_id()],
            vec![sender_nonce],
            vec![(recipient_id, recipient_keys.vpk(), epk)],
            output,
        )
        .unwrap();

        let witness_set = WitnessSet::for_message(&message, proof, &[&sender_keys.signing_key]);
        PrivacyPreservingTransaction::new(message, witness_set)
    }

    fn private_balance_transfer_for_tests(
        sender_keys: &TestPrivateKeys,
        sender_private_account: &Account,
        recipient_keys: &TestPrivateKeys,
        balance_to_move: u128,
        state: &V03State,
    ) -> PrivacyPreservingTransaction {
        let program = Program::authenticated_transfer_program();
        let sender_id = AccountId::private_account_id(&sender_keys.npk(), sender_keys.identifier);
        let recipient_id =
            AccountId::private_account_id(&recipient_keys.npk(), recipient_keys.identifier);
        let sender_commitment = Commitment::new(&sender_id, sender_private_account);
        let sender_pre = AccountWithMetadata::new(sender_private_account.clone(), true, sender_id);
        let recipient_pre = AccountWithMetadata::new(Account::default(), false, recipient_id);

        let esk_1 = [3; 32];
        let shared_secret_1 = SharedSecretKey::new(&esk_1, &sender_keys.vpk());
        let epk_1 = EphemeralPublicKey::from_scalar(esk_1);

        let esk_2 = [3; 32];
        let shared_secret_2 = SharedSecretKey::new(&esk_2, &recipient_keys.vpk());
        let epk_2 = EphemeralPublicKey::from_scalar(esk_2);

        let (output, proof) = circuit::execute_and_prove(
            vec![sender_pre, recipient_pre],
            Program::serialize_instruction(balance_to_move).unwrap(),
            vec![1, 2],
            vec![
                (sender_keys.npk(), shared_secret_1),
                (recipient_keys.npk(), shared_secret_2),
            ],
            vec![sender_keys.nsk],
            vec![sender_keys.identifier, recipient_keys.identifier],
            vec![state.get_proof_for_commitment(&sender_commitment), None],
            &program.into(),
        )
        .unwrap();

        let message = Message::try_from_circuit_output(
            vec![],
            vec![],
            vec![
                (sender_id, sender_keys.vpk(), epk_1),
                (recipient_id, recipient_keys.vpk(), epk_2),
            ],
            output,
        )
        .unwrap();

        let witness_set = WitnessSet::for_message(&message, proof, &[]);

        PrivacyPreservingTransaction::new(message, witness_set)
    }

    fn deshielded_balance_transfer_for_tests(
        sender_keys: &TestPrivateKeys,
        sender_private_account: &Account,
        recipient_account_id: &AccountId,
        balance_to_move: u128,
        state: &V03State,
    ) -> PrivacyPreservingTransaction {
        let program = Program::authenticated_transfer_program();
        let sender_id = AccountId::private_account_id(&sender_keys.npk(), sender_keys.identifier);
        let sender_commitment = Commitment::new(&sender_id, sender_private_account);
        let sender_pre = AccountWithMetadata::new(sender_private_account.clone(), true, sender_id);
        let recipient_pre = AccountWithMetadata::new(
            state.get_account_by_id(*recipient_account_id),
            false,
            *recipient_account_id,
        );

        let esk = [3; 32];
        let shared_secret = SharedSecretKey::new(&esk, &sender_keys.vpk());
        let epk = EphemeralPublicKey::from_scalar(esk);

        let (output, proof) = circuit::execute_and_prove(
            vec![sender_pre, recipient_pre],
            Program::serialize_instruction(balance_to_move).unwrap(),
            vec![1, 0],
            vec![(sender_keys.npk(), shared_secret)],
            vec![sender_keys.nsk],
            vec![sender_keys.identifier],
            vec![state.get_proof_for_commitment(&sender_commitment)],
            &program.into(),
        )
        .unwrap();

        let message = Message::try_from_circuit_output(
            vec![*recipient_account_id],
            vec![],
            vec![(sender_id, sender_keys.vpk(), epk)],
            output,
        )
        .unwrap();

        let witness_set = WitnessSet::for_message(&message, proof, &[]);

        PrivacyPreservingTransaction::new(message, witness_set)
    }

    #[test]
    fn transition_from_privacy_preserving_transaction_shielded() {
        let sender_keys = test_public_account_keys_1();
        let recipient_keys = test_private_account_keys_1();

        let mut state =
            V03State::new_with_genesis_accounts(&[(sender_keys.account_id(), 200)], &[], 0);

        let balance_to_move = 37;

        let tx = shielded_balance_transfer_for_tests(
            &sender_keys,
            &recipient_keys,
            balance_to_move,
            &state,
        );

        let expected_sender_post = {
            let mut this = state.get_account_by_id(sender_keys.account_id());
            this.balance -= balance_to_move;
            this.nonce.public_account_nonce_increment();
            this
        };

        let [expected_new_commitment] = tx.message().new_commitments.clone().try_into().unwrap();
        assert!(!state.private_state.0.contains(&expected_new_commitment));

        state
            .transition_from_privacy_preserving_transaction(&tx, 1, 0)
            .unwrap();

        let sender_post = state.get_account_by_id(sender_keys.account_id());
        assert_eq!(sender_post, expected_sender_post);
        assert!(state.private_state.0.contains(&expected_new_commitment));

        assert_eq!(
            state.get_account_by_id(sender_keys.account_id()).balance,
            200 - balance_to_move
        );
    }

    #[test]
    fn transition_from_privacy_preserving_transaction_private() {
        let sender_keys = test_private_account_keys_1();
        let sender_nonce = Nonce(0xdead_beef);

        let sender_private_account = Account {
            program_owner: Program::authenticated_transfer_program().id(),
            balance: 100,
            nonce: sender_nonce,
            data: Data::default(),
        };
        let recipient_keys = test_private_account_keys_2();

        let mut state = V03State::new_with_genesis_accounts(&[], &[], 0)
            .with_private_account(&sender_keys, &sender_private_account);

        let balance_to_move = 37;

        let tx = private_balance_transfer_for_tests(
            &sender_keys,
            &sender_private_account,
            &recipient_keys,
            balance_to_move,
            &state,
        );

        let sender_id = AccountId::private_account_id(&sender_keys.npk(), sender_keys.identifier);
        let recipient_id =
            AccountId::private_account_id(&recipient_keys.npk(), recipient_keys.identifier);

        let expected_new_commitment_1 = Commitment::new(
            &sender_id,
            &Account {
                program_owner: Program::authenticated_transfer_program().id(),
                nonce: sender_nonce.private_account_nonce_increment(&sender_keys.nsk),
                balance: sender_private_account.balance - balance_to_move,
                data: Data::default(),
            },
        );

        let sender_pre_commitment = Commitment::new(&sender_id, &sender_private_account);
        let expected_new_nullifier =
            Nullifier::for_account_update(&sender_pre_commitment, &sender_keys.nsk);

        let expected_new_commitment_2 = Commitment::new(
            &recipient_id,
            &Account {
                program_owner: Program::authenticated_transfer_program().id(),
                nonce: Nonce::private_account_nonce_init(&recipient_keys.npk()),
                balance: balance_to_move,
                ..Account::default()
            },
        );

        let previous_public_state = state.public_state.clone();
        assert!(state.private_state.0.contains(&sender_pre_commitment));
        assert!(!state.private_state.0.contains(&expected_new_commitment_1));
        assert!(!state.private_state.0.contains(&expected_new_commitment_2));
        assert!(!state.private_state.1.contains(&expected_new_nullifier));

        state
            .transition_from_privacy_preserving_transaction(&tx, 1, 0)
            .unwrap();

        assert_eq!(state.public_state, previous_public_state);
        assert!(state.private_state.0.contains(&sender_pre_commitment));
        assert!(state.private_state.0.contains(&expected_new_commitment_1));
        assert!(state.private_state.0.contains(&expected_new_commitment_2));
        assert!(state.private_state.1.contains(&expected_new_nullifier));
    }

    #[test]
    fn transition_from_privacy_preserving_transaction_deshielded() {
        let sender_keys = test_private_account_keys_1();
        let sender_id = AccountId::private_account_id(&sender_keys.npk(), sender_keys.identifier);
        let sender_nonce = Nonce(0xdead_beef);

        let sender_private_account = Account {
            program_owner: Program::authenticated_transfer_program().id(),
            balance: 100,
            nonce: sender_nonce,
            data: Data::default(),
        };
        let recipient_keys = test_public_account_keys_1();
        let recipient_initial_balance = 400;
        let mut state = V03State::new_with_genesis_accounts(
            &[(recipient_keys.account_id(), recipient_initial_balance)],
            &[],
            0,
        )
        .with_private_account(&sender_keys, &sender_private_account);

        let balance_to_move = 37;

        let expected_recipient_post = {
            let mut this = state.get_account_by_id(recipient_keys.account_id());
            this.balance += balance_to_move;
            this
        };

        let tx = deshielded_balance_transfer_for_tests(
            &sender_keys,
            &sender_private_account,
            &recipient_keys.account_id(),
            balance_to_move,
            &state,
        );

        let expected_new_commitment = Commitment::new(
            &sender_id,
            &Account {
                program_owner: Program::authenticated_transfer_program().id(),
                nonce: sender_nonce.private_account_nonce_increment(&sender_keys.nsk),
                balance: sender_private_account.balance - balance_to_move,
                data: Data::default(),
            },
        );

        let sender_pre_commitment = Commitment::new(&sender_id, &sender_private_account);
        let expected_new_nullifier =
            Nullifier::for_account_update(&sender_pre_commitment, &sender_keys.nsk);

        assert!(state.private_state.0.contains(&sender_pre_commitment));
        assert!(!state.private_state.0.contains(&expected_new_commitment));
        assert!(!state.private_state.1.contains(&expected_new_nullifier));

        state
            .transition_from_privacy_preserving_transaction(&tx, 1, 0)
            .unwrap();

        let recipient_post = state.get_account_by_id(recipient_keys.account_id());
        assert_eq!(recipient_post, expected_recipient_post);
        assert!(state.private_state.0.contains(&sender_pre_commitment));
        assert!(state.private_state.0.contains(&expected_new_commitment));
        assert!(state.private_state.1.contains(&expected_new_nullifier));
        assert_eq!(
            state.get_account_by_id(recipient_keys.account_id()).balance,
            recipient_initial_balance + balance_to_move
        );
    }

    #[test]
    fn burner_program_should_fail_in_privacy_preserving_circuit() {
        let program = Program::burner();
        let public_account = AccountWithMetadata::new(
            Account {
                program_owner: program.id(),
                balance: 100,
                ..Account::default()
            },
            true,
            AccountId::new([0; 32]),
        );

        let result = execute_and_prove(
            vec![public_account],
            Program::serialize_instruction(10_u128).unwrap(),
            vec![0],
            vec![],
            vec![],
            vec![],
            vec![],
            &program.into(),
        );

        assert!(matches!(result, Err(NssaError::CircuitProvingError(_))));
    }

    #[test]
    fn minter_program_should_fail_in_privacy_preserving_circuit() {
        let program = Program::minter();
        let public_account = AccountWithMetadata::new(
            Account {
                program_owner: program.id(),
                balance: 0,
                ..Account::default()
            },
            true,
            AccountId::new([0; 32]),
        );

        let result = execute_and_prove(
            vec![public_account],
            Program::serialize_instruction(10_u128).unwrap(),
            vec![0],
            vec![],
            vec![],
            vec![],
            vec![],
            &program.into(),
        );

        assert!(matches!(result, Err(NssaError::CircuitProvingError(_))));
    }

    #[test]
    fn nonce_changer_program_should_fail_in_privacy_preserving_circuit() {
        let program = Program::nonce_changer_program();
        let public_account = AccountWithMetadata::new(
            Account {
                program_owner: program.id(),
                balance: 0,
                ..Account::default()
            },
            true,
            AccountId::new([0; 32]),
        );

        let result = execute_and_prove(
            vec![public_account],
            Program::serialize_instruction(()).unwrap(),
            vec![0],
            vec![],
            vec![],
            vec![],
            vec![],
            &program.into(),
        );

        assert!(matches!(result, Err(NssaError::CircuitProvingError(_))));
    }

    #[test]
    fn data_changer_program_should_fail_for_non_owned_account_in_privacy_preserving_circuit() {
        let program = Program::data_changer();
        let public_account = AccountWithMetadata::new(
            Account {
                program_owner: [0, 1, 2, 3, 4, 5, 6, 7],
                balance: 0,
                ..Account::default()
            },
            true,
            AccountId::new([0; 32]),
        );

        let result = execute_and_prove(
            vec![public_account],
            Program::serialize_instruction(vec![0]).unwrap(),
            vec![0],
            vec![],
            vec![],
            vec![],
            vec![],
            &program.into(),
        );

        assert!(matches!(result, Err(NssaError::CircuitProvingError(_))));
    }

    #[test]
    fn data_changer_program_should_fail_for_too_large_data_in_privacy_preserving_circuit() {
        let program = Program::data_changer();
        let public_account = AccountWithMetadata::new(
            Account {
                program_owner: program.id(),
                balance: 0,
                ..Account::default()
            },
            true,
            AccountId::new([0; 32]),
        );

        let large_data: Vec<u8> =
            vec![
                0;
                usize::try_from(nssa_core::account::data::DATA_MAX_LENGTH.as_u64())
                    .expect("DATA_MAX_LENGTH fits in usize")
                    + 1
            ];

        let result = execute_and_prove(
            vec![public_account],
            Program::serialize_instruction(large_data).unwrap(),
            vec![0],
            vec![],
            vec![],
            vec![],
            vec![],
            &program.into(),
        );

        assert!(matches!(result, Err(NssaError::ProgramProveFailed(_))));
    }

    #[test]
    fn extra_output_program_should_fail_in_privacy_preserving_circuit() {
        let program = Program::extra_output_program();
        let public_account = AccountWithMetadata::new(
            Account {
                program_owner: program.id(),
                balance: 0,
                ..Account::default()
            },
            true,
            AccountId::new([0; 32]),
        );

        let result = execute_and_prove(
            vec![public_account],
            Program::serialize_instruction(()).unwrap(),
            vec![0],
            vec![],
            vec![],
            vec![],
            vec![],
            &program.into(),
        );

        assert!(matches!(result, Err(NssaError::CircuitProvingError(_))));
    }

    #[test]
    fn missing_output_program_should_fail_in_privacy_preserving_circuit() {
        let program = Program::missing_output_program();
        let public_account_1 = AccountWithMetadata::new(
            Account {
                program_owner: program.id(),
                balance: 0,
                ..Account::default()
            },
            true,
            AccountId::new([0; 32]),
        );
        let public_account_2 = AccountWithMetadata::new(
            Account {
                program_owner: program.id(),
                balance: 0,
                ..Account::default()
            },
            true,
            AccountId::new([1; 32]),
        );

        let result = execute_and_prove(
            vec![public_account_1, public_account_2],
            Program::serialize_instruction(()).unwrap(),
            vec![0, 0],
            vec![],
            vec![],
            vec![],
            vec![],
            &program.into(),
        );

        assert!(matches!(result, Err(NssaError::CircuitProvingError(_))));
    }

    #[test]
    fn program_owner_changer_should_fail_in_privacy_preserving_circuit() {
        let program = Program::program_owner_changer();
        let public_account = AccountWithMetadata::new(
            Account {
                program_owner: program.id(),
                balance: 0,
                ..Account::default()
            },
            true,
            AccountId::new([0; 32]),
        );

        let result = execute_and_prove(
            vec![public_account],
            Program::serialize_instruction(()).unwrap(),
            vec![0],
            vec![],
            vec![],
            vec![],
            vec![],
            &program.into(),
        );

        assert!(matches!(result, Err(NssaError::CircuitProvingError(_))));
    }

    #[test]
    fn transfer_from_non_owned_account_should_fail_in_privacy_preserving_circuit() {
        let program = Program::simple_balance_transfer();
        let public_account_1 = AccountWithMetadata::new(
            Account {
                program_owner: [0, 1, 2, 3, 4, 5, 6, 7],
                balance: 100,
                ..Account::default()
            },
            true,
            AccountId::new([0; 32]),
        );
        let public_account_2 = AccountWithMetadata::new(
            Account {
                program_owner: program.id(),
                balance: 0,
                ..Account::default()
            },
            true,
            AccountId::new([1; 32]),
        );

        let result = execute_and_prove(
            vec![public_account_1, public_account_2],
            Program::serialize_instruction(10_u128).unwrap(),
            vec![0, 0],
            vec![],
            vec![],
            vec![],
            vec![],
            &program.into(),
        );

        assert!(matches!(result, Err(NssaError::CircuitProvingError(_))));
    }

    #[test]
    fn circuit_fails_if_visibility_masks_have_incorrect_lenght() {
        let program = Program::simple_balance_transfer();
        let public_account_1 = AccountWithMetadata::new(
            Account {
                program_owner: program.id(),
                balance: 100,
                ..Account::default()
            },
            true,
            AccountId::new([0; 32]),
        );
        let public_account_2 = AccountWithMetadata::new(
            Account {
                program_owner: program.id(),
                balance: 0,
                ..Account::default()
            },
            true,
            AccountId::new([1; 32]),
        );

        // Setting only one visibility mask for a circuit execution with two pre_state accounts.
        let visibility_mask = [0];
        let result = execute_and_prove(
            vec![public_account_1, public_account_2],
            Program::serialize_instruction(10_u128).unwrap(),
            visibility_mask.to_vec(),
            vec![],
            vec![],
            vec![],
            vec![],
            &program.into(),
        );

        assert!(matches!(result, Err(NssaError::CircuitProvingError(_))));
    }

    #[test]
    fn circuit_fails_if_insufficient_nonces_are_provided() {
        let program = Program::simple_balance_transfer();
        let sender_keys = test_private_account_keys_1();
        let sender_account_id =
            AccountId::private_account_id(&sender_keys.npk(), Identifier(0_u128));
        let recipient_keys = test_private_account_keys_2();
        let recipient_account_id =
            AccountId::private_account_id(&recipient_keys.npk(), Identifier(0_u128));
        let private_account_1 = AccountWithMetadata::new(
            Account {
                program_owner: program.id(),
                balance: 100,
                ..Account::default()
            },
            true,
            sender_account_id,
        );
        let private_account_2 =
            AccountWithMetadata::new(Account::default(), false, recipient_account_id);

        let result = execute_and_prove(
            vec![private_account_1, private_account_2],
            Program::serialize_instruction(10_u128).unwrap(),
            vec![1, 2],
            vec![
                (
                    sender_keys.npk(),
                    SharedSecretKey::new(&[55; 32], &sender_keys.vpk()),
                ),
                (
                    recipient_keys.npk(),
                    SharedSecretKey::new(&[56; 32], &recipient_keys.vpk()),
                ),
            ],
            vec![sender_keys.nsk],
            vec![],
            vec![Some((0, vec![]))],
            &program.into(),
        );

        assert!(matches!(result, Err(NssaError::CircuitProvingError(_))));
    }

    #[test]
    fn circuit_fails_if_insufficient_keys_are_provided() {
        let program = Program::simple_balance_transfer();
        let sender_keys = test_private_account_keys_1();
        let sender_account_id =
            AccountId::private_account_id(&sender_keys.npk(), Identifier(0_u128));
        let private_account_1 = AccountWithMetadata::new(
            Account {
                program_owner: program.id(),
                balance: 100,
                ..Account::default()
            },
            true,
            sender_account_id,
        );
        let private_account_2 =
            AccountWithMetadata::new(Account::default(), false, AccountId::new([1; 32]));

        // Setting only one key for an execution with two private accounts.
        let private_account_keys = [(
            sender_keys.npk(),
            SharedSecretKey::new(&[55; 32], &sender_keys.vpk()),
        )];
        let result = execute_and_prove(
            vec![private_account_1, private_account_2],
            Program::serialize_instruction(10_u128).unwrap(),
            vec![1, 2],
            private_account_keys.to_vec(),
            vec![sender_keys.nsk],
            vec![],
            vec![Some((0, vec![]))],
            &program.into(),
        );

        assert!(matches!(result, Err(NssaError::CircuitProvingError(_))));
    }

    #[test]
    fn circuit_fails_if_insufficient_commitment_proofs_are_provided() {
        let program = Program::simple_balance_transfer();
        let sender_keys = test_private_account_keys_1();
        let sender_account_id =
            AccountId::private_account_id(&sender_keys.npk(), Identifier(0_u128));
        let recipient_keys = test_private_account_keys_2();
        let recipient_account_id =
            AccountId::private_account_id(&recipient_keys.npk(), Identifier(0_u128));
        let private_account_1 = AccountWithMetadata::new(
            Account {
                program_owner: program.id(),
                balance: 100,
                ..Account::default()
            },
            true,
            sender_account_id,
        );
        let private_account_2 =
            AccountWithMetadata::new(Account::default(), false, recipient_account_id);

        // Setting no second commitment proof.
        let private_account_membership_proofs = [Some((0, vec![]))];
        let result = execute_and_prove(
            vec![private_account_1, private_account_2],
            Program::serialize_instruction(10_u128).unwrap(),
            vec![1, 2],
            vec![
                (
                    sender_keys.npk(),
                    SharedSecretKey::new(&[55; 32], &sender_keys.vpk()),
                ),
                (
                    recipient_keys.npk(),
                    SharedSecretKey::new(&[56; 32], &recipient_keys.vpk()),
                ),
            ],
            vec![sender_keys.nsk],
            vec![],
            private_account_membership_proofs.to_vec(),
            &program.into(),
        );

        assert!(matches!(result, Err(NssaError::CircuitProvingError(_))));
    }

    #[test]
    fn circuit_fails_if_insufficient_auth_keys_are_provided() {
        let program = Program::simple_balance_transfer();
        let sender_keys = test_private_account_keys_1();
        let sender_account_id =
            AccountId::private_account_id(&sender_keys.npk(), Identifier(0_u128));
        let recipient_keys = test_private_account_keys_2();
        let recipient_account_id =
            AccountId::private_account_id(&recipient_keys.npk(), Identifier(0_u128));
        let private_account_1 = AccountWithMetadata::new(
            Account {
                program_owner: program.id(),
                balance: 100,
                ..Account::default()
            },
            true,
            sender_account_id,
        );
        let private_account_2 =
            AccountWithMetadata::new(Account::default(), false, recipient_account_id);

        // Setting no auth key for an execution with one non default private accounts.
        let private_account_nsks = [];
        let result = execute_and_prove(
            vec![private_account_1, private_account_2],
            Program::serialize_instruction(10_u128).unwrap(),
            vec![1, 2],
            vec![
                (
                    sender_keys.npk(),
                    SharedSecretKey::new(&[55; 32], &sender_keys.vpk()),
                ),
                (
                    recipient_keys.npk(),
                    SharedSecretKey::new(&[56; 32], &recipient_keys.vpk()),
                ),
            ],
            private_account_nsks.to_vec(),
            vec![],
            vec![],
            &program.into(),
        );

        assert!(matches!(result, Err(NssaError::CircuitProvingError(_))));
    }

    #[test]
    fn circuit_fails_if_invalid_auth_keys_are_provided() {
        let program = Program::simple_balance_transfer();
        let sender_keys = test_private_account_keys_1();
        let sender_account_id =
            AccountId::private_account_id(&sender_keys.npk(), Identifier(0_u128));
        let recipient_keys = test_private_account_keys_2();
        let recipient_account_id =
            AccountId::private_account_id(&recipient_keys.npk(), Identifier(0_u128));
        let private_account_1 = AccountWithMetadata::new(
            Account {
                program_owner: program.id(),
                balance: 100,
                ..Account::default()
            },
            true,
            sender_account_id,
        );
        let private_account_2 =
            AccountWithMetadata::new(Account::default(), false, recipient_account_id);

        let private_account_keys = [
            // First private account is the sender
            (
                sender_keys.npk(),
                SharedSecretKey::new(&[55; 32], &sender_keys.vpk()),
            ),
            // Second private account is the recipient
            (
                recipient_keys.npk(),
                SharedSecretKey::new(&[56; 32], &recipient_keys.vpk()),
            ),
        ];

        // Setting the recipient key to authorize the sender.
        // This should be set to the sender private account in
        // a normal circumstance. The recipient can't authorize this.
        let private_account_nsks = [recipient_keys.nsk];
        let private_account_membership_proofs = [Some((0, vec![]))];
        let result = execute_and_prove(
            vec![private_account_1, private_account_2],
            Program::serialize_instruction(10_u128).unwrap(),
            vec![1, 2],
            private_account_keys.to_vec(),
            private_account_nsks.to_vec(),
            vec![],
            private_account_membership_proofs.to_vec(),
            &program.into(),
        );

        assert!(matches!(result, Err(NssaError::CircuitProvingError(_))));
    }

    #[test]
    fn circuit_should_fail_if_new_private_account_with_non_default_balance_is_provided() {
        let program = Program::simple_balance_transfer();
        let sender_keys = test_private_account_keys_1();
        let sender_account_id =
            AccountId::private_account_id(&sender_keys.npk(), Identifier(0_u128));
        let recipient_keys = test_private_account_keys_2();
        let recipient_account_id =
            AccountId::private_account_id(&recipient_keys.npk(), Identifier(0_u128));
        let private_account_1 = AccountWithMetadata::new(
            Account {
                program_owner: program.id(),
                balance: 100,
                ..Account::default()
            },
            true,
            sender_account_id,
        );
        let private_account_2 = AccountWithMetadata::new(
            Account {
                // Non default balance
                balance: 1,
                ..Account::default()
            },
            false,
            recipient_account_id,
        );

        let result = execute_and_prove(
            vec![private_account_1, private_account_2],
            Program::serialize_instruction(10_u128).unwrap(),
            vec![1, 2],
            vec![
                (
                    sender_keys.npk(),
                    SharedSecretKey::new(&[55; 32], &sender_keys.vpk()),
                ),
                (
                    recipient_keys.npk(),
                    SharedSecretKey::new(&[56; 32], &recipient_keys.vpk()),
                ),
            ],
            vec![sender_keys.nsk],
            vec![],
            vec![Some((0, vec![]))],
            &program.into(),
        );

        assert!(matches!(result, Err(NssaError::CircuitProvingError(_))));
    }

    #[test]
    fn circuit_should_fail_if_new_private_account_with_non_default_program_owner_is_provided() {
        let program = Program::simple_balance_transfer();
        let sender_keys = test_private_account_keys_1();
        let sender_account_id =
            AccountId::private_account_id(&sender_keys.npk(), Identifier(0_u128));
        let recipient_keys = test_private_account_keys_2();
        let recipient_account_id =
            AccountId::private_account_id(&recipient_keys.npk(), Identifier(0_u128));
        let private_account_1 = AccountWithMetadata::new(
            Account {
                program_owner: program.id(),
                balance: 100,
                ..Account::default()
            },
            true,
            sender_account_id,
        );
        let private_account_2 = AccountWithMetadata::new(
            Account {
                // Non default program_owner
                program_owner: [0, 1, 2, 3, 4, 5, 6, 7],
                ..Account::default()
            },
            false,
            recipient_account_id,
        );

        let result = execute_and_prove(
            vec![private_account_1, private_account_2],
            Program::serialize_instruction(10_u128).unwrap(),
            vec![1, 2],
            vec![
                (
                    sender_keys.npk(),
                    SharedSecretKey::new(&[55; 32], &sender_keys.vpk()),
                ),
                (
                    recipient_keys.npk(),
                    SharedSecretKey::new(&[56; 32], &recipient_keys.vpk()),
                ),
            ],
            vec![sender_keys.nsk],
            vec![],
            vec![Some((0, vec![]))],
            &program.into(),
        );

        assert!(matches!(result, Err(NssaError::CircuitProvingError(_))));
    }

    #[test]
    fn circuit_should_fail_if_new_private_account_with_non_default_data_is_provided() {
        let program = Program::simple_balance_transfer();
        let sender_keys = test_private_account_keys_1();
        let sender_account_id =
            AccountId::private_account_id(&sender_keys.npk(), Identifier(0_u128));
        let recipient_keys = test_private_account_keys_2();
        let recipient_account_id =
            AccountId::private_account_id(&recipient_keys.npk(), Identifier(0_u128));
        let private_account_1 = AccountWithMetadata::new(
            Account {
                program_owner: program.id(),
                balance: 100,
                ..Account::default()
            },
            true,
            sender_account_id,
        );
        let private_account_2 = AccountWithMetadata::new(
            Account {
                // Non default data
                data: b"hola mundo".to_vec().try_into().unwrap(),
                ..Account::default()
            },
            false,
            recipient_account_id,
        );

        let result = execute_and_prove(
            vec![private_account_1, private_account_2],
            Program::serialize_instruction(10_u128).unwrap(),
            vec![1, 2],
            vec![
                (
                    sender_keys.npk(),
                    SharedSecretKey::new(&[55; 32], &sender_keys.vpk()),
                ),
                (
                    recipient_keys.npk(),
                    SharedSecretKey::new(&[56; 32], &recipient_keys.vpk()),
                ),
            ],
            vec![sender_keys.nsk],
            vec![],
            vec![Some((0, vec![]))],
            &program.into(),
        );

        assert!(matches!(result, Err(NssaError::CircuitProvingError(_))));
    }

    #[test]
    fn circuit_should_fail_if_new_private_account_with_non_default_nonce_is_provided() {
        let program = Program::simple_balance_transfer();
        let sender_keys = test_private_account_keys_1();
        let sender_account_id =
            AccountId::private_account_id(&sender_keys.npk(), Identifier(0_u128));
        let recipient_keys = test_private_account_keys_2();
        let recipient_account_id =
            AccountId::private_account_id(&recipient_keys.npk(), Identifier(0_u128));
        let private_account_1 = AccountWithMetadata::new(
            Account {
                program_owner: program.id(),
                balance: 100,
                ..Account::default()
            },
            true,
            sender_account_id,
        );
        let private_account_2 = AccountWithMetadata::new(
            Account {
                // Non default nonce
                nonce: Nonce(0xdead_beef),
                ..Account::default()
            },
            false,
            recipient_account_id,
        );

        let result = execute_and_prove(
            vec![private_account_1, private_account_2],
            Program::serialize_instruction(10_u128).unwrap(),
            vec![1, 2],
            vec![
                (
                    sender_keys.npk(),
                    SharedSecretKey::new(&[55; 32], &sender_keys.vpk()),
                ),
                (
                    recipient_keys.npk(),
                    SharedSecretKey::new(&[56; 32], &recipient_keys.vpk()),
                ),
            ],
            vec![sender_keys.nsk],
            vec![],
            vec![Some((0, vec![]))],
            &program.into(),
        );

        assert!(matches!(result, Err(NssaError::CircuitProvingError(_))));
    }

    #[test]
    fn circuit_should_fail_if_new_private_account_is_provided_with_default_values_but_marked_as_authorized()
     {
        let program = Program::simple_balance_transfer();
        let sender_keys = test_private_account_keys_1();
        let sender_account_id =
            AccountId::private_account_id(&sender_keys.npk(), Identifier(0_u128));
        let recipient_keys = test_private_account_keys_2();
        let recipient_account_id =
            AccountId::private_account_id(&recipient_keys.npk(), Identifier(1_u128));
        let private_account_1 = AccountWithMetadata::new(
            Account {
                program_owner: program.id(),
                balance: 100,
                ..Account::default()
            },
            true,
            sender_account_id,
        );
        let private_account_2 = AccountWithMetadata::new(
            Account::default(),
            // This should be set to false in normal circumstances
            true,
            recipient_account_id,
        );

        let result = execute_and_prove(
            vec![private_account_1, private_account_2],
            Program::serialize_instruction(10_u128).unwrap(),
            vec![1, 2],
            vec![
                (
                    sender_keys.npk(),
                    SharedSecretKey::new(&[55; 32], &sender_keys.vpk()),
                ),
                (
                    recipient_keys.npk(),
                    SharedSecretKey::new(&[56; 32], &recipient_keys.vpk()),
                ),
            ],
            vec![sender_keys.nsk],
            vec![Identifier(0_u128), Identifier(1_u128)],
            vec![Some((0, vec![]))],
            &program.into(),
        );

        assert!(matches!(result, Err(NssaError::CircuitProvingError(_))));
    }

    #[test]
    fn circuit_should_fail_with_invalid_visibility_mask_value() {
        let program = Program::simple_balance_transfer();
        let public_account_1 = AccountWithMetadata::new(
            Account {
                program_owner: program.id(),
                balance: 100,
                ..Account::default()
            },
            true,
            AccountId::new([0; 32]),
        );
        let public_account_2 =
            AccountWithMetadata::new(Account::default(), false, AccountId::new([1; 32]));

        let visibility_mask = [0, 3];
        let result = execute_and_prove(
            vec![public_account_1, public_account_2],
            Program::serialize_instruction(10_u128).unwrap(),
            visibility_mask.to_vec(),
            vec![],
            vec![],
            vec![],
            vec![],
            &program.into(),
        );

        assert!(matches!(result, Err(NssaError::CircuitProvingError(_))));
    }

    #[test]
    fn circuit_should_fail_with_too_many_nonces() {
        let program = Program::simple_balance_transfer();
        let sender_keys = test_private_account_keys_1();
        let sender_account_id =
            AccountId::private_account_id(&sender_keys.npk(), Identifier(0_u128));
        let recipient_keys = test_private_account_keys_2();
        let recipient_account_id =
            AccountId::private_account_id(&recipient_keys.npk(), Identifier(0_u128));
        let private_account_1 = AccountWithMetadata::new(
            Account {
                program_owner: program.id(),
                balance: 100,
                ..Account::default()
            },
            true,
            sender_account_id,
        );
        let private_account_2 =
            AccountWithMetadata::new(Account::default(), false, recipient_account_id);

        let result = execute_and_prove(
            vec![private_account_1, private_account_2],
            Program::serialize_instruction(10_u128).unwrap(),
            vec![1, 2],
            vec![
                (
                    sender_keys.npk(),
                    SharedSecretKey::new(&[55; 32], &sender_keys.vpk()),
                ),
                (
                    recipient_keys.npk(),
                    SharedSecretKey::new(&[56; 32], &recipient_keys.vpk()),
                ),
            ],
            vec![sender_keys.nsk],
            vec![],
            vec![Some((0, vec![]))],
            &program.into(),
        );

        assert!(matches!(result, Err(NssaError::CircuitProvingError(_))));
    }

    #[test]
    fn circuit_should_fail_with_too_many_private_account_keys() {
        let program = Program::simple_balance_transfer();
        let sender_keys = test_private_account_keys_1();
        let sender_account_id =
            AccountId::private_account_id(&sender_keys.npk(), Identifier(0_u128));
        let recipient_keys = test_private_account_keys_2();
        let recipient_account_id =
            AccountId::private_account_id(&recipient_keys.npk(), Identifier(0_u128));
        let private_account_1 = AccountWithMetadata::new(
            Account {
                program_owner: program.id(),
                balance: 100,
                ..Account::default()
            },
            true,
            sender_account_id,
        );
        let private_account_2 =
            AccountWithMetadata::new(Account::default(), false, recipient_account_id);

        // Setting three private account keys for a circuit execution with only two private
        // accounts.
        let private_account_keys = [
            (
                sender_keys.npk(),
                SharedSecretKey::new(&[55; 32], &sender_keys.vpk()),
            ),
            (
                recipient_keys.npk(),
                SharedSecretKey::new(&[56; 32], &recipient_keys.vpk()),
            ),
            (
                sender_keys.npk(),
                SharedSecretKey::new(&[57; 32], &sender_keys.vpk()),
            ),
        ];
        let result = execute_and_prove(
            vec![private_account_1, private_account_2],
            Program::serialize_instruction(10_u128).unwrap(),
            vec![1, 2],
            private_account_keys.to_vec(),
            vec![sender_keys.nsk],
            vec![],
            vec![Some((0, vec![]))],
            &program.into(),
        );

        assert!(matches!(result, Err(NssaError::CircuitProvingError(_))));
    }

    #[test]
    fn circuit_should_fail_with_too_many_private_account_auth_keys() {
        let program = Program::simple_balance_transfer();
        let sender_keys = test_private_account_keys_1();
        let sender_account_id =
            AccountId::private_account_id(&sender_keys.npk(), Identifier(0_u128));
        let recipient_keys = test_private_account_keys_2();
        let recipient_account_id =
            AccountId::private_account_id(&recipient_keys.npk(), Identifier(0_u128));
        let private_account_1 = AccountWithMetadata::new(
            Account {
                program_owner: program.id(),
                balance: 100,
                ..Account::default()
            },
            true,
            sender_account_id,
        );
        let private_account_2 =
            AccountWithMetadata::new(Account::default(), false, recipient_account_id);

        // Setting two private account keys for a circuit execution with only one non default
        // private account (visibility mask equal to 1 means that auth keys are expected).
        let visibility_mask = [1, 2];
        let private_account_nsks = [sender_keys.nsk, recipient_keys.nsk];
        let private_account_membership_proofs = [Some((0, vec![])), Some((1, vec![]))];
        let result = execute_and_prove(
            vec![private_account_1, private_account_2],
            Program::serialize_instruction(10_u128).unwrap(),
            visibility_mask.to_vec(),
            vec![
                (
                    sender_keys.npk(),
                    SharedSecretKey::new(&[55; 32], &sender_keys.vpk()),
                ),
                (
                    recipient_keys.npk(),
                    SharedSecretKey::new(&[56; 32], &recipient_keys.vpk()),
                ),
            ],
            private_account_nsks.to_vec(),
            vec![],
            private_account_membership_proofs.to_vec(),
            &program.into(),
        );

        assert!(matches!(result, Err(NssaError::CircuitProvingError(_))));
    }

    #[test]
    fn private_accounts_can_only_be_initialized_once() {
        let sender_keys = test_private_account_keys_1();
        let sender_nonce = Nonce(0xdead_beef);

        let sender_private_account = Account {
            program_owner: Program::authenticated_transfer_program().id(),
            balance: 100,
            nonce: sender_nonce,
            data: Data::default(),
        };
        let recipient_keys = test_private_account_keys_2();

        let mut state = V03State::new_with_genesis_accounts(&[], &[], 0)
            .with_private_account(&sender_keys, &sender_private_account);

        let balance_to_move = 37;
        let balance_to_move_2 = 30;

        let tx = private_balance_transfer_for_tests(
            &sender_keys,
            &sender_private_account,
            &recipient_keys,
            balance_to_move,
            &state,
        );

        state
            .transition_from_privacy_preserving_transaction(&tx, 1, 0)
            .unwrap();

        let _sender_private_account = Account {
            program_owner: Program::authenticated_transfer_program().id(),
            balance: 100,
            nonce: sender_nonce,
            data: Data::default(),
        };

        let tx = private_balance_transfer_for_tests(
            &sender_keys,
            &sender_private_account,
            &recipient_keys,
            balance_to_move_2,
            &state,
        );

        let result = state.transition_from_privacy_preserving_transaction(&tx, 1, 0);

        assert!(matches!(result, Err(NssaError::InvalidInput(_))));
        let NssaError::InvalidInput(error_message) = result.err().unwrap() else {
            panic!("Incorrect message error");
        };
        let expected_error_message = "Nullifier already seen".to_owned();
        assert_eq!(error_message, expected_error_message);
    }

    #[test]
    fn circuit_should_fail_if_there_are_repeated_ids() {
        let program = Program::simple_balance_transfer();
        let sender_keys = test_private_account_keys_1();
        let sender_account_id =
            AccountId::private_account_id(&sender_keys.npk(), Identifier(0_u128));
        let private_account_1 = AccountWithMetadata::new(
            Account {
                program_owner: program.id(),
                balance: 100,
                ..Account::default()
            },
            true,
            sender_account_id,
        );

        let visibility_mask = [1, 1];
        let private_account_nsks = [sender_keys.nsk, sender_keys.nsk];
        let private_account_membership_proofs = [Some((1, vec![])), Some((1, vec![]))];
        let shared_secret = SharedSecretKey::new(&[55; 32], &sender_keys.vpk());
        let result = execute_and_prove(
            vec![private_account_1.clone(), private_account_1],
            Program::serialize_instruction(100_u128).unwrap(),
            visibility_mask.to_vec(),
            vec![
                (sender_keys.npk(), shared_secret),
                (sender_keys.npk(), shared_secret),
            ],
            private_account_nsks.to_vec(),
            vec![],
            private_account_membership_proofs.to_vec(),
            &program.into(),
        );

        assert!(matches!(result, Err(NssaError::CircuitProvingError(_))));
    }

    #[test]
    fn claiming_mechanism() {
        let program = Program::authenticated_transfer_program();
        let from_key = PrivateKey::try_new([1; 32]).unwrap();
        let from = AccountId::public_account_id(&PublicKey::new_from_private_key(&from_key));
        let initial_balance = 100;
        let initial_data = [(from, initial_balance)];
        let mut state =
            V03State::new_with_genesis_accounts(&initial_data, &[], 0).with_test_programs();
        let to_key = PrivateKey::try_new([2; 32]).unwrap();
        let to = AccountId::public_account_id(&PublicKey::new_from_private_key(&to_key));
        let amount: u128 = 37;

        // Check the recipient is an uninitialized account
        assert_eq!(state.get_account_by_id(to), Account::default());

        let expected_recipient_post = Account {
            program_owner: program.id(),
            balance: amount,
            nonce: Nonce(1),
            ..Account::default()
        };

        let message = public_transaction::Message::try_new(
            program.id(),
            vec![from, to],
            vec![Nonce(0), Nonce(0)],
            amount,
        )
        .unwrap();
        let witness_set =
            public_transaction::WitnessSet::for_message(&message, &[&from_key, &to_key]);
        let tx = PublicTransaction::new(message, witness_set);

        state.transition_from_public_transaction(&tx, 1, 0).unwrap();

        let recipient_post = state.get_account_by_id(to);

        assert_eq!(recipient_post, expected_recipient_post);
    }

    #[test]
    fn unauthorized_public_account_claiming_fails() {
        let program = Program::authenticated_transfer_program();
        let account_key = PrivateKey::try_new([9; 32]).unwrap();
        let account_id = AccountId::from(&PublicKey::new_from_private_key(&account_key));
        let mut state = V03State::new_with_genesis_accounts(&[], &[], 0);

        assert_eq!(state.get_account_by_id(account_id), Account::default());

        let message =
            public_transaction::Message::try_new(program.id(), vec![account_id], vec![], 0_u128)
                .unwrap();
        let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
        let tx = PublicTransaction::new(message, witness_set);

        let result = state.transition_from_public_transaction(&tx, 1, 0);

        assert!(matches!(result, Err(NssaError::ProgramExecutionFailed(_))));
        assert_eq!(state.get_account_by_id(account_id), Account::default());
    }

    #[test]
    fn authorized_public_account_claiming_succeeds() {
        let program = Program::authenticated_transfer_program();
        let account_key = PrivateKey::try_new([10; 32]).unwrap();
        let account_id = AccountId::from(&PublicKey::new_from_private_key(&account_key));
        let mut state = V03State::new_with_genesis_accounts(&[], &[], 0);

        assert_eq!(state.get_account_by_id(account_id), Account::default());

        let message = public_transaction::Message::try_new(
            program.id(),
            vec![account_id],
            vec![Nonce(0)],
            0_u128,
        )
        .unwrap();
        let witness_set = public_transaction::WitnessSet::for_message(&message, &[&account_key]);
        let tx = PublicTransaction::new(message, witness_set);

        state.transition_from_public_transaction(&tx, 1, 0).unwrap();

        assert_eq!(
            state.get_account_by_id(account_id),
            Account {
                program_owner: program.id(),
                nonce: Nonce(1),
                ..Account::default()
            }
        );
    }

    #[test]
    fn public_chained_call() {
        let program = Program::chain_caller();
        let key = PrivateKey::try_new([1; 32]).unwrap();
        let from = AccountId::public_account_id(&PublicKey::new_from_private_key(&key));
        let to = AccountId::new([2; 32]);
        let initial_balance = 1000;
        let initial_data = [(from, initial_balance), (to, 0)];
        let mut state =
            V03State::new_with_genesis_accounts(&initial_data, &[], 0).with_test_programs();
        let from_key = key;
        let amount: u128 = 37;
        let instruction: (u128, ProgramId, u32, Option<PdaSeed>) = (
            amount,
            Program::authenticated_transfer_program().id(),
            2,
            None,
        );

        let expected_to_post = Account {
            program_owner: Program::authenticated_transfer_program().id(),
            balance: amount * 2, // The `chain_caller` chains the program twice
            ..Account::default()
        };

        let message = public_transaction::Message::try_new(
            program.id(),
            vec![to, from], // The chain_caller program permutes the account order in the chain
            // call
            vec![Nonce(0)],
            instruction,
        )
        .unwrap();
        let witness_set = public_transaction::WitnessSet::for_message(&message, &[&from_key]);
        let tx = PublicTransaction::new(message, witness_set);

        state.transition_from_public_transaction(&tx, 1, 0).unwrap();

        let from_post = state.get_account_by_id(from);
        let to_post = state.get_account_by_id(to);
        // The `chain_caller` program calls the program twice
        assert_eq!(from_post.balance, initial_balance - 2 * amount);
        assert_eq!(to_post, expected_to_post);
    }

    #[test]
    fn execution_fails_if_chained_calls_exceeds_depth() {
        let program = Program::chain_caller();
        let key = PrivateKey::try_new([1; 32]).unwrap();
        let from = AccountId::public_account_id(&PublicKey::new_from_private_key(&key));
        let to = AccountId::new([2; 32]);
        let initial_balance = 100;
        let initial_data = [(from, initial_balance), (to, 0)];
        let mut state =
            V03State::new_with_genesis_accounts(&initial_data, &[], 0).with_test_programs();
        let from_key = key;
        let amount: u128 = 0;
        let instruction: (u128, ProgramId, u32, Option<PdaSeed>) = (
            amount,
            Program::authenticated_transfer_program().id(),
            u32::try_from(MAX_NUMBER_CHAINED_CALLS).expect("MAX_NUMBER_CHAINED_CALLS fits in u32")
                + 1,
            None,
        );

        let message = public_transaction::Message::try_new(
            program.id(),
            vec![to, from], // The chain_caller program permutes the account order in the chain
            // call
            vec![Nonce(0)],
            instruction,
        )
        .unwrap();
        let witness_set = public_transaction::WitnessSet::for_message(&message, &[&from_key]);
        let tx = PublicTransaction::new(message, witness_set);

        let result = state.transition_from_public_transaction(&tx, 1, 0);
        assert!(matches!(
            result,
            Err(NssaError::MaxChainedCallsDepthExceeded)
        ));
    }

    #[test]
    fn execution_that_requires_authentication_of_a_program_derived_account_id_succeeds() {
        let chain_caller = Program::chain_caller();
        let pda_seed = PdaSeed::new([37; 32]);
        let from = AccountId::from((&chain_caller.id(), &pda_seed));
        let to = AccountId::new([2; 32]);
        let initial_balance = 1000;
        let initial_data = [(from, initial_balance), (to, 0)];
        let mut state =
            V03State::new_with_genesis_accounts(&initial_data, &[], 0).with_test_programs();
        let amount: u128 = 58;
        let instruction: (u128, ProgramId, u32, Option<PdaSeed>) = (
            amount,
            Program::authenticated_transfer_program().id(),
            1,
            Some(pda_seed),
        );

        let expected_to_post = Account {
            program_owner: Program::authenticated_transfer_program().id(),
            balance: amount, // The `chain_caller` chains the program twice
            ..Account::default()
        };
        let message = public_transaction::Message::try_new(
            chain_caller.id(),
            vec![to, from], // The chain_caller program permutes the account order in the chain
            // call
            vec![],
            instruction,
        )
        .unwrap();
        let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
        let tx = PublicTransaction::new(message, witness_set);

        state.transition_from_public_transaction(&tx, 1, 0).unwrap();

        let from_post = state.get_account_by_id(from);
        let to_post = state.get_account_by_id(to);
        assert_eq!(from_post.balance, initial_balance - amount);
        assert_eq!(to_post, expected_to_post);
    }

    #[test]
    fn claiming_mechanism_within_chain_call() {
        // This test calls the authenticated transfer program through the chain_caller program.
        // The transfer is made from an initialized sender to an uninitialized recipient. And
        // it is expected that the recipient account is claimed by the authenticated transfer
        // program and not the chained_caller program.
        let chain_caller = Program::chain_caller();
        let auth_transfer = Program::authenticated_transfer_program();
        let from_key = PrivateKey::try_new([1; 32]).unwrap();
        let from = AccountId::public_account_id(&PublicKey::new_from_private_key(&from_key));
        let initial_balance = 100;
        let initial_data = [(from, initial_balance)];
        let mut state =
            V03State::new_with_genesis_accounts(&initial_data, &[], 0).with_test_programs();
        let to_key = PrivateKey::try_new([2; 32]).unwrap();
        let to = AccountId::public_account_id(&PublicKey::new_from_private_key(&to_key));
        let amount: u128 = 37;

        // Check the recipient is an uninitialized account
        assert_eq!(state.get_account_by_id(to), Account::default());

        let expected_to_post = Account {
            // The expected program owner is the authenticated transfer program
            program_owner: auth_transfer.id(),
            balance: amount,
            nonce: Nonce(1),
            ..Account::default()
        };

        // The transaction executes the chain_caller program, which internally calls the
        // authenticated_transfer program
        let instruction: (u128, ProgramId, u32, Option<PdaSeed>) = (
            amount,
            Program::authenticated_transfer_program().id(),
            1,
            None,
        );
        let message = public_transaction::Message::try_new(
            chain_caller.id(),
            vec![to, from], // The chain_caller program permutes the account order in the chain
            // call
            vec![Nonce(0), Nonce(0)],
            instruction,
        )
        .unwrap();
        let witness_set =
            public_transaction::WitnessSet::for_message(&message, &[&from_key, &to_key]);
        let tx = PublicTransaction::new(message, witness_set);

        state.transition_from_public_transaction(&tx, 1, 0).unwrap();

        let from_post = state.get_account_by_id(from);
        let to_post = state.get_account_by_id(to);
        assert_eq!(from_post.balance, initial_balance - amount);
        assert_eq!(to_post, expected_to_post);
    }

    #[test]
    fn unauthorized_public_account_claiming_fails_when_executed_privately() {
        let program = Program::authenticated_transfer_program();
        let account_id = AccountId::new([11; 32]);
        let public_account = AccountWithMetadata::new(Account::default(), false, account_id);

        let result = execute_and_prove(
            vec![public_account],
            Program::serialize_instruction(0_u128).unwrap(),
            vec![0],
            vec![],
            vec![],
            vec![],
            vec![],
            &program.into(),
        );

        assert!(matches!(result, Err(NssaError::ProgramProveFailed(_))));
    }

    #[test]
    fn authorized_public_account_claiming_succeeds_when_executed_privately() {
        let program = Program::authenticated_transfer_program();
        let program_id = program.id();
        let sender_keys = test_private_account_keys_1();
        let sender_account_id =
            AccountId::private_account_id(&sender_keys.npk(), sender_keys.identifier);
        let sender_private_account = Account {
            program_owner: program_id,
            balance: 100,
            ..Account::default()
        };
        let sender_commitment = Commitment::new(&sender_account_id, &sender_private_account);
        let mut state =
            V03State::new_with_genesis_accounts(&[], std::slice::from_ref(&sender_commitment), 0);
        let sender_pre = AccountWithMetadata::new(sender_private_account, true, sender_account_id);
        let recipient_private_key = PrivateKey::try_new([2; 32]).unwrap();
        let recipient_account_id =
            AccountId::public_account_id(&PublicKey::new_from_private_key(&recipient_private_key));
        let recipient_pre =
            AccountWithMetadata::new(Account::default(), true, recipient_account_id);

        let esk = [5; 32];
        let shared_secret = SharedSecretKey::new(&esk, &sender_keys.vpk());
        let epk = EphemeralPublicKey::from_scalar(esk);

        let (output, proof) = execute_and_prove(
            vec![sender_pre, recipient_pre],
            Program::serialize_instruction(37_u128).unwrap(),
            vec![1, 0],
            vec![(sender_keys.npk(), shared_secret)],
            vec![sender_keys.nsk],
            vec![sender_keys.identifier],
            vec![state.get_proof_for_commitment(&sender_commitment)],
            &program.into(),
        )
        .unwrap();

        let message = Message::try_from_circuit_output(
            vec![recipient_account_id],
            vec![Nonce(0)],
            vec![(sender_account_id, sender_keys.vpk(), epk)],
            output,
        )
        .unwrap();

        let witness_set = WitnessSet::for_message(&message, proof, &[&recipient_private_key]);
        let tx = PrivacyPreservingTransaction::new(message, witness_set);

        state
            .transition_from_privacy_preserving_transaction(&tx, 1, 0)
            .unwrap();

        let nullifier = Nullifier::for_account_update(&sender_commitment, &sender_keys.nsk);
        assert!(state.private_state.1.contains(&nullifier));

        assert_eq!(
            state.get_account_by_id(recipient_account_id),
            Account {
                program_owner: program_id,
                balance: 37,
                nonce: Nonce(1),
                ..Account::default()
            }
        );
    }

    // TODO: Marvin check this
    #[test_case::test_case(1; "single call")]
    #[test_case::test_case(2; "two calls")]
    fn private_chained_call(number_of_calls: u32) {
        // Arrange
        let chain_caller = Program::chain_caller();
        let auth_transfers = Program::authenticated_transfer_program();
        let from_keys = test_private_account_keys_1();
        let to_keys = test_private_account_keys_2();
        let initial_balance = 100;
        let from_account_id = AccountId::private_account_id(&from_keys.npk(), from_keys.identifier);
        let from_account = AccountWithMetadata::new(
            Account {
                program_owner: auth_transfers.id(),
                balance: initial_balance,
                ..Account::default()
            },
            true,
            from_account_id,
        );
        let to_account_id = AccountId::private_account_id(&to_keys.npk(), to_keys.identifier);
        let to_account = AccountWithMetadata::new(
            Account {
                program_owner: auth_transfers.id(),
                ..Account::default()
            },
            true,
            to_account_id,
        );

        let from_commitment = Commitment::new(&from_account_id, &from_account.account);
        let to_commitment = Commitment::new(&to_account_id, &to_account.account);
        let mut state = V03State::new_with_genesis_accounts(
            &[],
            &[from_commitment.clone(), to_commitment.clone()],
            0,
        )
        .with_test_programs();
        let amount: u128 = 37;
        let instruction: (u128, ProgramId, u32, Option<PdaSeed>) = (
            amount,
            Program::authenticated_transfer_program().id(),
            number_of_calls,
            None,
        );

        let from_esk = [3; 32];
        let from_ss = SharedSecretKey::new(&from_esk, &from_keys.vpk());
        let from_epk = EphemeralPublicKey::from_scalar(from_esk);

        let to_esk = [3; 32];
        let to_ss = SharedSecretKey::new(&to_esk, &to_keys.vpk());
        let to_epk = EphemeralPublicKey::from_scalar(to_esk);

        let mut dependencies = HashMap::new();

        dependencies.insert(auth_transfers.id(), auth_transfers);
        let program_with_deps = ProgramWithDependencies::new(chain_caller, dependencies);

        let from_new_nonce = Nonce::default().private_account_nonce_increment(&from_keys.nsk);
        let to_new_nonce = Nonce::default().private_account_nonce_increment(&to_keys.nsk);

        let from_expected_post = Account {
            balance: initial_balance - u128::from(number_of_calls) * amount,
            nonce: from_new_nonce,
            ..from_account.account.clone()
        };
        let from_expected_commitment = Commitment::new(&from_account_id, &from_expected_post);

        let to_expected_post = Account {
            balance: u128::from(number_of_calls) * amount,
            nonce: to_new_nonce,
            ..to_account.account.clone()
        };
        let to_expected_commitment = Commitment::new(&to_account_id, &to_expected_post);

        // Act
        let (output, proof) = execute_and_prove(
            vec![to_account, from_account],
            Program::serialize_instruction(instruction).unwrap(),
            vec![1, 1],
            vec![(from_keys.npk(), to_ss), (to_keys.npk(), from_ss)],
            vec![from_keys.nsk, to_keys.nsk],
            vec![from_keys.identifier, to_keys.identifier],
            vec![
                state.get_proof_for_commitment(&from_commitment),
                state.get_proof_for_commitment(&to_commitment),
            ],
            &program_with_deps,
        )
        .unwrap();

        let message = Message::try_from_circuit_output(
            vec![],
            vec![],
            vec![
                (to_account_id, to_keys.vpk(), to_epk),
                (from_account_id, from_keys.vpk(), from_epk),
            ],
            output,
        )
        .unwrap();
        let witness_set = WitnessSet::for_message(&message, proof, &[]);
        let transaction = PrivacyPreservingTransaction::new(message, witness_set);

        state
            .transition_from_privacy_preserving_transaction(&transaction, 1, 0)
            .unwrap();

        // Assert
        assert!(
            state
                .get_proof_for_commitment(&from_expected_commitment)
                .is_some()
        );
        assert!(
            state
                .get_proof_for_commitment(&to_expected_commitment)
                .is_some()
        );
    }

    #[test]
    fn pda_mechanism_with_pinata_token_program() {
        let pinata_token = Program::pinata_token();
        let token = Program::token();

        let pinata_definition_id = AccountId::new([1; 32]);
        let pinata_token_definition_id = AccountId::new([2; 32]);
        // Total supply of pinata token will be in an account under a PDA.
        let pinata_token_holding_id = AccountId::from((&pinata_token.id(), &PdaSeed::new([0; 32])));
        let winner_token_holding_id = AccountId::new([3; 32]);

        let expected_winner_account_holding = token_core::TokenHolding::Fungible {
            definition_id: pinata_token_definition_id,
            balance: 150,
        };
        let expected_winner_token_holding_post = Account {
            program_owner: token.id(),
            data: Data::from(&expected_winner_account_holding),
            ..Account::default()
        };

        let mut state = V03State::new_with_genesis_accounts(&[], &[], 0);
        state.add_pinata_token_program(pinata_definition_id);

        // Set up the token accounts directly (bypassing public transactions which
        // would require signers for Claim::Authorized). The focus of this test is
        // the PDA mechanism in the pinata program's chained call, not token creation.
        let total_supply: u128 = 10_000_000;
        let token_definition = token_core::TokenDefinition::Fungible {
            name: String::from("PINATA"),
            total_supply,
            metadata_id: None,
        };
        let token_holding = token_core::TokenHolding::Fungible {
            definition_id: pinata_token_definition_id,
            balance: total_supply,
        };
        let winner_holding = token_core::TokenHolding::Fungible {
            definition_id: pinata_token_definition_id,
            balance: 0,
        };
        state.force_insert_account(
            pinata_token_definition_id,
            Account {
                program_owner: token.id(),
                data: Data::from(&token_definition),
                ..Account::default()
            },
        );
        state.force_insert_account(
            pinata_token_holding_id,
            Account {
                program_owner: token.id(),
                data: Data::from(&token_holding),
                ..Account::default()
            },
        );
        state.force_insert_account(
            winner_token_holding_id,
            Account {
                program_owner: token.id(),
                data: Data::from(&winner_holding),
                ..Account::default()
            },
        );

        // Submit a solution to the pinata program to claim the prize
        let solution: u128 = 989_106;
        let message = public_transaction::Message::try_new(
            pinata_token.id(),
            vec![
                pinata_definition_id,
                pinata_token_holding_id,
                winner_token_holding_id,
            ],
            vec![],
            solution,
        )
        .unwrap();
        let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
        let tx = PublicTransaction::new(message, witness_set);
        state.transition_from_public_transaction(&tx, 1, 0).unwrap();

        let winner_token_holding_post = state.get_account_by_id(winner_token_holding_id);
        assert_eq!(
            winner_token_holding_post,
            expected_winner_token_holding_post
        );
    }

    #[test]
    fn claiming_mechanism_cannot_claim_initialied_accounts() {
        let claimer = Program::claimer();
        let mut state = V03State::new_with_genesis_accounts(&[], &[], 0).with_test_programs();
        let account_id = AccountId::new([2; 32]);

        // Insert an account with non-default program owner
        state.force_insert_account(
            account_id,
            Account {
                program_owner: [1, 2, 3, 4, 5, 6, 7, 8],
                ..Account::default()
            },
        );

        let message =
            public_transaction::Message::try_new(claimer.id(), vec![account_id], vec![], ())
                .unwrap();
        let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
        let tx = PublicTransaction::new(message, witness_set);

        let result = state.transition_from_public_transaction(&tx, 1, 0);

        assert!(matches!(result, Err(NssaError::InvalidProgramBehavior)));
    }

    /// This test ensures that even if a malicious program tries to perform overflow of balances
    /// it will not be able to break the balance validation.
    #[test]
    fn malicious_program_cannot_break_balance_validation() {
        let sender_key = PrivateKey::try_new([37; 32]).unwrap();
        let sender_id = AccountId::public_account_id(&PublicKey::new_from_private_key(&sender_key));
        let sender_init_balance: u128 = 10;

        let recipient_key = PrivateKey::try_new([42; 32]).unwrap();
        let recipient_id =
            AccountId::public_account_id(&PublicKey::new_from_private_key(&recipient_key));
        let recipient_init_balance: u128 = 10;

        let mut state = V03State::new_with_genesis_accounts(
            &[
                (sender_id, sender_init_balance),
                (recipient_id, recipient_init_balance),
            ],
            &[],
            0,
        );

        state.insert_program(Program::modified_transfer_program());

        let balance_to_move: u128 = 4;

        let sender = AccountWithMetadata::new(state.get_account_by_id(sender_id), true, sender_id);

        let sender_nonce = sender.account.nonce;

        let _recipient =
            AccountWithMetadata::new(state.get_account_by_id(recipient_id), false, sender_id);

        let message = public_transaction::Message::try_new(
            Program::modified_transfer_program().id(),
            vec![sender_id, recipient_id],
            vec![sender_nonce],
            balance_to_move,
        )
        .unwrap();

        let witness_set = public_transaction::WitnessSet::for_message(&message, &[&sender_key]);
        let tx = PublicTransaction::new(message, witness_set);
        let res = state.transition_from_public_transaction(&tx, 1, 0);
        assert!(matches!(res, Err(NssaError::InvalidProgramBehavior)));

        let sender_post = state.get_account_by_id(sender_id);
        let recipient_post = state.get_account_by_id(recipient_id);

        let expected_sender_post = {
            let mut this = state.get_account_by_id(sender_id);
            this.balance = sender_init_balance;
            this.nonce = Nonce(0);
            this
        };

        let expected_recipient_post = {
            let mut this = state.get_account_by_id(sender_id);
            this.balance = recipient_init_balance;
            this.nonce = Nonce(0);
            this
        };

        assert_eq!(expected_sender_post, sender_post);
        assert_eq!(expected_recipient_post, recipient_post);
    }

    #[test]
    fn private_authorized_uninitialized_account() {
        let mut state = V03State::new_with_genesis_accounts(&[], &[], 0);

        // Set up keys for the authorized private account
        let private_keys = test_private_account_keys_1();
        let account_id =
            AccountId::private_account_id(&private_keys.npk(), private_keys.identifier);

        // Create an authorized private account with default values (new account being initialized)
        let authorized_account = AccountWithMetadata::new(Account::default(), true, account_id);

        let program = Program::authenticated_transfer_program();

        // Set up parameters for the new account
        let esk = [3; 32];
        let shared_secret = SharedSecretKey::new(&esk, &private_keys.vpk());
        let epk = EphemeralPublicKey::from_scalar(esk);

        // Balance to initialize the account with (0 for a new account)
        let balance: u128 = 0;

        // Execute and prove the circuit with the authorized account but no commitment proof
        let (output, proof) = execute_and_prove(
            vec![authorized_account],
            Program::serialize_instruction(balance).unwrap(),
            vec![1],
            vec![(private_keys.npk(), shared_secret)],
            vec![private_keys.nsk],
            vec![private_keys.identifier],
            vec![None],
            &program.into(),
        )
        .unwrap();

        // Create message from circuit output
        let message = Message::try_from_circuit_output(
            vec![],
            vec![],
            vec![(account_id, private_keys.vpk(), epk)],
            output,
        )
        .unwrap();

        let witness_set = WitnessSet::for_message(&message, proof, &[]);

        let tx = PrivacyPreservingTransaction::new(message, witness_set);
        let result = state.transition_from_privacy_preserving_transaction(&tx, 1, 0);
        assert!(result.is_ok());

        let nullifier = Nullifier::for_account_initialization(&account_id);
        assert!(state.private_state.1.contains(&nullifier));
    }

    #[test]
    fn private_unauthorized_uninitialized_account_can_still_be_claimed() {
        let mut state = V03State::new_with_genesis_accounts(&[], &[], 0).with_test_programs();

        let private_keys = test_private_account_keys_1();
        let account_id =
            AccountId::private_account_id(&private_keys.npk(), private_keys.identifier);
        // This is intentional: claim authorization was introduced to protect public accounts,
        // especially PDAs. Private PDAs are not useful in practice because there is no way to
        // operate them without the corresponding private keys, so unauthorized private claiming
        // remains allowed.
        let unauthorized_account = AccountWithMetadata::new(Account::default(), false, account_id);

        let program = Program::claimer();
        let esk = [5; 32];
        let shared_secret = SharedSecretKey::new(&esk, &private_keys.vpk());
        let epk = EphemeralPublicKey::from_scalar(esk);

        let (output, proof) = execute_and_prove(
            vec![unauthorized_account],
            Program::serialize_instruction(0_u128).unwrap(),
            vec![2],
            vec![(private_keys.npk(), shared_secret)],
            vec![],
            vec![private_keys.identifier],
            vec![None],
            &program.into(),
        )
        .unwrap();

        let message = Message::try_from_circuit_output(
            vec![],
            vec![],
            vec![(account_id, private_keys.vpk(), epk)],
            output,
        )
        .unwrap();

        let witness_set = WitnessSet::for_message(&message, proof, &[]);
        let tx = PrivacyPreservingTransaction::new(message, witness_set);

        state
            .transition_from_privacy_preserving_transaction(&tx, 1, 0)
            .unwrap();

        let nullifier = Nullifier::for_account_initialization(&account_id);
        assert!(state.private_state.1.contains(&nullifier));
    }

    #[test]
    fn private_account_claimed_then_used_without_init_flag_should_fail() {
        let mut state = V03State::new_with_genesis_accounts(&[], &[], 0).with_test_programs();

        // Set up keys for the private account
        let private_keys = test_private_account_keys_1();
        let account_id =
            AccountId::private_account_id(&private_keys.npk(), private_keys.identifier);

        // Step 1: Create a new private account with authorization
        let authorized_account = AccountWithMetadata::new(Account::default(), true, account_id);

        let claimer_program = Program::claimer();

        // Set up parameters for claiming the new account
        let esk = [3; 32];
        let shared_secret = SharedSecretKey::new(&esk, &private_keys.vpk());
        let epk = EphemeralPublicKey::from_scalar(esk);

        let balance: u128 = 0;

        // Step 2: Execute claimer program to claim the account with authentication
        let (output, proof) = execute_and_prove(
            vec![authorized_account.clone()],
            Program::serialize_instruction(balance).unwrap(),
            vec![1],
            vec![(private_keys.npk(), shared_secret)],
            vec![private_keys.nsk],
            vec![private_keys.identifier],
            vec![None],
            &claimer_program.into(),
        )
        .unwrap();

        let message = Message::try_from_circuit_output(
            vec![],
            vec![],
            vec![(account_id, private_keys.vpk(), epk)],
            output,
        )
        .unwrap();

        let witness_set = WitnessSet::for_message(&message, proof, &[]);
        let tx = PrivacyPreservingTransaction::new(message, witness_set);

        // Claim should succeed
        assert!(
            state
                .transition_from_privacy_preserving_transaction(&tx, 1, 0)
                .is_ok()
        );

        // Verify the account is now initialized (nullifier exists)
        let nullifier = Nullifier::for_account_initialization(&account_id);
        assert!(state.private_state.1.contains(&nullifier));

        // Prepare new state of account
        let account_metadata = {
            let mut acc = authorized_account;
            acc.account.program_owner = Program::claimer().id();
            acc
        };

        let noop_program = Program::noop();
        let esk2 = [4; 32];
        let shared_secret2 = SharedSecretKey::new(&esk2, &private_keys.vpk());

        // Step 3: Try to execute noop program with authentication but without initialization
        let res = execute_and_prove(
            vec![account_metadata],
            Program::serialize_instruction(()).unwrap(),
            vec![1],
            vec![(private_keys.npk(), shared_secret2)],
            vec![private_keys.nsk],
            vec![],
            vec![None],
            &noop_program.into(),
        );

        assert!(matches!(res, Err(NssaError::CircuitProvingError(_))));
    }

    #[test]
    fn public_changer_claimer_no_data_change_no_claim_succeeds() {
        let initial_data = [];
        let mut state =
            V03State::new_with_genesis_accounts(&initial_data, &[], 0).with_test_programs();
        let account_id = AccountId::new([1; 32]);
        let program_id = Program::changer_claimer().id();
        // Don't change data (None) and don't claim (false)
        let instruction: (Option<Vec<u8>>, bool) = (None, false);

        let message =
            public_transaction::Message::try_new(program_id, vec![account_id], vec![], instruction)
                .unwrap();
        let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
        let tx = PublicTransaction::new(message, witness_set);

        let result = state.transition_from_public_transaction(&tx, 1, 0);

        // Should succeed - no changes made, no claim needed
        assert!(result.is_ok());
        // Account should remain default/unclaimed
        assert_eq!(state.get_account_by_id(account_id), Account::default());
    }

    #[test]
    fn public_changer_claimer_data_change_no_claim_fails() {
        let initial_data = [];
        let mut state =
            V03State::new_with_genesis_accounts(&initial_data, &[], 0).with_test_programs();
        let account_id = AccountId::new([1; 32]);
        let program_id = Program::changer_claimer().id();
        // Change data but don't claim (false) - should fail
        let new_data = vec![1, 2, 3, 4, 5];
        let instruction: (Option<Vec<u8>>, bool) = (Some(new_data), false);

        let message =
            public_transaction::Message::try_new(program_id, vec![account_id], vec![], instruction)
                .unwrap();
        let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
        let tx = PublicTransaction::new(message, witness_set);

        let result = state.transition_from_public_transaction(&tx, 1, 0);

        // Should fail - cannot modify data without claiming the account
        assert!(matches!(result, Err(NssaError::InvalidProgramBehavior)));
    }

    #[test]
    fn private_changer_claimer_no_data_change_no_claim_succeeds() {
        let program = Program::changer_claimer();
        let sender_keys = test_private_account_keys_1();
        let sender_id = AccountId::private_account_id(&sender_keys.npk(), sender_keys.identifier);
        let private_account = AccountWithMetadata::new(Account::default(), true, sender_id);
        // Don't change data (None) and don't claim (false)
        let instruction: (Option<Vec<u8>>, bool) = (None, false);

        let result = execute_and_prove(
            vec![private_account],
            Program::serialize_instruction(instruction).unwrap(),
            vec![1],
            vec![(
                sender_keys.npk(),
                SharedSecretKey::new(&[3; 32], &sender_keys.vpk()),
            )],
            vec![sender_keys.nsk],
            vec![sender_keys.identifier],
            vec![Some((0, vec![]))],
            &program.into(),
        );

        // Should succeed - no changes made, no claim needed
        assert!(result.is_ok());
    }

    #[test]
    fn private_changer_claimer_data_change_no_claim_fails() {
        let program = Program::changer_claimer();
        let sender_keys = test_private_account_keys_1();
        let sender_id = AccountId::private_account_id(&sender_keys.npk(), Identifier(0_u128));
        let private_account = AccountWithMetadata::new(Account::default(), true, sender_id);
        // Change data but don't claim (false) - should fail
        let new_data = vec![1, 2, 3, 4, 5];
        let instruction: (Option<Vec<u8>>, bool) = (Some(new_data), false);

        let result = execute_and_prove(
            vec![private_account],
            Program::serialize_instruction(instruction).unwrap(),
            vec![1],
            vec![(
                sender_keys.npk(),
                SharedSecretKey::new(&[3; 32], &sender_keys.vpk()),
            )],
            vec![sender_keys.nsk],
            vec![],
            vec![Some((0, vec![]))],
            &program.into(),
        );

        // Should fail - cannot modify data without claiming the account
        assert!(matches!(result, Err(NssaError::CircuitProvingError(_))));
    }

    #[test]
    fn malicious_authorization_changer_should_fail_in_privacy_preserving_circuit() {
        // Arrange
        let malicious_program = Program::malicious_authorization_changer();
        let auth_transfers = Program::authenticated_transfer_program();
        let sender_keys = test_public_account_keys_1();
        let recipient_keys = test_private_account_keys_1();
        let recipient_account_id =
            AccountId::private_account_id(&recipient_keys.npk(), Identifier(0_u128));

        let sender_account = AccountWithMetadata::new(
            Account {
                program_owner: auth_transfers.id(),
                balance: 100,
                ..Default::default()
            },
            false,
            sender_keys.account_id(),
        );
        let recipient_account =
            AccountWithMetadata::new(Account::default(), true, recipient_account_id);
        let recipient_commitment =
            Commitment::new(&recipient_account_id, &recipient_account.account);
        let state = V03State::new_with_genesis_accounts(
            &[(sender_account.account_id, sender_account.account.balance)],
            std::slice::from_ref(&recipient_commitment),
            0,
        )
        .with_test_programs();

        let balance_to_transfer = 10_u128;
        let instruction = (balance_to_transfer, auth_transfers.id());

        let recipient_esk = [3; 32];
        let recipient = SharedSecretKey::new(&recipient_esk, &recipient_keys.vpk());

        let mut dependencies = HashMap::new();
        dependencies.insert(auth_transfers.id(), auth_transfers);
        let program_with_deps = ProgramWithDependencies::new(malicious_program, dependencies);

        // Act - execute the malicious program - this should fail during proving
        let result = execute_and_prove(
            vec![sender_account, recipient_account],
            Program::serialize_instruction(instruction).unwrap(),
            vec![0, 1],
            vec![(recipient_keys.npk(), recipient)],
            vec![recipient_keys.nsk],
            vec![],
            vec![state.get_proof_for_commitment(&recipient_commitment)],
            &program_with_deps,
        );

        // Assert - should fail because the malicious program tries to manipulate is_authorized
        assert!(matches!(result, Err(NssaError::CircuitProvingError(_))));
    }

    #[test_case::test_case((Some(1), Some(3)), 3; "at upper bound")]
    #[test_case::test_case((Some(1), Some(3)), 2; "inside range")]
    #[test_case::test_case((Some(1), Some(3)), 0; "below range")]
    #[test_case::test_case((Some(1), Some(3)), 1; "at lower bound")]
    #[test_case::test_case((Some(1), Some(3)), 4; "above range")]
    #[test_case::test_case((Some(1), None), 1; "lower bound only - at bound")]
    #[test_case::test_case((Some(1), None), 10; "lower bound only - above")]
    #[test_case::test_case((Some(1), None), 0; "lower bound only - below")]
    #[test_case::test_case((None, Some(3)), 3; "upper bound only - at bound")]
    #[test_case::test_case((None, Some(3)), 0; "upper bound only - below")]
    #[test_case::test_case((None, Some(3)), 4; "upper bound only - above")]
    #[test_case::test_case((None, None), 0; "no bounds - always valid")]
    #[test_case::test_case((None, None), 100; "no bounds - always valid 2")]
    fn validity_window_works_in_public_transactions(
        validity_window: (Option<BlockId>, Option<BlockId>),
        block_id: BlockId,
    ) {
        let block_validity_window: BlockValidityWindow = validity_window.try_into().unwrap();
        let validity_window_program = Program::validity_window();
        let account_keys = test_public_account_keys_1();
        let pre = AccountWithMetadata::new(Account::default(), false, account_keys.account_id());
        let mut state = V03State::new_with_genesis_accounts(&[], &[], 0).with_test_programs();
        let tx = {
            let account_ids = vec![pre.account_id];
            let nonces = vec![];
            let program_id = validity_window_program.id();
            let instruction = (
                block_validity_window,
                TimestampValidityWindow::new_unbounded(),
            );
            let message =
                public_transaction::Message::try_new(program_id, account_ids, nonces, instruction)
                    .unwrap();
            let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
            PublicTransaction::new(message, witness_set)
        };
        let result = state.transition_from_public_transaction(&tx, block_id, 0);
        let is_inside_validity_window =
            match (block_validity_window.start(), block_validity_window.end()) {
                (Some(s), Some(e)) => s <= block_id && block_id < e,
                (Some(s), None) => s <= block_id,
                (None, Some(e)) => block_id < e,
                (None, None) => true,
            };
        if is_inside_validity_window {
            assert!(result.is_ok());
        } else {
            assert!(matches!(result, Err(NssaError::OutOfValidityWindow)));
        }
    }

    #[test_case::test_case((Some(1), Some(3)), 3; "at upper bound")]
    #[test_case::test_case((Some(1), Some(3)), 2; "inside range")]
    #[test_case::test_case((Some(1), Some(3)), 0; "below range")]
    #[test_case::test_case((Some(1), Some(3)), 1; "at lower bound")]
    #[test_case::test_case((Some(1), Some(3)), 4; "above range")]
    #[test_case::test_case((Some(1), None), 1; "lower bound only - at bound")]
    #[test_case::test_case((Some(1), None), 10; "lower bound only - above")]
    #[test_case::test_case((Some(1), None), 0; "lower bound only - below")]
    #[test_case::test_case((None, Some(3)), 3; "upper bound only - at bound")]
    #[test_case::test_case((None, Some(3)), 0; "upper bound only - below")]
    #[test_case::test_case((None, Some(3)), 4; "upper bound only - above")]
    #[test_case::test_case((None, None), 0; "no bounds - always valid")]
    #[test_case::test_case((None, None), 100; "no bounds - always valid 2")]
    fn timestamp_validity_window_works_in_public_transactions(
        validity_window: (Option<Timestamp>, Option<Timestamp>),
        timestamp: Timestamp,
    ) {
        let timestamp_validity_window: TimestampValidityWindow =
            validity_window.try_into().unwrap();
        let validity_window_program = Program::validity_window();
        let account_keys = test_public_account_keys_1();
        let pre = AccountWithMetadata::new(Account::default(), false, account_keys.account_id());
        let mut state = V03State::new_with_genesis_accounts(&[], &[], 0).with_test_programs();
        let tx = {
            let account_ids = vec![pre.account_id];
            let nonces = vec![];
            let program_id = validity_window_program.id();
            let instruction = (
                BlockValidityWindow::new_unbounded(),
                timestamp_validity_window,
            );
            let message =
                public_transaction::Message::try_new(program_id, account_ids, nonces, instruction)
                    .unwrap();
            let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
            PublicTransaction::new(message, witness_set)
        };
        let result = state.transition_from_public_transaction(&tx, 1, timestamp);
        let is_inside_validity_window = match (
            timestamp_validity_window.start(),
            timestamp_validity_window.end(),
        ) {
            (Some(s), Some(e)) => s <= timestamp && timestamp < e,
            (Some(s), None) => s <= timestamp,
            (None, Some(e)) => timestamp < e,
            (None, None) => true,
        };
        if is_inside_validity_window {
            assert!(result.is_ok());
        } else {
            assert!(matches!(result, Err(NssaError::OutOfValidityWindow)));
        }
    }

    #[test_case::test_case((Some(1), Some(3)), 3; "at upper bound")]
    #[test_case::test_case((Some(1), Some(3)), 2; "inside range")]
    #[test_case::test_case((Some(1), Some(3)), 0; "below range")]
    #[test_case::test_case((Some(1), Some(3)), 1; "at lower bound")]
    #[test_case::test_case((Some(1), Some(3)), 4; "above range")]
    #[test_case::test_case((Some(1), None), 1; "lower bound only - at bound")]
    #[test_case::test_case((Some(1), None), 10; "lower bound only - above")]
    #[test_case::test_case((Some(1), None), 0; "lower bound only - below")]
    #[test_case::test_case((None, Some(3)), 3; "upper bound only - at bound")]
    #[test_case::test_case((None, Some(3)), 0; "upper bound only - below")]
    #[test_case::test_case((None, Some(3)), 4; "upper bound only - above")]
    #[test_case::test_case((None, None), 0; "no bounds - always valid")]
    #[test_case::test_case((None, None), 100; "no bounds - always valid 2")]
    fn validity_window_works_in_privacy_preserving_transactions(
        validity_window: (Option<BlockId>, Option<BlockId>),
        block_id: BlockId,
    ) {
        let block_validity_window: BlockValidityWindow = validity_window.try_into().unwrap();
        let validity_window_program = Program::validity_window();
        let account_keys = test_private_account_keys_1();
        let account_id =
            AccountId::private_account_id(&account_keys.npk(), account_keys.identifier);
        let pre = AccountWithMetadata::new(Account::default(), false, account_id);
        let mut state = V03State::new_with_genesis_accounts(&[], &[], 0).with_test_programs();
        let tx = {
            let esk = [3; 32];
            let shared_secret = SharedSecretKey::new(&esk, &account_keys.vpk());
            let epk = EphemeralPublicKey::from_scalar(esk);

            let instruction = (
                block_validity_window,
                TimestampValidityWindow::new_unbounded(),
            );
            let (output, proof) = circuit::execute_and_prove(
                vec![pre],
                Program::serialize_instruction(instruction).unwrap(),
                vec![2],
                vec![(account_keys.npk(), shared_secret)],
                vec![],
                vec![account_keys.identifier],
                vec![None],
                &validity_window_program.into(),
            )
            .unwrap();

            let message = Message::try_from_circuit_output(
                vec![],
                vec![],
                vec![(account_id, account_keys.vpk(), epk)],
                output,
            )
            .unwrap();
            //
            let witness_set = WitnessSet::for_message(&message, proof, &[]);
            PrivacyPreservingTransaction::new(message, witness_set)
        };
        let result = state.transition_from_privacy_preserving_transaction(&tx, block_id, 0);
        let is_inside_validity_window =
            match (block_validity_window.start(), block_validity_window.end()) {
                (Some(s), Some(e)) => s <= block_id && block_id < e,
                (Some(s), None) => s <= block_id,
                (None, Some(e)) => block_id < e,
                (None, None) => true,
            };
        if is_inside_validity_window {
            assert!(result.is_ok());
        } else {
            assert!(matches!(result, Err(NssaError::OutOfValidityWindow)));
        }
    }

    #[test_case::test_case((Some(1), Some(3)), 3; "at upper bound")]
    #[test_case::test_case((Some(1), Some(3)), 2; "inside range")]
    #[test_case::test_case((Some(1), Some(3)), 0; "below range")]
    #[test_case::test_case((Some(1), Some(3)), 1; "at lower bound")]
    #[test_case::test_case((Some(1), Some(3)), 4; "above range")]
    #[test_case::test_case((Some(1), None), 1; "lower bound only - at bound")]
    #[test_case::test_case((Some(1), None), 10; "lower bound only - above")]
    #[test_case::test_case((Some(1), None), 0; "lower bound only - below")]
    #[test_case::test_case((None, Some(3)), 3; "upper bound only - at bound")]
    #[test_case::test_case((None, Some(3)), 0; "upper bound only - below")]
    #[test_case::test_case((None, Some(3)), 4; "upper bound only - above")]
    #[test_case::test_case((None, None), 0; "no bounds - always valid")]
    #[test_case::test_case((None, None), 100; "no bounds - always valid 2")]
    fn timestamp_validity_window_works_in_privacy_preserving_transactions(
        validity_window: (Option<Timestamp>, Option<Timestamp>),
        timestamp: Timestamp,
    ) {
        let timestamp_validity_window: TimestampValidityWindow =
            validity_window.try_into().unwrap();
        let validity_window_program = Program::validity_window();
        let account_keys = test_private_account_keys_1();
        let account_id =
            AccountId::private_account_id(&account_keys.npk(), account_keys.identifier);
        let pre = AccountWithMetadata::new(Account::default(), false, account_id);
        let mut state = V03State::new_with_genesis_accounts(&[], &[], 0).with_test_programs();
        let tx = {
            let esk = [3; 32];
            let shared_secret = SharedSecretKey::new(&esk, &account_keys.vpk());
            let epk = EphemeralPublicKey::from_scalar(esk);

            let instruction = (
                BlockValidityWindow::new_unbounded(),
                timestamp_validity_window,
            );
            let (output, proof) = circuit::execute_and_prove(
                vec![pre],
                Program::serialize_instruction(instruction).unwrap(),
                vec![2],
                vec![(account_keys.npk(), shared_secret)],
                vec![],
                vec![account_keys.identifier],
                vec![None],
                &validity_window_program.into(),
            )
            .unwrap();

            let message = Message::try_from_circuit_output(
                vec![],
                vec![],
                vec![(account_id, account_keys.vpk(), epk)],
                output,
            )
            .unwrap();

            let witness_set = WitnessSet::for_message(&message, proof, &[]);
            PrivacyPreservingTransaction::new(message, witness_set)
        };
        let result = state.transition_from_privacy_preserving_transaction(&tx, 1, timestamp);
        let is_inside_validity_window = match (
            timestamp_validity_window.start(),
            timestamp_validity_window.end(),
        ) {
            (Some(s), Some(e)) => s <= timestamp && timestamp < e,
            (Some(s), None) => s <= timestamp,
            (None, Some(e)) => timestamp < e,
            (None, None) => true,
        };
        if is_inside_validity_window {
            assert!(result.is_ok());
        } else {
            assert!(matches!(result, Err(NssaError::OutOfValidityWindow)));
        }
    }

    fn time_locked_transfer_transaction(
        from: AccountId,
        from_key: &PrivateKey,
        from_nonce: u128,
        to: AccountId,
        clock_account_id: AccountId,
        amount: u128,
        deadline: u64,
    ) -> PublicTransaction {
        let program_id = Program::time_locked_transfer().id();
        let message = public_transaction::Message::try_new(
            program_id,
            vec![from, to, clock_account_id],
            vec![Nonce(from_nonce)],
            (amount, deadline),
        )
        .unwrap();
        let witness_set = public_transaction::WitnessSet::for_message(&message, &[from_key]);
        PublicTransaction::new(message, witness_set)
    }

    #[test]
    fn time_locked_transfer_succeeds_when_deadline_has_passed() {
        let recipient_id = AccountId::new([42; 32]);
        let genesis_timestamp = 500_u64;
        let mut state =
            V03State::new_with_genesis_accounts(&[(recipient_id, 0)], &[], genesis_timestamp)
                .with_test_programs();
        let key1 = PrivateKey::try_new([1; 32]).unwrap();
        let sender_id = AccountId::from(&PublicKey::new_from_private_key(&key1));
        state.force_insert_account(
            sender_id,
            Account {
                program_owner: Program::time_locked_transfer().id(),
                balance: 100,
                ..Account::default()
            },
        );

        let amount = 100_u128;
        // Deadline in the past: transfer should succeed.
        let deadline = 0_u64;

        let tx = time_locked_transfer_transaction(
            sender_id,
            &key1,
            0,
            recipient_id,
            CLOCK_01_PROGRAM_ACCOUNT_ID,
            amount,
            deadline,
        );

        let block_id = 1;
        let timestamp = genesis_timestamp + 100;
        state
            .transition_from_public_transaction(&tx, block_id, timestamp)
            .unwrap();

        // Balances changed.
        assert_eq!(state.get_account_by_id(sender_id).balance, 0);
        assert_eq!(state.get_account_by_id(recipient_id).balance, 100);
    }

    #[test]
    fn time_locked_transfer_fails_when_deadline_is_in_the_future() {
        let recipient_id = AccountId::new([42; 32]);
        let genesis_timestamp = 500_u64;
        let mut state =
            V03State::new_with_genesis_accounts(&[(recipient_id, 0)], &[], genesis_timestamp)
                .with_test_programs();
        let key1 = PrivateKey::try_new([1; 32]).unwrap();
        let sender_id = AccountId::from(&PublicKey::new_from_private_key(&key1));
        state.force_insert_account(
            sender_id,
            Account {
                program_owner: Program::time_locked_transfer().id(),
                balance: 100,
                ..Account::default()
            },
        );

        let amount = 100_u128;
        // Far-future deadline: program should panic.
        let deadline = u64::MAX;

        let tx = time_locked_transfer_transaction(
            sender_id,
            &key1,
            0,
            recipient_id,
            CLOCK_01_PROGRAM_ACCOUNT_ID,
            amount,
            deadline,
        );

        let block_id = 1;
        let timestamp = genesis_timestamp + 100;
        let result = state.transition_from_public_transaction(&tx, block_id, timestamp);

        assert!(
            result.is_err(),
            "Transfer should fail when deadline is in the future"
        );
        // Balances unchanged.
        assert_eq!(state.get_account_by_id(sender_id).balance, 100);
        assert_eq!(state.get_account_by_id(recipient_id).balance, 0);
    }

    fn pinata_cooldown_data(prize: u128, cooldown_ms: u64, last_claim_timestamp: u64) -> Vec<u8> {
        let mut buf = Vec::with_capacity(32);
        buf.extend_from_slice(&prize.to_le_bytes());
        buf.extend_from_slice(&cooldown_ms.to_le_bytes());
        buf.extend_from_slice(&last_claim_timestamp.to_le_bytes());
        buf
    }

    fn pinata_cooldown_transaction(
        pinata_id: AccountId,
        winner_id: AccountId,
        clock_account_id: AccountId,
    ) -> PublicTransaction {
        let program_id = Program::pinata_cooldown().id();
        let message = public_transaction::Message::try_new(
            program_id,
            vec![pinata_id, winner_id, clock_account_id],
            vec![],
            (),
        )
        .unwrap();
        let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
        PublicTransaction::new(message, witness_set)
    }

    #[test]
    fn pinata_cooldown_claim_succeeds_after_cooldown() {
        let winner_id = AccountId::new([11; 32]);
        let pinata_id = AccountId::new([99; 32]);

        let genesis_timestamp = 1000_u64;
        let mut state =
            V03State::new_with_genesis_accounts(&[(winner_id, 0)], &[], genesis_timestamp)
                .with_test_programs();

        let prize = 50_u128;
        let cooldown_ms = 500_u64;
        // Last claim was at genesis, so any timestamp >= genesis + cooldown should work.
        let last_claim_timestamp = genesis_timestamp;

        state.force_insert_account(
            pinata_id,
            Account {
                program_owner: Program::pinata_cooldown().id(),
                balance: 1000,
                data: pinata_cooldown_data(prize, cooldown_ms, last_claim_timestamp)
                    .try_into()
                    .unwrap(),
                ..Account::default()
            },
        );

        let tx = pinata_cooldown_transaction(pinata_id, winner_id, CLOCK_01_PROGRAM_ACCOUNT_ID);

        let block_id = 1;
        let block_timestamp = genesis_timestamp + cooldown_ms;
        // Advance clock so the cooldown check reads an updated timestamp.
        let clock_tx = clock_transaction(block_timestamp);
        state
            .transition_from_public_transaction(&clock_tx, block_id, block_timestamp)
            .unwrap();

        state
            .transition_from_public_transaction(&tx, block_id, block_timestamp)
            .unwrap();

        assert_eq!(state.get_account_by_id(pinata_id).balance, 1000 - prize);
        assert_eq!(state.get_account_by_id(winner_id).balance, prize);
    }

    #[test]
    fn pinata_cooldown_claim_fails_during_cooldown() {
        let winner_id = AccountId::new([11; 32]);
        let pinata_id = AccountId::new([99; 32]);

        let genesis_timestamp = 1000_u64;
        let mut state =
            V03State::new_with_genesis_accounts(&[(winner_id, 0)], &[], genesis_timestamp)
                .with_test_programs();

        let prize = 50_u128;
        let cooldown_ms = 500_u64;
        let last_claim_timestamp = genesis_timestamp;

        state.force_insert_account(
            pinata_id,
            Account {
                balance: 1000,
                data: pinata_cooldown_data(prize, cooldown_ms, last_claim_timestamp)
                    .try_into()
                    .unwrap(),
                ..Account::default()
            },
        );

        let tx = pinata_cooldown_transaction(pinata_id, winner_id, CLOCK_01_PROGRAM_ACCOUNT_ID);

        let block_id = 1;
        // Timestamp is only 100ms after last claim, well within the 500ms cooldown.
        let block_timestamp = genesis_timestamp + 100;
        let clock_tx = clock_transaction(block_timestamp);
        state
            .transition_from_public_transaction(&clock_tx, block_id, block_timestamp)
            .unwrap();

        let result = state.transition_from_public_transaction(&tx, block_id, block_timestamp);

        assert!(result.is_err(), "Claim should fail during cooldown period");
        assert_eq!(state.get_account_by_id(pinata_id).balance, 1000);
        assert_eq!(state.get_account_by_id(winner_id).balance, 0);
    }

    #[test]
    fn state_serialization_roundtrip() {
        let account_id_1 = AccountId::new([1; 32]);
        let account_id_2 = AccountId::new([2; 32]);
        let initial_data = [(account_id_1, 100_u128), (account_id_2, 151_u128)];
        let state = V03State::new_with_genesis_accounts(&initial_data, &[], 0).with_test_programs();
        let bytes = borsh::to_vec(&state).unwrap();
        let state_from_bytes: V03State = borsh::from_slice(&bytes).unwrap();
        assert_eq!(state, state_from_bytes);
    }

    #[test]
    fn flash_swap_successful() {
        let initiator = Program::flash_swap_initiator();
        let callback = Program::flash_swap_callback();
        let token = Program::authenticated_transfer_program();

        let vault_id = AccountId::from((&initiator.id(), &PdaSeed::new([0_u8; 32])));
        let receiver_id = AccountId::from((&callback.id(), &PdaSeed::new([1_u8; 32])));

        let initial_balance: u128 = 1000;
        let amount_out: u128 = 100;

        let vault_account = Account {
            program_owner: token.id(),
            balance: initial_balance,
            ..Account::default()
        };
        let receiver_account = Account {
            program_owner: token.id(),
            balance: 0,
            ..Account::default()
        };

        let mut state = V03State::new_with_genesis_accounts(&[], &[], 0).with_test_programs();
        state.force_insert_account(vault_id, vault_account);
        state.force_insert_account(receiver_id, receiver_account);

        // Callback instruction: return funds
        let cb_instruction = CallbackInstruction {
            return_funds: true,
            token_program_id: token.id(),
            amount: amount_out,
        };
        let cb_data = Program::serialize_instruction(cb_instruction).unwrap();

        let instruction = FlashSwapInstruction::Initiate {
            token_program_id: token.id(),
            callback_program_id: callback.id(),
            amount_out,
            callback_instruction_data: cb_data,
        };

        let tx = build_flash_swap_tx(&initiator, vault_id, receiver_id, instruction);
        let result = state.transition_from_public_transaction(&tx, 1, 0);
        assert!(result.is_ok(), "flash swap should succeed: {result:?}");

        // Vault balance restored, receiver back to 0
        assert_eq!(state.get_account_by_id(vault_id).balance, initial_balance);
        assert_eq!(state.get_account_by_id(receiver_id).balance, 0);
    }

    #[test]
    fn flash_swap_callback_keeps_funds_rollback() {
        let initiator = Program::flash_swap_initiator();
        let callback = Program::flash_swap_callback();
        let token = Program::authenticated_transfer_program();

        let vault_id = AccountId::from((&initiator.id(), &PdaSeed::new([0_u8; 32])));
        let receiver_id = AccountId::from((&callback.id(), &PdaSeed::new([1_u8; 32])));

        let initial_balance: u128 = 1000;
        let amount_out: u128 = 100;

        let vault_account = Account {
            program_owner: token.id(),
            balance: initial_balance,
            ..Account::default()
        };
        let receiver_account = Account {
            program_owner: token.id(),
            balance: 0,
            ..Account::default()
        };

        let mut state = V03State::new_with_genesis_accounts(&[], &[], 0).with_test_programs();
        state.force_insert_account(vault_id, vault_account);
        state.force_insert_account(receiver_id, receiver_account);

        // Callback instruction: do NOT return funds
        let cb_instruction = CallbackInstruction {
            return_funds: false,
            token_program_id: token.id(),
            amount: amount_out,
        };
        let cb_data = Program::serialize_instruction(cb_instruction).unwrap();

        let instruction = FlashSwapInstruction::Initiate {
            token_program_id: token.id(),
            callback_program_id: callback.id(),
            amount_out,
            callback_instruction_data: cb_data,
        };

        let tx = build_flash_swap_tx(&initiator, vault_id, receiver_id, instruction);
        let result = state.transition_from_public_transaction(&tx, 1, 0);

        // Invariant check fails → entire tx rolls back
        assert!(
            result.is_err(),
            "flash swap should fail when callback keeps funds"
        );

        // State unchanged (rollback)
        assert_eq!(state.get_account_by_id(vault_id).balance, initial_balance);
        assert_eq!(state.get_account_by_id(receiver_id).balance, 0);
    }

    #[test]
    fn flash_swap_self_call_targets_correct_program() {
        // Zero-amount flash swap: the invariant self-call still runs and succeeds
        // because vault balance doesn't decrease.
        let initiator = Program::flash_swap_initiator();
        let callback = Program::flash_swap_callback();
        let token = Program::authenticated_transfer_program();

        let vault_id = AccountId::from((&initiator.id(), &PdaSeed::new([0_u8; 32])));
        let receiver_id = AccountId::from((&callback.id(), &PdaSeed::new([1_u8; 32])));

        let initial_balance: u128 = 1000;

        let vault_account = Account {
            program_owner: token.id(),
            balance: initial_balance,
            ..Account::default()
        };
        let receiver_account = Account {
            program_owner: token.id(),
            balance: 0,
            ..Account::default()
        };

        let mut state = V03State::new_with_genesis_accounts(&[], &[], 0).with_test_programs();
        state.force_insert_account(vault_id, vault_account);
        state.force_insert_account(receiver_id, receiver_account);

        let cb_instruction = CallbackInstruction {
            return_funds: true,
            token_program_id: token.id(),
            amount: 0,
        };
        let cb_data = Program::serialize_instruction(cb_instruction).unwrap();

        let instruction = FlashSwapInstruction::Initiate {
            token_program_id: token.id(),
            callback_program_id: callback.id(),
            amount_out: 0,
            callback_instruction_data: cb_data,
        };

        let tx = build_flash_swap_tx(&initiator, vault_id, receiver_id, instruction);
        let result = state.transition_from_public_transaction(&tx, 1, 0);
        assert!(
            result.is_ok(),
            "zero-amount flash swap should succeed: {result:?}"
        );
    }

    #[test]
    fn flash_swap_standalone_invariant_check_rejected() {
        // Calling InvariantCheck directly (not as a chained self-call) should fail
        // because caller_program_id will be None.
        let initiator = Program::flash_swap_initiator();
        let token = Program::authenticated_transfer_program();

        let vault_id = AccountId::from((&initiator.id(), &PdaSeed::new([0_u8; 32])));

        let vault_account = Account {
            program_owner: token.id(),
            balance: 1000,
            ..Account::default()
        };

        let mut state = V03State::new_with_genesis_accounts(&[], &[], 0).with_test_programs();
        state.force_insert_account(vault_id, vault_account);

        let instruction = FlashSwapInstruction::InvariantCheck {
            min_vault_balance: 1000,
        };

        let message = public_transaction::Message::try_new(
            initiator.id(),
            vec![vault_id],
            vec![],
            instruction,
        )
        .unwrap();
        let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
        let tx = PublicTransaction::new(message, witness_set);

        let result = state.transition_from_public_transaction(&tx, 1, 0);
        assert!(
            result.is_err(),
            "standalone InvariantCheck should be rejected (caller_program_id is None)"
        );
    }

    #[test]
    fn malicious_self_program_id_rejected_in_public_execution() {
        let program = Program::malicious_self_program_id();
        let acc_id = AccountId::new([99; 32]);
        let account = Account::default();

        let mut state = V03State::new_with_genesis_accounts(&[], &[], 0).with_test_programs();
        state.force_insert_account(acc_id, account);

        let message =
            public_transaction::Message::try_new(program.id(), vec![acc_id], vec![], ()).unwrap();
        let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
        let tx = PublicTransaction::new(message, witness_set);

        let result = state.transition_from_public_transaction(&tx, 1, 0);
        assert!(
            result.is_err(),
            "program with wrong self_program_id in output should be rejected"
        );
    }

    #[test]
    fn malicious_caller_program_id_rejected_in_public_execution() {
        let program = Program::malicious_caller_program_id();
        let acc_id = AccountId::new([99; 32]);
        let account = Account::default();

        let mut state = V03State::new_with_genesis_accounts(&[], &[], 0).with_test_programs();
        state.force_insert_account(acc_id, account);

        let message =
            public_transaction::Message::try_new(program.id(), vec![acc_id], vec![], ()).unwrap();
        let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
        let tx = PublicTransaction::new(message, witness_set);

        let result = state.transition_from_public_transaction(&tx, 1, 0);
        assert!(
            result.is_err(),
            "program with spoofed caller_program_id in output should be rejected"
        );
    }
}