miden-testing 0.17.0

use alloc::string::String;
use alloc::sync::Arc;
use alloc::vec::Vec;
use std::collections::BTreeMap;

use anyhow::Context;
use assert_matches::assert_matches;
use miden_crypto::rand::{FeltRng, RandomCoin};
use miden_processor::{ExecutionError, Word};
use miden_protocol::Hasher;
use miden_protocol::account::auth::AuthScheme;
use miden_protocol::account::component::AccountComponentMetadata;
use miden_protocol::account::{
    Account,
    AccountBuilder,
    AccountCode,
    AccountCodeUpgrade,
    AccountComponent,
    AccountId,
    AccountProcedureRoot,
    AccountStorage,
    AccountType,
    StorageMap,
    StorageMapKey,
    StorageMapPatch,
    StorageSlot,
    StorageSlotContent,
    StorageSlotId,
    StorageSlotName,
    StorageSlotPatch,
    StorageSlotType,
    StorageValuePatch,
};
use miden_protocol::assembly::diagnostics::reporting::PrintDiagnostic;
use miden_protocol::assembly::{DefaultSourceManager, Linkage, ModuleKind, ModuleParser, Path};
use miden_protocol::asset::{Asset, AssetId, FungibleAsset};
use miden_protocol::errors::MasmError;
use miden_protocol::errors::tx_kernel::{
    ERR_ACCOUNT_AUTH_PROCEDURE_MUST_NOT_BE_DUPLICATED,
    ERR_ACCOUNT_CODE_COMMITMENT_MISMATCH,
    ERR_ACCOUNT_ID_SUFFIX_LEAST_SIGNIFICANT_BYTE_MUST_BE_ZERO,
    ERR_ACCOUNT_ID_SUFFIX_MOST_SIGNIFICANT_BIT_MUST_BE_ZERO,
    ERR_ACCOUNT_ID_UNKNOWN_VERSION,
    ERR_ACCOUNT_NONCE_AT_MAX,
    ERR_ACCOUNT_NONCE_CAN_ONLY_BE_INCREMENTED_ONCE,
    ERR_ACCOUNT_NOT_ENOUGH_PROCEDURES,
    ERR_ACCOUNT_PATCH_NONCE_MUST_BE_INCREMENTED_IF_STATE_CHANGED,
    ERR_ACCOUNT_PROCEDURES_MUST_BE_SORTED_AND_UNIQUE,
    ERR_ACCOUNT_STORAGE_SLOT_TYPE_IS_INVALID,
    ERR_ACCOUNT_UNKNOWN_STORAGE_SLOT_NAME,
    ERR_ACCOUNT_UPGRADE_CODE_ALREADY_SET,
    ERR_ACCOUNT_UPGRADE_NOT_ALLOWED_FOR_NEW_ACCOUNT,
    ERR_ACCOUNT_UPGRADE_STORAGE_UPGRADES_UNSUPPORTED,
};
use miden_protocol::field::PrimeField64;
use miden_protocol::note::NoteType;
use miden_protocol::testing::account_id::{
    ACCOUNT_ID_PRIVATE_NON_FUNGIBLE_FAUCET,
    ACCOUNT_ID_PUBLIC_FUNGIBLE_FAUCET,
    ACCOUNT_ID_PUBLIC_FUNGIBLE_FAUCET_1,
    ACCOUNT_ID_REGULAR_PRIVATE_ACCOUNT_UPDATABLE_CODE,
    ACCOUNT_ID_REGULAR_PUBLIC_ACCOUNT_IMMUTABLE_CODE,
    ACCOUNT_ID_REGULAR_PUBLIC_ACCOUNT_UPDATABLE_CODE,
    ACCOUNT_ID_SENDER,
};
use miden_protocol::testing::add_component::AddComponent;
use miden_protocol::testing::storage::{MOCK_MAP_SLOT, MOCK_VALUE_SLOT0, MOCK_VALUE_SLOT1};
use miden_protocol::transaction::memory::{
    CODE_UPGRADE_COMMITMENT_PTR,
    STORAGE_UPGRADE_COMMITMENT_PTR,
};
use miden_protocol::transaction::{RawOutputNote, TransactionKernel, TransactionScript};
use miden_protocol::utils::sync::LazyLock;
use miden_protocol::vm::Package;
use miden_standards::code_builder::CodeBuilder;
use miden_standards::testing::account_component::{IncrNonceAuthComponent, MockAccountComponent};
use miden_standards::testing::mock_account::MockAccountExt;
use miden_tx::{LocalTransactionProver, TransactionKernelError};
use rand::random;

use super::{Felt, StackInputs, ZERO};
use crate::executor::CodeExecutor;
use crate::kernel_tests::tx::ExecutionOutputExt;
use crate::utils::create_public_p2any_note;
use crate::{
    Auth,
    ExecError,
    MockChain,
    TestTransactionBuilder,
    assert_execution_error,
    assert_transaction_executor_error,
};

// ACCOUNT COMMITMENT TESTS
// ================================================================================================

#[tokio::test]
pub async fn compute_commitment() -> anyhow::Result<()> {
    let account = Account::mock(ACCOUNT_ID_REGULAR_PRIVATE_ACCOUNT_UPDATABLE_CODE, Auth::IncrNonce);

    // Precompute a commitment to a changed account so we can assert it during tx script execution.
    let mut account_clone = account.clone();
    let key = StorageMapKey::from_array([1, 2, 3, 4]);
    let value = Word::from([2, 3, 4, 5u32]);
    let mock_map_slot = &*MOCK_MAP_SLOT;
    account_clone.storage_mut().set_map_item(mock_map_slot, key, value).unwrap();
    let expected_commitment = account_clone.to_commitment();

    let tx_script = format!(
        r#"
        use miden::core::word

        use mock::account as mock_account

        const MOCK_MAP_SLOT = word("{mock_map_slot}")

        @transaction_script
        pub proc main
            call.mock_account::get_initial_commitment
            # => [INITIAL_COMMITMENT, pad(12)]
            swapdw dropw dropw swapw dropw
            # => [INITIAL_COMMITMENT]

            call.mock_account::compute_commitment
            # => [CURRENT_COMMITMENT, INITIAL_COMMITMENT, pad(8)]
            swapdw dropw dropw
            # => [CURRENT_COMMITMENT, INITIAL_COMMITMENT]

            assert_eqw.err="initial and current commitment should be equal when no changes have been made"
            # => []

            call.mock_account::compute_storage_commitment
            # => [STORAGE_COMMITMENT0, pad(12)]
            swapdw dropw dropw swapw dropw
            # => [STORAGE_COMMITMENT0]

            # update a value in the storage map
            padw push.0.0.0
            push.{value}
            push.{key}
            push.MOCK_MAP_SLOT[0..2]
            # => [slot_id_suffix, slot_id_prefix, KEY, VALUE, pad(7)]
            call.mock_account::set_map_item
            dropw dropw dropw dropw
            # => [STORAGE_COMMITMENT0]

            # compute the commitment which will recompute the storage commitment
            call.mock_account::compute_commitment
            # => [CURRENT_COMMITMENT, STORAGE_COMMITMENT0, pad(8)]
            swapdw dropw dropw
            # => [CURRENT_COMMITMENT, STORAGE_COMMITMENT0]

            push.{expected_commitment}
            assert_eqw.err="current commitment should match expected one"
            # => [STORAGE_COMMITMENT0]

            padw padw padw padw
            call.mock_account::compute_storage_commitment
            # => [STORAGE_COMMITMENT1, pad(12), STORAGE_COMMITMENT0]
            swapdw dropw dropw swapw dropw
            # => [STORAGE_COMMITMENT1, STORAGE_COMMITMENT0]

            # assert that the commitment has changed
            exec.word::eq
            assertz.err="storage commitment should have been updated by compute_commitment"
            # => []
        end
    "#,
        key = key,
        value = value,
        expected_commitment = expected_commitment,
    );

    let mock_tx_builder = TestTransactionBuilder::new(account);
    let tx_script = CodeBuilder::with_mock_packages().compile_tx_script(tx_script)?;
    let mock_tx = mock_tx_builder.tx_script(tx_script).build()?;

    mock_tx
        .execute()
        .await
        .map_err(|err| anyhow::anyhow!("failed to execute transaction: {err}"))?;

    Ok(())
}

/// State changes are visible before authentication increments the nonce, and reverting storage
/// restores the initial commitment.
#[tokio::test]
async fn has_state_changed() -> anyhow::Result<()> {
    let account = Account::mock(ACCOUNT_ID_REGULAR_PRIVATE_ACCOUNT_UPDATABLE_CODE, Auth::IncrNonce);
    let mock_map_slot = &*MOCK_MAP_SLOT;
    let key = StorageMapKey::from_array([1, 2, 3, 4]);
    let original_value = account.storage().get_map_item(mock_map_slot, key)?;
    let changed_value = Word::from([2, 3, 4, 5u32]);
    assert_ne!(original_value, changed_value);

    let tx_script = format!(
        r#"
        use mock::account as mock_account

        const MOCK_MAP_SLOT = word("{mock_map_slot}")

        @transaction_script
        pub proc main
            call.mock_account::has_state_changed
            assertz.err="the native account should initially be unchanged"
            dropw dropw dropw drop drop drop

            push.{changed_value} push.{key} push.MOCK_MAP_SLOT[0..2]
            call.mock_account::set_map_item
            dropw dropw dropw dropw

            call.mock_account::has_state_changed
            assert.err="storage changes should be detected before incrementing the nonce"
            dropw dropw dropw drop drop drop

            push.{original_value} push.{key} push.MOCK_MAP_SLOT[0..2]
            call.mock_account::set_map_item
            dropw dropw dropw dropw

            call.mock_account::has_state_changed
            assertz.err="restoring storage should restore the initial commitment"
            dropw dropw dropw drop drop drop
        end
        "#,
    );
    let tx_script = CodeBuilder::with_mock_packages().compile_tx_script(tx_script)?;

    TestTransactionBuilder::new(account)
        .tx_script(tx_script)
        .build()?
        .execute()
        .await?;

    Ok(())
}

/// Native-account commitment APIs still require an account procedure as their caller.
#[rstest::rstest]
#[case::compute_commitment("compute_commitment")]
#[case::has_state_changed("has_state_changed")]
#[tokio::test]
async fn native_account_commitment_from_tx_script_is_rejected(
    #[case] procedure: &str,
) -> anyhow::Result<()> {
    let tx_script = CodeBuilder::default().compile_tx_script(format!(
        r#"
        use miden::core::sys
        use miden::protocol::native_account

        @transaction_script
        pub proc main
            # keep this script's root distinct from the mock account's wrapper
            push.1 drop
            exec.native_account::{procedure}
            exec.sys::truncate_stack
        end
        "#,
    ))?;

    let result = TestTransactionBuilder::with_existing_mock_account()
        .tx_script(tx_script)
        .build()?
        .execute()
        .await;

    assert_transaction_executor_error!(
        result,
        matches ExecutionError::EventError { error: ref event_err, .. }
            if matches!(
                event_err.downcast_ref::<TransactionKernelError>(),
                Some(TransactionKernelError::UnknownAccountProcedure(_))
            )
    );

    Ok(())
}

// ACCOUNT ID TESTS
// ================================================================================================

/// Splits a raw account ID into its suffix and prefix felts.
fn account_id_felts(account_id: &u128) -> anyhow::Result<(Felt, Felt)> {
    let prefix = Felt::try_from((account_id / (1u128 << 64)) as u64)?;
    let suffix = Felt::try_from((account_id % (1u128 << 64)) as u64)?;
    Ok((prefix, suffix))
}

#[tokio::test]
async fn test_account_validate_id() -> anyhow::Result<()> {
    let test_cases = [
        (ACCOUNT_ID_REGULAR_PUBLIC_ACCOUNT_IMMUTABLE_CODE, None),
        (ACCOUNT_ID_REGULAR_PRIVATE_ACCOUNT_UPDATABLE_CODE, None),
        (ACCOUNT_ID_PUBLIC_FUNGIBLE_FAUCET, None),
        (ACCOUNT_ID_PRIVATE_NON_FUNGIBLE_FAUCET, None),
        (
            // The zero account ID should be invalid by construction.
            0,
            Some(ERR_ACCOUNT_ID_UNKNOWN_VERSION),
        ),
        (
            // Set version to another unsupported value (10).
            (ACCOUNT_ID_REGULAR_PRIVATE_ACCOUNT_UPDATABLE_CODE & !(0x0f << 64)) | (0x0a << 64),
            Some(ERR_ACCOUNT_ID_UNKNOWN_VERSION),
        ),
        (
            // Set most significant bit to `1`.
            ACCOUNT_ID_PUBLIC_FUNGIBLE_FAUCET | (0x80 << 56),
            Some(ERR_ACCOUNT_ID_SUFFIX_MOST_SIGNIFICANT_BIT_MUST_BE_ZERO),
        ),
        (
            // Set lower 8 bits to a non-zero value (1).
            ACCOUNT_ID_PRIVATE_NON_FUNGIBLE_FAUCET | 1,
            Some(ERR_ACCOUNT_ID_SUFFIX_LEAST_SIGNIFICANT_BYTE_MUST_BE_ZERO),
        ),
    ];

    for (account_id, expected_error) in test_cases.iter() {
        // Manually split the account ID into prefix and suffix since we can't use AccountId methods
        // on invalid ids.
        let (prefix, suffix) = account_id_felts(account_id)?;

        let code = "
            use miden::tx_kernel_core::account

            begin
                exec.account::validate_id
            end
            ";

        let result = CodeExecutor::with_default_host()
            .stack_inputs(StackInputs::new(&[suffix, prefix]).unwrap())
            .run(code)
            .await;

        match (result.map_err(ExecError::into_execution_error), expected_error) {
            (Ok(_), None) => (),
            (Ok(_), Some(err)) => {
                anyhow::bail!("expected error {err} but validation was successful")
            },
            (
                Err(ExecutionError::OperationError {
                    err:
                        miden_processor::operation::OperationError::FailedAssertion {
                            err_code,
                            err_msg,
                        },
                    ..
                }),
                Some(err),
            ) => {
                if err_code != err.code() {
                    anyhow::bail!(
                        "actual error \"{}\" (code: {err_code}) did not match expected error {err}",
                        err_msg.as_ref().map(AsRef::as_ref).unwrap_or("<no message>")
                    );
                }
            },
            (Err(err), None) => {
                return Err(anyhow::anyhow!(
                    "validation is supposed to succeed but error occurred: {}",
                    PrintDiagnostic::new(&err)
                ));
            },
            (Err(err), Some(_)) => {
                return Err(anyhow::anyhow!(
                    "unexpected different error than expected: {}",
                    PrintDiagnostic::new(&err)
                ));
            },
        }
    }

    Ok(())
}

/// `account_id::validate_structure` enforces the version-independent structural requirements of an
/// account ID without constraining the version, so IDs carrying a future (currently unsupported)
/// version pass as long as their suffix is well-formed.
#[tokio::test]
async fn test_account_validate_structure_ignores_version() -> anyhow::Result<()> {
    let test_cases = [
        // A regular version-one ID is structurally valid.
        (ACCOUNT_ID_REGULAR_PUBLIC_ACCOUNT_IMMUTABLE_CODE, None),
        (
            // An unsupported version (10) is still accepted: the structure is validated, not the
            // version.
            (ACCOUNT_ID_REGULAR_PRIVATE_ACCOUNT_UPDATABLE_CODE & !(0x0f << 64)) | (0x0a << 64),
            None,
        ),
        (
            // Set most significant bit of the suffix to `1`.
            ACCOUNT_ID_PUBLIC_FUNGIBLE_FAUCET | (0x80 << 56),
            Some(ERR_ACCOUNT_ID_SUFFIX_MOST_SIGNIFICANT_BIT_MUST_BE_ZERO),
        ),
        (
            // Set lower 8 bits of the suffix to a non-zero value (1).
            ACCOUNT_ID_PRIVATE_NON_FUNGIBLE_FAUCET | 1,
            Some(ERR_ACCOUNT_ID_SUFFIX_LEAST_SIGNIFICANT_BYTE_MUST_BE_ZERO),
        ),
    ];

    for (account_id, expected_error) in test_cases.iter() {
        // Manually split the account ID into prefix and suffix since we can't use AccountId methods
        // on invalid ids.
        let (prefix, suffix) = account_id_felts(account_id)?;

        let code = "
            use miden::protocol::account_id

            begin
                exec.account_id::validate_structure
            end
            ";

        let result = CodeExecutor::with_default_host()
            .stack_inputs(StackInputs::new(&[suffix, prefix]).unwrap())
            .run(code)
            .await;

        match (result.map_err(ExecError::into_execution_error), expected_error) {
            (Ok(_), None) => (),
            (Ok(_), Some(err)) => {
                anyhow::bail!("expected error {err} but validation was successful")
            },
            (
                Err(ExecutionError::OperationError {
                    err:
                        miden_processor::operation::OperationError::FailedAssertion {
                            err_code,
                            err_msg,
                        },
                    ..
                }),
                Some(err),
            ) => {
                if err_code != err.code() {
                    anyhow::bail!(
                        "actual error \"{}\" (code: {err_code}) did not match expected error {err}",
                        err_msg.as_ref().map(AsRef::as_ref).unwrap_or("<no message>")
                    );
                }
            },
            (Err(err), None) => {
                return Err(anyhow::anyhow!(
                    "validation is supposed to succeed but error occurred: {}",
                    PrintDiagnostic::new(&err)
                ));
            },
            (Err(err), Some(_)) => {
                return Err(anyhow::anyhow!(
                    "unexpected different error than expected: {}",
                    PrintDiagnostic::new(&err)
                ));
            },
        }
    }

    Ok(())
}

/// Exercises the account ID comparison helpers -- `eq`, `eqz`, and `testz` -- in a single program.
#[tokio::test]
async fn test_account_id_comparison() -> anyhow::Result<()> {
    let (prefix_1, suffix_1) = account_id_felts(&ACCOUNT_ID_REGULAR_PUBLIC_ACCOUNT_IMMUTABLE_CODE)?;
    let (prefix_2, suffix_2) = account_id_felts(&ACCOUNT_ID_PUBLIC_FUNGIBLE_FAUCET_1)?;

    let code = format!(
        r#"
        use miden::protocol::account_id

        begin
            # eq: identical IDs are equal
            push.{prefix_1}.{suffix_1}.{prefix_1}.{suffix_1}
            exec.account_id::eq
            assert.err="eq: identical IDs should be equal"
            # => [pad(16)]

            # eq: different IDs are not equal
            push.{prefix_1}.{suffix_1}.{prefix_2}.{suffix_2}
            exec.account_id::eq
            assertz.err="eq: different IDs should not be equal"
            # => [pad(16)]

            # eqz: the zero address is zero
            push.0.0
            exec.account_id::eqz
            assert.err="eqz: the zero address should be zero"
            # => [pad(16)]

            # eqz: a valid ID is not zero
            push.{prefix_1}.{suffix_1}
            exec.account_id::eqz
            assertz.err="eqz: a valid ID should not be zero"
            # => [pad(16)]

            # testz: the zero address is zero
            push.0.0
            exec.account_id::testz
            assert.err="testz: the zero address should be zero"
            # => [pad(18)]

            # testz: a valid ID is not zero, leaving [suffix_1, prefix_1] on the stack
            push.{prefix_1}.{suffix_1}
            exec.account_id::testz
            assertz.err="testz: a valid ID should not be zero"
            # => [suffix_1, prefix_1, pad(18)]

            # truncate the stack
            swapw dropw
            # => [suffix_1, prefix_1, pad(14)]
        end
        "#
    );

    let exec_output = CodeExecutor::with_default_host()
        .run(&code)
        .await
        .map_err(ExecError::into_execution_error)?;

    // testz must preserve the account ID beneath the (already consumed) flag.
    assert_eq!(exec_output.get_stack_element(0), suffix_1, "testz must preserve the ID suffix");
    assert_eq!(exec_output.get_stack_element(1), prefix_1, "testz must preserve the ID prefix");

    Ok(())
}

// ACCOUNT CODE TESTS
// ================================================================================================

// TODO: update this test once the ability to change the account code will be implemented
#[tokio::test]
pub async fn test_compute_code_commitment() -> anyhow::Result<()> {
    let mock_tx = TestTransactionBuilder::with_existing_mock_account().build().unwrap();
    let account = mock_tx.account();

    let code = format!(
        r#"
        use miden::tx_kernel_core::prologue
        use mock::account as mock_account

        begin
            exec.prologue::prepare_transaction
            # get the code commitment
            call.mock_account::get_code_commitment
            push.{expected_code_commitment}
            assert_eqw.err="actual code commitment is not equal to the expected one"
        end
        "#,
        expected_code_commitment = account.code().commitment()
    );

    mock_tx.execute_code(&code).await?;

    Ok(())
}

// ACCOUNT STORAGE TESTS
// ================================================================================================

#[tokio::test]
async fn test_get_item() -> anyhow::Result<()> {
    for storage_item in [AccountStorage::mock_value_slot0(), AccountStorage::mock_value_slot1()] {
        let mock_tx = TestTransactionBuilder::with_existing_mock_account().build().unwrap();

        let code = format!(
            r#"
            use miden::tx_kernel_core::account
            use miden::tx_kernel_core::prologue

            const SLOT_NAME = word("{slot_name}")

            begin
                exec.prologue::prepare_transaction

                # push the account storage item index
                push.SLOT_NAME[0..2]
                # => [slot_id_suffix, slot_id_prefix]

                # assert the item value is correct
                exec.account::get_item
                push.{item_value}
                assert_eqw.err="expected item to have value {item_value}"
            end
            "#,
            slot_name = storage_item.name(),
            item_value = storage_item.content().value(),
        );

        mock_tx.execute_code(&code).await.unwrap();
    }

    Ok(())
}

#[tokio::test]
async fn test_get_map_item() -> anyhow::Result<()> {
    let slot = AccountStorage::mock_map_slot();
    let account = AccountBuilder::new(rand::random())
        .with_components(Auth::IncrNonce)
        .with_component(MockAccountComponent::with_slots(vec![slot.clone()]))
        .build_existing()
        .unwrap();

    let mock_tx = TestTransactionBuilder::new(account).build().unwrap();

    let StorageSlotContent::Map(map) = slot.content() else {
        panic!("expected map")
    };

    for (key, expected_value) in map.entries() {
        let code = format!(
            r#"
            use miden::tx_kernel_core::prologue
            use mock::account

            const SLOT_NAME = word("{slot_name}")

            begin
                exec.prologue::prepare_transaction

                # get the map item
                push.{key}
                push.SLOT_NAME[0..2]
                call.account::get_map_item
                # => [VALUE]

                push.{expected_value}
                assert_eqw.err="value did not match {expected_value}"

                exec.::miden::core::sys::truncate_stack
            end
            "#,
            slot_name = slot.name(),
        );

        mock_tx.execute_code(&code).await?;
    }

    Ok(())
}

#[tokio::test]
async fn test_get_native_storage_slot_type() -> anyhow::Result<()> {
    for slot_name in [
        AccountStorage::mock_value_slot0().name(),
        AccountStorage::mock_value_slot1().name(),
        AccountStorage::mock_map_slot().name(),
    ] {
        let mock_tx = TestTransactionBuilder::with_existing_mock_account().build().unwrap();
        let (slot_idx, slot) = mock_tx
            .account()
            .storage()
            .slots()
            .iter()
            .enumerate()
            .find(|(_, slot)| slot.name() == slot_name)
            .unwrap();

        let code = format!(
            "
            use miden::tx_kernel_core::account
            use miden::tx_kernel_core::prologue

            begin
                exec.prologue::prepare_transaction

                # push the account storage slot index
                push.{slot_idx}

                # get the type of the respective storage slot
                exec.account::get_native_storage_slot_type

                # truncate the stack
                swap drop
            end
            ",
        );

        let exec_output = &mock_tx.execute_code(&code).await.unwrap();

        assert_eq!(
            slot.slot_type(),
            StorageSlotType::try_from(
                u8::try_from(exec_output.get_stack_element(0).as_canonical_u64()).unwrap()
            )
            .unwrap()
        );
        assert_eq!(exec_output.get_stack_element(1), ZERO, "the rest of the stack is empty");
        assert_eq!(exec_output.get_stack_element(2), ZERO, "the rest of the stack is empty");
        assert_eq!(exec_output.get_stack_element(3), ZERO, "the rest of the stack is empty");
        assert_eq!(exec_output.get_stack_word(4), Word::empty(), "the rest of the stack is empty");
        assert_eq!(exec_output.get_stack_word(8), Word::empty(), "the rest of the stack is empty");
        assert_eq!(exec_output.get_stack_word(12), Word::empty(), "the rest of the stack is empty");
    }

    Ok(())
}

/// Tests that `validate_storage` rejects a storage slot whose type is outside the supported set
/// (value or map) instead of silently committing it as a map (audit finding L-11).
#[tokio::test]
async fn validate_storage_rejects_unsupported_slot_type() -> anyhow::Result<()> {
    let mock_tx = TestTransactionBuilder::with_existing_mock_account().build().unwrap();

    // Overwrite the type element (offset 1) of the first storage slot with an unsupported type,
    // then run the new-account storage validation which must reject it.
    let code = "
        use miden::tx_kernel_core::account
        use miden::tx_kernel_core::memory
        use miden::tx_kernel_core::prologue

        begin
            exec.prologue::prepare_transaction

            push.2
            exec.memory::get_native_account_active_storage_slots_ptr add.1
            mem_store

            exec.account::validate_storage
        end
        ";

    let exec_output = mock_tx.execute_code(code).await;
    assert_execution_error!(exec_output, ERR_ACCOUNT_STORAGE_SLOT_TYPE_IS_INVALID);

    Ok(())
}

/// Tests that accessing an unknown slot fails with the expected error message.
///
/// This tests both accounts with empty storage and non-empty storage.
#[tokio::test]
async fn test_account_get_item_fails_on_unknown_slot() -> anyhow::Result<()> {
    let mut builder = MockChain::builder();

    let account_empty_storage = builder.add_existing_mock_account(Auth::IncrNonce)?;
    assert_eq!(account_empty_storage.storage().num_slots(), 0);

    let account_non_empty_storage = builder.add_existing_mock_account(Auth::BasicAuth {
        auth_scheme: AuthScheme::Falcon512Poseidon2,
    })?;
    assert_eq!(account_non_empty_storage.storage().num_slots(), 2);

    let chain = builder.build()?;

    let code = r#"
            use mock::account

            const UNKNOWN_SLOT_NAME = word("unknown::slot::name")

            @transaction_script
            pub proc main
                push.UNKNOWN_SLOT_NAME[0..2]
                call.account::get_item
            end
            "#;
    let tx_script = CodeBuilder::with_mock_packages().compile_tx_script(code)?;

    let result = chain
        .build_transaction(account_empty_storage)
        .tx_script(tx_script.clone())
        .build()?
        .execute()
        .await;
    assert_transaction_executor_error!(result, ERR_ACCOUNT_UNKNOWN_STORAGE_SLOT_NAME);

    let result = chain
        .build_transaction(account_non_empty_storage)
        .tx_script(tx_script)
        .build()?
        .execute()
        .await;
    assert_transaction_executor_error!(result, ERR_ACCOUNT_UNKNOWN_STORAGE_SLOT_NAME);

    Ok(())
}

#[tokio::test]
async fn test_is_slot_id_lt() -> anyhow::Result<()> {
    // Note that the slot IDs derived from the names are essentially randomly sorted, so these cover
    // "less than" and "greater than" outcomes.
    let mut test_cases = (0..100)
        .map(|i| {
            let prev_slot = StorageSlotName::mock(i).id();
            let curr_slot = StorageSlotName::mock(i + 1).id();
            (prev_slot, curr_slot)
        })
        .collect::<Vec<_>>();

    // Extend with special case where prefix matches and suffix determines the outcome.
    let prefix = Felt::from(100u32);
    test_cases.extend([
        // prev_slot == curr_slot
        (
            StorageSlotId::new(Felt::from(50u32), prefix),
            StorageSlotId::new(Felt::from(50u32), prefix),
        ),
        // prev_slot < curr_slot
        (
            StorageSlotId::new(Felt::from(50u32), prefix),
            StorageSlotId::new(Felt::from(51u32), prefix),
        ),
        // prev_slot > curr_slot
        (
            StorageSlotId::new(Felt::from(51u32), prefix),
            StorageSlotId::new(Felt::from(50u32), prefix),
        ),
    ]);

    for (prev_slot, curr_slot) in test_cases {
        let code = format!(
            r#"
            use miden::tx_kernel_core::account

            begin
                push.{curr_prefix}.{curr_suffix}.{prev_prefix}.{prev_suffix}
                # => [prev_slot_id_suffix, prev_slot_id_prefix, curr_slot_id_suffix, curr_slot_id_prefix]

                exec.account::is_slot_id_lt
                # => [is_slot_id_lt]

                push.{is_lt}
                assert_eq.err="is_slot_id_lt was not {is_lt}"
                # => []
            end
            "#,
            prev_prefix = prev_slot.prefix(),
            prev_suffix = prev_slot.suffix(),
            curr_prefix = curr_slot.prefix(),
            curr_suffix = curr_slot.suffix(),
            is_lt = u8::from(prev_slot < curr_slot)
        );

        CodeExecutor::with_default_host().run(&code).await?;
    }

    Ok(())
}

/// Returns the requested number of distinct procedure roots sorted in ascending order.
fn sorted_procedure_roots(num_procedures: usize) -> Vec<AccountProcedureRoot> {
    let mut rng = RandomCoin::new([num_procedures as u32, 0, 0, 0].into());
    let mut roots = (0..num_procedures as u32)
        .map(|_| AccountProcedureRoot::from_raw(rng.draw_word()))
        .collect::<Vec<_>>();
    roots.sort_unstable();

    roots
}

/// Returns a program that writes the provided procedure roots into the native account's code
/// section and validates them.
/// TODO: Refactor to create the custom procedure section in local memory and call
/// validate_procedures with a pointer, so it becomes simpler.
fn validate_procedures_program(procedure_roots: &[AccountProcedureRoot]) -> String {
    let procedure_writes = procedure_roots
        .iter()
        .enumerate()
        .map(|(index, root)| {
            format!(
                "push.{root}
                exec.memory::get_account_procedures_section_ptr add.{offset} mem_storew_le dropw",
                offset = index * Word::NUM_ELEMENTS
            )
        })
        .collect::<Vec<_>>()
        .join("\n");

    format!(
        r#"
        use miden::tx_kernel_core::account
        use miden::tx_kernel_core::memory

        begin
            exec.memory::set_active_account_data_ptr_to_native_account

            push.{num_procedures} exec.memory::set_num_account_procedures

            {procedure_writes}

            exec.memory::get_active_account_code_section_ptr
            exec.account::validate_procedures
        end
        "#,
        num_procedures = procedure_roots.len()
    )
}

/// The auth procedure at index 0 is exempt from the ordering, so it is validated against roots that
/// sort before, in between and after the other procedures.
#[rstest::rstest]
#[case::minimum_number_of_procedures(2, 0)]
#[case::auth_procedure_sorts_first(8, 0)]
#[case::auth_procedure_sorts_in_between(8, 4)]
#[case::auth_procedure_sorts_last(8, 7)]
#[tokio::test]
async fn test_validate_procedures_accepts_sorted_and_unique_procedures(
    #[case] num_procedures: usize,
    #[case] auth_procedure_index: usize,
) -> anyhow::Result<()> {
    let mut procedure_roots = sorted_procedure_roots(num_procedures);
    let auth_procedure_root = procedure_roots.remove(auth_procedure_index);
    procedure_roots.insert(0, auth_procedure_root);

    CodeExecutor::with_default_host()
        .run(&validate_procedures_program(&procedure_roots))
        .await?;

    Ok(())
}

#[rstest::rstest]
#[case::first_pair_swapped(1, 2)]
#[case::last_pair_swapped(6, 7)]
#[case::first_and_last_swapped(1, 7)]
#[tokio::test]
async fn test_validate_procedures_rejects_unsorted_procedures(
    #[case] left_index: usize,
    #[case] right_index: usize,
) -> anyhow::Result<()> {
    let mut procedure_roots = sorted_procedure_roots(8);
    procedure_roots.swap(left_index, right_index);

    let result = CodeExecutor::with_default_host()
        .run(&validate_procedures_program(&procedure_roots))
        .await;

    assert_execution_error!(result, ERR_ACCOUNT_PROCEDURES_MUST_BE_SORTED_AND_UNIQUE);

    Ok(())
}

/// The duplicated procedure is always the one preceding it, which must be at index 1 or greater so
/// the duplicate does not involve the auth procedure at index 0.
#[rstest::rstest]
#[case::first_procedure(2)]
#[case::middle_procedure(4)]
#[case::last_procedure(7)]
#[tokio::test]
async fn test_validate_procedures_rejects_duplicated_procedure(
    #[case] duplicated_index: usize,
) -> anyhow::Result<()> {
    let mut procedure_roots = sorted_procedure_roots(8);
    procedure_roots[duplicated_index] = procedure_roots[duplicated_index - 1];

    let result = CodeExecutor::with_default_host()
        .run(&validate_procedures_program(&procedure_roots))
        .await;

    assert_execution_error!(result, ERR_ACCOUNT_PROCEDURES_MUST_BE_SORTED_AND_UNIQUE);

    Ok(())
}

/// A duplicated auth procedure is not caught by the ordering check, since the auth procedure is
/// exempt from it, so it must be rejected by the explicit comparison against the root at index 0.
#[rstest::rstest]
#[case::minimum_number_of_procedures(2, 1)]
#[case::first_procedure(8, 1)]
#[case::middle_procedure(8, 4)]
#[case::last_procedure(8, 7)]
#[tokio::test]
async fn test_validate_procedures_rejects_duplicated_auth_procedure(
    #[case] num_procedures: usize,
    #[case] duplicated_index: usize,
) -> anyhow::Result<()> {
    let mut procedure_roots = sorted_procedure_roots(num_procedures);
    procedure_roots[0] = procedure_roots[duplicated_index];

    let result = CodeExecutor::with_default_host()
        .run(&validate_procedures_program(&procedure_roots))
        .await;

    assert_execution_error!(result, ERR_ACCOUNT_AUTH_PROCEDURE_MUST_NOT_BE_DUPLICATED);

    Ok(())
}

#[tokio::test]
async fn test_set_item() -> anyhow::Result<()> {
    let mock_tx = TestTransactionBuilder::with_existing_mock_account().build().unwrap();

    let slot_name = &*MOCK_VALUE_SLOT0;
    let new_value = Word::from([91, 92, 93, 94u32]);
    let old_value = mock_tx.account().storage().get_item(slot_name)?;

    let code = format!(
        r#"
        use miden::tx_kernel_core::account
        use miden::tx_kernel_core::prologue

        const MOCK_VALUE_SLOT0 = word("{slot_name}")

        begin
            exec.prologue::prepare_transaction

            # set the storage item
            push.{new_value}
            push.MOCK_VALUE_SLOT0[0..2]
            # => [slot_id_suffix, slot_id_prefix, NEW_VALUE]

            exec.account::set_item

            # assert old value was correctly returned
            push.{old_value}
            assert_eqw.err="old value did not match"

            # assert new value has been correctly set
            push.MOCK_VALUE_SLOT0[0..2]
            # => [slot_id_suffix, slot_id_prefix]

            exec.account::get_item
            push.{new_value}
            assert_eqw.err="new value did not match"
        end
        "#,
    );

    mock_tx.execute_code(&code).await?;

    Ok(())
}

#[tokio::test]
async fn test_set_map_item() -> anyhow::Result<()> {
    let (new_key, new_value) = (
        StorageMapKey::from_array([109, 110, 111, 112u32]),
        Word::from([9, 10, 11, 12u32]),
    );

    let slot = AccountStorage::mock_map_slot();
    let account = AccountBuilder::new(rand::random())
        .with_components(Auth::IncrNonce)
        .with_component(MockAccountComponent::with_slots(vec![slot.clone()]))
        .build_existing()
        .unwrap();

    let mock_tx = TestTransactionBuilder::new(account).build().unwrap();

    let code = format!(
        r#"
        use miden::core::sys

        use miden::tx_kernel_core::prologue
        use mock::account as mock_account

        const SLOT_NAME=word("{slot_name}")

        begin
            exec.prologue::prepare_transaction

            # set the map item
            push.{new_value}
            push.{new_key}
            push.SLOT_NAME[0..2]
            call.mock_account::set_map_item

            # double check that the storage slot is indeed the new map
            push.SLOT_NAME[0..2]
            # => [slot_id_suffix, slot_id_prefix, OLD_VALUE]

            # pad the stack
            repeat.14 push.0 movdn.2 end
            # => [slot_id_suffix, slot_id_prefix, pad(14), OLD_VALUE]

            call.mock_account::get_item
            # => [MAP_ROOT, pad(12), OLD_VALUE]

            # truncate the stack
            repeat.3 swapw dropw end
            # => [MAP_ROOT, OLD_VALUE]

            exec.sys::truncate_stack
        end
        "#,
        slot_name = slot.name(),
        new_key = new_key,
        new_value = new_value,
    );

    let exec_output = &mock_tx.execute_code(&code).await?;

    let mut new_storage_map = AccountStorage::mock_map();
    new_storage_map.insert(new_key, new_value).unwrap();

    assert_eq!(
        new_storage_map.root(),
        exec_output.get_stack_word(0),
        "get_item should return the updated root",
    );

    let old_value_for_key = match slot.content() {
        StorageSlotContent::Map(original_map) => original_map.get(&new_key),
        _ => panic!("expected map"),
    };
    assert_eq!(
        old_value_for_key,
        exec_output.get_stack_word(4),
        "set_map_item must return the old value for the key (empty word for new key)",
    );

    Ok(())
}

/// Tests that we can successfully create regular and faucet accounts with empty storage.
#[tokio::test]
async fn create_account_with_empty_storage_slots() -> anyhow::Result<()> {
    let account = AccountBuilder::new([5; 32])
        .with_components(Auth::IncrNonce)
        .with_component(MockAccountComponent::with_empty_slots())
        .build()
        .context("failed to build account")?;

    TestTransactionBuilder::new(account).build()?.execute().await?;

    Ok(())
}

#[tokio::test]
async fn test_get_initial_storage_commitment() -> anyhow::Result<()> {
    let mock_tx = TestTransactionBuilder::with_existing_mock_account().build()?;

    let code = format!(
        r#"
        use mock::account as mock_account
        use miden::tx_kernel_core::prologue

        begin
            exec.prologue::prepare_transaction

            # get the initial storage commitment
            call.mock_account::get_initial_storage_commitment
            # => [INIT_STORAGE_COMMITMENT, pad(12)]
            swapdw dropw dropw swapw dropw
            # => [INIT_STORAGE_COMMITMENT]
            push.{expected_storage_commitment}
            assert_eqw.err="actual storage commitment is not equal to the expected one"
        end
        "#,
        expected_storage_commitment = mock_tx.account().storage().to_commitment(),
    );
    mock_tx.execute_code(&code).await?;

    Ok(())
}

// ACCOUNT UPGRADE TESTS
// ================================================================================================

/// Returns the existing mock account that the upgrade tests start from.
fn existing_mock_account() -> Account {
    Account::mock(ACCOUNT_ID_REGULAR_PUBLIC_ACCOUNT_UPDATABLE_CODE, [IncrNonceAuthComponent])
}

/// Returns the code of the mock account extended with the [`AddComponent`], which the upgrade tests
/// upgrade to.
fn upgraded_mock_account_code() -> anyhow::Result<AccountCode> {
    Ok(AccountCode::from_components(&[
        IncrNonceAuthComponent.into(),
        MockAccountComponent::with_slots(vec![]).into(),
        AddComponent.into(),
    ])?)
}

/// Returns a transaction script that initializes one account upgrade per provided pair of new code
/// and new storage commitments.
fn upgrade_tx_script(
    upgrades: impl IntoIterator<Item = (Word, Word)>,
) -> anyhow::Result<TransactionScript> {
    let upgrade_calls: String = upgrades
        .into_iter()
        .map(|(new_code_commitment, storage_upgrade_commitment)| {
            format!(
                "
                push.{storage_upgrade_commitment}
                push.{new_code_commitment}
                call.mock_account::upgrade
                dropw dropw
                "
            )
        })
        .collect();

    let tx_script = format!(
        r#"
        use mock::account as mock_account

        @transaction_script
        pub proc main
            {upgrade_calls}
        end
        "#
    );

    Ok(CodeBuilder::with_mock_packages().compile_tx_script(tx_script)?)
}

/// Returns a program that runs the prologue and then initializes an upgrade of the native account's
/// code to the code with `new_code_commitment`.
fn upgrade_program(new_code_commitment: Word) -> String {
    format!(
        r#"
        use mock::account as mock_account
        use miden::tx_kernel_core::prologue

        begin
            exec.prologue::prepare_transaction

            padw push.{new_code_commitment}
            # => [NEW_CODE_COMMITMENT, STORAGE_UPGRADE_COMMITMENT, pad(8)]

            call.mock_account::upgrade
            # => [pad(16)]
            dropw dropw dropw dropw
        end
        "#,
    )
}

/// Tests that `native_account::upgrade` stores the new code commitment in kernel memory, and the
/// empty word if the upgrade is a no-op.
#[rstest::rstest]
#[case::upgraded_code(upgraded_mock_account_code()?.commitment(), upgraded_mock_account_code()?.commitment())]
#[case::current_code(existing_mock_account().code().commitment(), Word::empty())]
#[case::empty_word(Word::empty(), Word::empty())]
#[tokio::test]
async fn test_account_upgrade_stores_code_commitment(
    #[case] new_code_commitment: Word,
    #[case] expected_code_upgrade_commitment: Word,
) -> anyhow::Result<()> {
    let mock_tx = TestTransactionBuilder::new(existing_mock_account())
        .account_code_upgrade(AccountCodeUpgrade::new(upgraded_mock_account_code()?))
        .build()?;

    let exec_output = &mock_tx.execute_code(&upgrade_program(new_code_commitment)).await?;

    assert_eq!(
        exec_output.get_kernel_mem_word(CODE_UPGRADE_COMMITMENT_PTR),
        expected_code_upgrade_commitment,
    );
    assert_eq!(exec_output.get_kernel_mem_word(STORAGE_UPGRADE_COMMITMENT_PTR), Word::empty());

    Ok(())
}

/// Tests that a transaction replaces the code of the native account in its final state and patch,
/// and that an upgrade to the current code leaves the code unchanged.
///
/// A pending upgrade must also count as a state change before the nonce is incremented.
#[rstest::rstest]
#[case::upgraded_code(true)]
#[case::current_code(false)]
#[tokio::test]
async fn test_account_upgrade_replaces_code(#[case] is_upgrade: bool) -> anyhow::Result<()> {
    let account = existing_mock_account();
    let upgraded_code = upgraded_mock_account_code()?;
    let expected_code = if is_upgrade {
        upgraded_code.clone()
    } else {
        account.code().clone()
    };

    let tx_script = format!(
        r#"
        use mock::account as mock_account

        @transaction_script
        pub proc main
            padw push.{new_code_commitment}
            # => [NEW_CODE_COMMITMENT, STORAGE_UPGRADE_COMMITMENT]

            call.mock_account::upgrade
            dropw dropw
            # => []

            call.mock_account::has_state_changed
            # => [has_state_changed]

            eq.{is_upgrade} assert.err="a pending code upgrade should count as a state change"
            dropw dropw dropw drop drop drop
            # => []
        end
        "#,
        new_code_commitment = expected_code.commitment(),
        is_upgrade = u8::from(is_upgrade),
    );
    let tx_script = CodeBuilder::with_mock_packages().compile_tx_script(tx_script)?;

    let executed_tx = TestTransactionBuilder::new(account.clone())
        .tx_script(tx_script)
        .account_code_upgrade(AccountCodeUpgrade::new(upgraded_code))
        .build()?
        .execute()
        .await?;

    assert_eq!(executed_tx.final_account().code_commitment(), expected_code.commitment());
    assert_eq!(
        executed_tx.account_patch().code().as_code(),
        is_upgrade.then_some(&expected_code)
    );

    let mut expected_account = account;
    expected_account.apply_patch(executed_tx.account_patch())?;
    assert_eq!(expected_account.to_commitment(), executed_tx.final_account().to_commitment());

    // Check account upgrades work under the prover's re-execution.
    LocalTransactionProver::default().prove(executed_tx)?;

    Ok(())
}

/// Tests that an auth component that signs the delta commitment signs the upgraded code, which
/// requires the delta commitments of the kernel and the host to agree on the code section.
#[tokio::test]
async fn test_account_upgrade_with_signature_auth() -> anyhow::Result<()> {
    let mut builder = MockChain::builder();
    let account = builder.add_existing_mock_account(Auth::BasicAuth {
        auth_scheme: AuthScheme::Falcon512Poseidon2,
    })?;
    let mock_chain = builder.build()?;

    let upgraded_code = upgraded_mock_account_code()?;
    let executed_tx = mock_chain
        .build_transaction(account.id())
        .tx_script(upgrade_tx_script([(upgraded_code.commitment(), Word::empty())])?)
        .account_code_upgrade(AccountCodeUpgrade::new(upgraded_code.clone()))
        .build()?
        .execute()
        .await?;

    assert_eq!(executed_tx.final_account().code_commitment(), upgraded_code.commitment());
    assert_eq!(executed_tx.account_patch().code().as_code(), Some(&upgraded_code));

    Ok(())
}

/// Tests that the kernel rejects upgrades that are invalid on their own.
#[rstest::rstest]
#[case::storage_upgrade(
    vec![(upgraded_mock_account_code()?.commitment(), Word::from([1, 2, 3, 4u32]))],
    ERR_ACCOUNT_UPGRADE_STORAGE_UPGRADES_UNSUPPORTED,
)]
#[case::second_upgrade(
    vec![(upgraded_mock_account_code()?.commitment(), Word::empty()); 2],
    ERR_ACCOUNT_UPGRADE_CODE_ALREADY_SET,
)]
#[tokio::test]
async fn test_account_upgrade_rejects_invalid_initialization(
    #[case] upgrades: Vec<(Word, Word)>,
    #[case] expected_error: MasmError,
) -> anyhow::Result<()> {
    let result = TestTransactionBuilder::new(existing_mock_account())
        .tx_script(upgrade_tx_script(upgrades)?)
        .account_code_upgrade(AccountCodeUpgrade::new(upgraded_mock_account_code()?))
        .build()?
        .execute()
        .await;

    assert_transaction_executor_error!(result, expected_error);

    Ok(())
}

/// Tests that a new account cannot be upgraded, since its code is part of its creation.
#[tokio::test]
async fn test_account_upgrade_rejects_new_account() -> anyhow::Result<()> {
    let mut builder = MockChain::builder();
    let account = builder.create_new_mock_account(Auth::IncrNonce)?;
    let mut mock_chain = builder.build()?;
    mock_chain.prove_next_block()?;

    let result = mock_chain
        .build_transaction(account)
        .tx_script(upgrade_tx_script([(
            upgraded_mock_account_code()?.commitment(),
            Word::empty(),
        )])?)
        .build()?
        .execute()
        .await;

    assert_transaction_executor_error!(result, ERR_ACCOUNT_UPGRADE_NOT_ALLOWED_FOR_NEW_ACCOUNT);

    Ok(())
}

/// Tests that the host rejects an upgrade to code that the advice map does not provide.
#[tokio::test]
async fn test_account_upgrade_rejects_missing_code() -> anyhow::Result<()> {
    let upgraded_code = upgraded_mock_account_code()?;

    let result = TestTransactionBuilder::new(existing_mock_account())
        .tx_script(upgrade_tx_script([(upgraded_code.commitment(), Word::empty())])?)
        .build()?
        .execute()
        .await;

    assert_transaction_executor_error!(
        result,
        matches ExecutionError::EventError { error: ref event_err, .. }
            if matches!(
                event_err.downcast_ref::<TransactionKernelError>(),
                Some(TransactionKernelError::AccountCodeUpgradeMissing(new_code_commitment))
                    if *new_code_commitment == upgraded_code.commitment()
            )
    );

    Ok(())
}

/// Tests that the host rejects an upgrade whose advice map entry does not encode valid code.
#[tokio::test]
async fn test_account_upgrade_rejects_invalid_code() -> anyhow::Result<()> {
    let upgraded_code = upgraded_mock_account_code()?;

    let result = TestTransactionBuilder::new(existing_mock_account())
        .tx_script(upgrade_tx_script([(upgraded_code.commitment(), Word::empty())])?)
        .add_advice_map_entry(
            AccountCodeUpgrade::advice_map_key(upgraded_code.commitment()),
            vec![Felt::ONE],
        )
        .build()?
        .execute()
        .await;

    assert_transaction_executor_error!(
        result,
        matches ExecutionError::EventError { error: ref event_err, .. }
            if matches!(
                event_err.downcast_ref::<TransactionKernelError>(),
                Some(TransactionKernelError::AccountCodeUpgradeInvalid { new_code_commitment, .. })
                    if *new_code_commitment == upgraded_code.commitment()
            )
    );

    Ok(())
}

/// Tests that the host rejects an upgrade whose advice map entry provides code with a different
/// commitment.
#[tokio::test]
async fn test_account_upgrade_rejects_code_not_matching_commitment() -> anyhow::Result<()> {
    let upgraded_code = upgraded_mock_account_code()?;
    let provided_code = AccountCode::mock();
    assert_ne!(upgraded_code.commitment(), provided_code.commitment());

    let (_, provided_code_data) =
        AccountCodeUpgrade::new(provided_code.clone()).to_advice_map_entry();
    let result = TestTransactionBuilder::new(existing_mock_account())
        .tx_script(upgrade_tx_script([(upgraded_code.commitment(), Word::empty())])?)
        .add_advice_map_entry(
            AccountCodeUpgrade::advice_map_key(upgraded_code.commitment()),
            provided_code_data,
        )
        .build()?
        .execute()
        .await;

    assert_transaction_executor_error!(
        result,
        matches ExecutionError::EventError { error: ref event_err, .. }
            if matches!(
                event_err.downcast_ref::<TransactionKernelError>(),
                Some(TransactionKernelError::AccountCodeUpgradeCommitmentMismatch {
                    expected,
                    actual,
                }) if *expected == upgraded_code.commitment()
                    && *actual == provided_code.commitment()
            )
    );

    Ok(())
}

/// Tests that the upgrade rejects procedures that do not match the new code commitment, so that a
/// malicious host cannot swap the code the transaction commits to.
#[tokio::test]
async fn test_account_upgrade_rejects_procedures_not_matching_commitment() -> anyhow::Result<()> {
    let upgraded_code = upgraded_mock_account_code()?;

    // Without lazy loading, the host does not provide the procedures of the new code, so the kernel
    // loads the mismatching procedures from the advice map.
    let mock_tx = TestTransactionBuilder::new(existing_mock_account())
        .add_advice_map_entry(upgraded_code.commitment(), AccountCode::mock().to_elements())
        .build()?;

    let result = mock_tx
        .execute_code_without_lazy_loading(&upgrade_program(upgraded_code.commitment()))
        .await;

    assert_execution_error!(result, ERR_ACCOUNT_CODE_COMMITMENT_MISMATCH);

    Ok(())
}

/// Tests that the upgrade rejects procedures that match the new code commitment but are not valid
/// account procedures, which only a malicious host could provide.
#[rstest::rstest]
#[case::not_enough_procedures(sorted_procedure_roots(1), ERR_ACCOUNT_NOT_ENOUGH_PROCEDURES)]
#[case::unsorted_procedures(
    {
        let mut procedure_roots = sorted_procedure_roots(4);
        procedure_roots.swap(1, 2);
        procedure_roots
    },
    ERR_ACCOUNT_PROCEDURES_MUST_BE_SORTED_AND_UNIQUE,
)]
#[case::duplicated_auth_procedure(
    {
        let mut procedure_roots = sorted_procedure_roots(4);
        procedure_roots[3] = procedure_roots[0];
        procedure_roots
    },
    ERR_ACCOUNT_AUTH_PROCEDURE_MUST_NOT_BE_DUPLICATED,
)]
#[tokio::test]
async fn test_account_upgrade_rejects_invalid_procedures(
    #[case] procedure_roots: Vec<AccountProcedureRoot>,
    #[case] expected_error: MasmError,
) -> anyhow::Result<()> {
    let procedure_elements: Vec<Felt> = procedure_roots
        .iter()
        .flat_map(AccountProcedureRoot::as_elements)
        .copied()
        .collect();
    let new_code_commitment = Hasher::hash_elements(&procedure_elements);

    // Without lazy loading, the host does not provide the procedures of the new code, so the kernel
    // loads the invalid procedures from the advice map.
    let mock_tx = TestTransactionBuilder::new(existing_mock_account())
        .add_advice_map_entry(new_code_commitment, procedure_elements)
        .build()?;

    let result = mock_tx
        .execute_code_without_lazy_loading(&upgrade_program(new_code_commitment))
        .await;

    assert_execution_error!(result, expected_error);

    Ok(())
}

/// Tests that a code upgrade without a nonce increment is rejected.
#[tokio::test]
async fn test_account_upgrade_requires_nonce_increment() -> anyhow::Result<()> {
    let upgraded_code = upgraded_mock_account_code()?;

    let result = TestTransactionBuilder::with_noop_auth_account()
        .tx_script(upgrade_tx_script([(upgraded_code.commitment(), Word::empty())])?)
        .account_code_upgrade(AccountCodeUpgrade::new(upgraded_code))
        .build()?
        .execute()
        .await;

    assert_transaction_executor_error!(
        result,
        ERR_ACCOUNT_PATCH_NONCE_MUST_BE_INCREMENTED_IF_STATE_CHANGED
    );

    Ok(())
}

/// Tests that `account_upgrade` is gated by the account context: invoking it from outside an
/// account procedure (so that `caller` is not a procedure of the account) must be rejected by the
/// authenticator. This models what an untrusted transaction script could attempt.
#[tokio::test]
async fn test_native_account_upgrade_from_tx_script_is_rejected() -> anyhow::Result<()> {
    let code_upgrade_commitment = Word::from([1, 2, 3, 4u32]);
    let storage_upgrade_commitment = Word::from([5, 6, 7, 8u32]);

    // A transaction script invokes `native_account::upgrade` directly. Its caller is the tx script,
    // not a procedure of the account, so authentication must reject the invocation.
    let tx_script_source = format!(
        r#"
        use miden::protocol::native_account

        @transaction_script
        pub proc main
            push.{storage_upgrade_commitment}
            push.{code_upgrade_commitment}
            # => [NEW_CODE_COMMITMENT, STORAGE_UPGRADE_COMMITMENT]

            exec.native_account::upgrade
        end
        "#,
        code_upgrade_commitment = code_upgrade_commitment,
        storage_upgrade_commitment = storage_upgrade_commitment,
    );
    let tx_script = CodeBuilder::with_mock_packages().compile_tx_script(tx_script_source)?;

    let mock_tx = TestTransactionBuilder::with_existing_mock_account()
        .tx_script(tx_script)
        .build()?;

    let execution_result = mock_tx.execute().await;

    // The tx-script root is not a procedure of the account, so `authenticate_account_origin` fails
    // while resolving the caller's procedure index.
    assert_transaction_executor_error!(
        execution_result,
        matches ExecutionError::EventError { error: ref event_err, .. }
            if matches!(
                event_err.downcast_ref::<TransactionKernelError>(),
                Some(TransactionKernelError::UnknownAccountProcedure(_))
            )
    );

    Ok(())
}

/// This test creates an account with mock storage slots and calls the
/// `compute_storage_commitment` procedure each time the storage is updated.
///
/// Namely, we invoke the `mock_account::compute_storage_commitment` procedure:
/// - Right after the account creation.
/// - After updating the 0th storage slot (value slot).
/// - Right after the previous call to make sure it returns the same commitment from the cached
///   data.
/// - After updating the 2nd storage slot (map slot).
#[tokio::test]
async fn test_compute_storage_commitment() -> anyhow::Result<()> {
    let mock_tx = TestTransactionBuilder::with_existing_mock_account().build().unwrap();
    let mut account_clone = mock_tx.account().clone();
    let account_storage = account_clone.storage_mut();

    let init_storage_commitment = account_storage.to_commitment();

    let mock_value_slot0 = &*MOCK_VALUE_SLOT0;
    let value_slot0 = Word::from([9, 10, 11, 12u32]);

    let mock_map_slot = &*MOCK_MAP_SLOT;
    let map_key = StorageMapKey::from_array([101, 102, 103, 104u32]);
    let map_value = Word::from([5, 6, 7, 8u32]);

    account_storage.set_item(mock_value_slot0, value_slot0)?;
    let storage_commitment_value = account_storage.to_commitment();

    account_storage.set_map_item(mock_map_slot, map_key, map_value)?;
    let storage_commitment_map = account_storage.to_commitment();

    let code = format!(
        r#"
        use miden::tx_kernel_core::prologue
        use mock::account as mock_account

        const MOCK_VALUE_SLOT0=word("{mock_value_slot0}")
        const MOCK_MAP_SLOT=word("{mock_map_slot}")

        begin
            exec.prologue::prepare_transaction

            # assert the correctness of the initial storage commitment
            call.mock_account::compute_storage_commitment
            push.{init_storage_commitment}
            assert_eqw.err="storage commitment at the beginning of the transaction is not equal to the expected one"

            # update the value storage slot
            push.{value_slot0}
            push.MOCK_VALUE_SLOT0[0..2]
            call.mock_account::set_item dropw drop
            # => []

            # assert the correctness of the storage commitment after the value slot was updated
            call.mock_account::compute_storage_commitment
            push.{storage_commitment_value}
            assert_eqw.err="storage commitment after the value slot was updated is not equal to the expected one"

            # get the storage commitment once more to get the cached data and assert that this data
            # didn't change
            call.mock_account::compute_storage_commitment
            push.{storage_commitment_value}
            assert_eqw.err="storage commitment should remain the same"

            # update the map storage slot
            push.{map_value}
            push.{map_key}
            push.MOCK_MAP_SLOT[0..2]
            # => [slot_id_suffix, slot_id_prefix, KEY, VALUE]

            call.mock_account::set_map_item dropw dropw
            # => []

            # assert the correctness of the storage commitment after the map slot was updated
            call.mock_account::compute_storage_commitment
            push.{storage_commitment_map}
            assert_eqw.err="storage commitment after the map slot was updated is not equal to the expected one"
        end
        "#,
    );

    mock_tx.execute_code(&code).await?;

    Ok(())
}

/// Tests that an account with a non-empty map can be created.
///
/// In particular, this tests the account delta logic for (non-empty) storage slots for _new_
/// accounts.
#[tokio::test]
async fn prove_account_creation_with_non_empty_storage() -> anyhow::Result<()> {
    let slot_name0 = StorageSlotName::mock(0);
    let slot_name1 = StorageSlotName::mock(1);
    let slot_name2 = StorageSlotName::mock(2);

    let slot0 = StorageSlot::with_value(slot_name0.clone(), Word::from([1, 2, 3, 4u32]));
    let slot1 = StorageSlot::with_value(slot_name1.clone(), Word::from([10, 20, 30, 40u32]));
    let mut map_entries = Vec::new();
    for _ in 0..10 {
        map_entries.push((StorageMapKey::from_raw(random::<Word>()), random::<Word>()));
    }
    let map_slot =
        StorageSlot::with_map(slot_name2.clone(), StorageMap::with_entries(map_entries.clone())?);

    let account = AccountBuilder::new([6; 32])
        .account_type(AccountType::Public)
        .with_components(Auth::IncrNonce)
        .with_component(MockAccountComponent::with_slots(vec![
            slot0.clone(),
            slot1.clone(),
            map_slot,
        ]))
        .build()?;

    let tx = TestTransactionBuilder::new(account)
        .build()?
        .execute()
        .await
        .context("failed to execute account-creating transaction")?;

    assert_eq!(
        tx.account_patch().final_nonce(),
        Some(Felt::ONE),
        "new account should have nonce 1"
    );

    assert_matches!(
        tx.account_patch().storage().get(&slot_name0).unwrap(),
        StorageSlotPatch::Value(StorageValuePatch::Create { value }) => {
            assert_eq!(*value, slot0.value())
        }
    );
    assert_matches!(
        tx.account_patch().storage().get(&slot_name1).unwrap(),
        StorageSlotPatch::Value(StorageValuePatch::Create { value }) => {
            assert_eq!(*value, slot1.value())
        }
    );
    assert_matches!(
        tx.account_patch().storage().get(&slot_name2).unwrap(),
        StorageSlotPatch::Map(StorageMapPatch::Create { entries }) => {
            let expected = &BTreeMap::from_iter(map_entries);
            assert_eq!(expected, entries.as_map())
        }
    );

    assert!(tx.account_patch().vault().is_empty());
    assert_eq!(tx.final_account().nonce(), Felt::ONE);

    let proven_tx = LocalTransactionProver::default().prove(tx.clone())?;

    // The patch should be present on the proven tx.
    let patch = proven_tx.account_update().details().unwrap_public();
    assert_eq!(patch, tx.account_patch());

    Ok(())
}

// ACCOUNT VAULT TESTS
// ================================================================================================

#[tokio::test]
async fn test_get_vault_root() -> anyhow::Result<()> {
    let mock_tx = TestTransactionBuilder::with_existing_mock_account().build()?;

    let mut account = mock_tx.account().clone();

    let fungible_asset = Asset::from(
        FungibleAsset::new(
            AccountId::try_from(ACCOUNT_ID_PUBLIC_FUNGIBLE_FAUCET).context("id should be valid")?,
            5,
        )
        .context("fungible_asset_0 is invalid")?,
    );

    // get the initial vault root
    let code = format!(
        r#"
        use mock::account as mock_account
        use miden::tx_kernel_core::prologue

        begin
            exec.prologue::prepare_transaction

            # get the initial vault root
            call.mock_account::get_initial_vault_root
            # => [INIT_VAULT_ROOT, pad(12)]
            swapdw dropw dropw swapw dropw
            # => [INIT_VAULT_ROOT]
            push.{expected_vault_root}
            assert_eqw.err="initial vault root mismatch"
        end
        "#,
        expected_vault_root = account.vault().root(),
    );
    mock_tx.execute_code(&code).await?;

    // get the current vault root
    account.vault_mut().add_asset(fungible_asset)?;

    let code = format!(
        r#"
        use miden::protocol::active_account
        use miden::tx_kernel_core::prologue
        use mock::account as mock_account

        begin
            exec.prologue::prepare_transaction

            # add an asset to the account
            push.{FUNGIBLE_ASSET_VALUE}
            push.{FUNGIBLE_ASSET_ID}
            call.mock_account::add_asset
            dropw dropw
            # => []

            # get the current vault root
            call.mock_account::get_vault_root
            # => [VAULT_ROOT, pad(12)]
            swapdw dropw dropw swapw dropw
            # => [VAULT_ROOT]
            push.{expected_vault_root}
            assert_eqw.err="vault root mismatch"
        end
        "#,
        FUNGIBLE_ASSET_VALUE = fungible_asset.to_value_word(),
        FUNGIBLE_ASSET_ID = fungible_asset.to_id_word(),
        expected_vault_root = account.vault().root(),
    );
    mock_tx.execute_code(&code).await?;

    Ok(())
}

/// This test checks the correctness of the
/// `miden::standards::assets::fungible_asset::get_initial_native_account_balance` procedure in two
/// cases:
/// - when a note adds the asset which already exists in the account vault.
/// - when a note adds the asset which doesn't exist in the account vault.
///
/// As part of the test pipeline it also checks the correctness of the
/// `miden::standards::assets::fungible_asset::get_active_account_balance` procedure.
#[tokio::test]
async fn test_get_init_balance_addition() -> anyhow::Result<()> {
    // prepare the testing data
    // ------------------------------------------
    let mut builder = MockChain::builder();

    let faucet_existing_asset =
        AccountId::try_from(ACCOUNT_ID_PUBLIC_FUNGIBLE_FAUCET).context("id should be valid")?;
    let faucet_new_asset =
        AccountId::try_from(ACCOUNT_ID_PUBLIC_FUNGIBLE_FAUCET_1).context("id should be valid")?;

    let fungible_asset_for_account = Asset::from(
        FungibleAsset::new(faucet_existing_asset, 10).context("fungible_asset_0 is invalid")?,
    );
    let account = builder.add_existing_mock_account_with_assets(
        crate::Auth::BasicAuth {
            auth_scheme: AuthScheme::Falcon512Poseidon2,
        },
        [fungible_asset_for_account],
    )?;

    let fungible_asset_for_note_existing = Asset::from(
        FungibleAsset::new(faucet_existing_asset, 7).context("fungible_asset_0 is invalid")?,
    );

    let fungible_asset_for_note_new = Asset::from(
        FungibleAsset::new(faucet_new_asset, 20).context("fungible_asset_1 is invalid")?,
    );

    let p2id_note_existing_asset = builder.add_p2id_note(
        ACCOUNT_ID_SENDER.try_into().unwrap(),
        account.id(),
        &[fungible_asset_for_note_existing],
        NoteType::Public,
    )?;
    let p2id_note_new_asset = builder.add_p2id_note(
        ACCOUNT_ID_SENDER.try_into().unwrap(),
        account.id(),
        &[fungible_asset_for_note_new],
        NoteType::Public,
    )?;

    let mut mock_chain = builder.build()?;
    mock_chain.prove_next_block()?;

    // case 1: existing asset was added to the account
    // ------------------------------------------

    let asset_id = AssetId::new_fungible(faucet_existing_asset);
    let initial_balance = account.vault().get_balance(asset_id)?.as_u64();

    let add_existing_source = format!(
        r#"
        use miden::core::sys
        use mock::account as mock_account

        @transaction_script
        pub proc main
            # get the current asset balance
            push.{ASSET_ID}
            call.mock_account::get_active_account_balance
            # => [final_balance, pad(15)]

            # assert final balance is correct
            push.{final_balance}
            assert_eq.err="final balance is incorrect"
            # => []

            # get the initial asset balance
            push.{ASSET_ID}
            call.mock_account::get_initial_native_account_balance
            # => [init_balance, pad(15)]

            # assert initial balance is correct
            push.{initial_balance}
            assert_eq.err="initial balance is incorrect"

            # truncate the stack
            exec.sys::truncate_stack
        end
    "#,
        ASSET_ID = asset_id.to_word(),
        final_balance =
            initial_balance + fungible_asset_for_note_existing.unwrap_fungible().amount().as_u64(),
    );

    let tx_script = CodeBuilder::with_mock_packages().compile_tx_script(add_existing_source)?;

    let mock_tx = mock_chain
        .build_transaction(account.id())
        .unauthenticated_input_note(p2id_note_existing_asset)
        .tx_script(tx_script)
        .build()?;

    mock_tx.execute().await?;

    // case 2: new asset was added to the account
    // ------------------------------------------

    let asset_id = AssetId::new_fungible(faucet_new_asset);
    let initial_balance = account.vault().get_balance(asset_id)?.as_u64();

    let add_new_source = format!(
        r#"
        use miden::core::sys
        use mock::account as mock_account

        @transaction_script
        pub proc main
            # get the current asset balance
            push.{ASSET_ID}
            call.mock_account::get_active_account_balance
            # => [final_balance, pad(15)]

            # assert final balance is correct
            push.{final_balance}
            assert_eq.err="final balance is incorrect"
            # => []

            # get the initial asset balance
            push.{ASSET_ID}
            call.mock_account::get_initial_native_account_balance
            # => [init_balance, pad(15)]

            # assert initial balance is correct
            push.{initial_balance}
            assert_eq.err="initial balance is incorrect"

            # truncate the stack
            exec.sys::truncate_stack
        end
    "#,
        ASSET_ID = asset_id.to_word(),
        final_balance =
            initial_balance + fungible_asset_for_note_new.unwrap_fungible().amount().as_u64(),
    );

    let tx_script = CodeBuilder::with_mock_packages().compile_tx_script(add_new_source)?;

    let mock_tx = mock_chain
        .build_transaction(account.id())
        .unauthenticated_input_note(p2id_note_new_asset)
        .tx_script(tx_script)
        .build()?;

    mock_tx.execute().await?;

    Ok(())
}

/// This test checks the correctness of the
/// `miden::standards::assets::fungible_asset::get_initial_native_account_balance` procedure when an
/// asset is moved from the vault to a note.
///
/// As part of the test pipeline it also checks the correctness of the
/// `miden::standards::assets::fungible_asset::get_active_account_balance` procedure.
#[tokio::test]
async fn test_get_init_balance_subtraction() -> anyhow::Result<()> {
    let mut builder = MockChain::builder();

    let faucet_existing_asset =
        AccountId::try_from(ACCOUNT_ID_PUBLIC_FUNGIBLE_FAUCET).context("id should be valid")?;

    let fungible_asset_for_account = Asset::from(
        FungibleAsset::new(faucet_existing_asset, 10).context("fungible_asset_0 is invalid")?,
    );
    let account = builder.add_existing_mock_account_with_assets(
        crate::Auth::BasicAuth {
            auth_scheme: AuthScheme::Falcon512Poseidon2,
        },
        [fungible_asset_for_account],
    )?;

    let fungible_asset_for_note_existing = Asset::from(
        FungibleAsset::new(faucet_existing_asset, 7).context("fungible_asset_0 is invalid")?,
    );

    let mut mock_chain = builder.build()?;
    mock_chain.prove_next_block()?;

    let asset_id = AssetId::new_fungible(faucet_existing_asset);
    let initial_balance = account.vault().get_balance(asset_id)?.as_u64();

    let expected_output_note =
        create_public_p2any_note(ACCOUNT_ID_SENDER.try_into()?, [fungible_asset_for_note_existing]);

    let remove_existing_source = format!(
        r#"
        use miden::core::sys
        use mock::account as mock_account
        use mock::util

        @transaction_script
        pub proc main
            # create random note and move the asset into it
            exec.util::create_default_note
            # => [note_idx]

            push.{REMOVED_ASSET_VALUE}
            push.{REMOVED_ASSET_ID}
            exec.util::move_asset_to_note
            # => []

            # get the current asset balance
            push.{ASSET_ID}
            call.mock_account::get_active_account_balance
            # => [final_balance, pad(15)]

            # assert final balance is correct
            push.{final_balance}
            assert_eq.err="final balance is incorrect"
            # => []

            # get the initial asset balance
            push.{ASSET_ID}
            call.mock_account::get_initial_native_account_balance
            # => [init_balance, pad(15)]

            # assert initial balance is correct
            push.{initial_balance}
            assert_eq.err="initial balance is incorrect"

            # truncate the stack
            exec.sys::truncate_stack
        end
    "#,
        REMOVED_ASSET_ID = fungible_asset_for_note_existing.to_id_word(),
        REMOVED_ASSET_VALUE = fungible_asset_for_note_existing.to_value_word(),
        ASSET_ID = asset_id.to_word(),
        final_balance =
            initial_balance - fungible_asset_for_note_existing.unwrap_fungible().amount().as_u64(),
    );

    let tx_script = CodeBuilder::with_mock_packages().compile_tx_script(remove_existing_source)?;

    let mock_tx = mock_chain
        .build_transaction(account.id())
        .tx_script(tx_script)
        .expected_output_note(RawOutputNote::Full(expected_output_note))
        .build()?;

    mock_tx.execute().await?;

    Ok(())
}

/// This test checks the correctness of the `miden::protocol::native_account::get_initial_asset`
/// procedure creating a note which removes an asset from the account vault.
///
/// As part of the test pipeline it also checks the correctness of the
/// `miden::protocol::active_account::get_asset` procedure.
#[tokio::test]
async fn test_get_init_asset() -> anyhow::Result<()> {
    let mut builder = MockChain::builder();

    let faucet_existing_asset =
        AccountId::try_from(ACCOUNT_ID_PUBLIC_FUNGIBLE_FAUCET).context("id should be valid")?;

    let fungible_asset_for_account = Asset::from(
        FungibleAsset::new(faucet_existing_asset, 10).context("fungible_asset_0 is invalid")?,
    );
    let account = builder.add_existing_mock_account_with_assets(
        crate::Auth::BasicAuth {
            auth_scheme: AuthScheme::Falcon512Poseidon2,
        },
        [fungible_asset_for_account],
    )?;

    let fungible_asset_for_note_existing = Asset::from(
        FungibleAsset::new(faucet_existing_asset, 7).context("fungible_asset_0 is invalid")?,
    );

    let mut mock_chain = builder.build()?;
    mock_chain.prove_next_block()?;

    let final_asset = fungible_asset_for_account
        .unwrap_fungible()
        .sub(fungible_asset_for_note_existing.unwrap_fungible())?;

    let expected_output_note =
        create_public_p2any_note(ACCOUNT_ID_SENDER.try_into()?, [fungible_asset_for_note_existing]);

    let remove_existing_source = format!(
        r#"
        use miden::core::sys
        use miden::protocol::active_account
        use miden::protocol::native_account
        use mock::account as mock_account
        use mock::util

        @transaction_script
        pub proc main
            # create default note and move the asset into it
            exec.util::create_default_note
            # => [note_idx]

            push.{REMOVED_ASSET_VALUE}
            push.{ASSET_ID}
            exec.util::move_asset_to_note
            # => []

            # get the current asset
            push.{ASSET_ID} call.mock_account::get_asset
            # => [ASSET_VALUE, pad(12)]

            push.{FINAL_ASSET}
            assert_eqw.err="final asset is incorrect"
            # => [pad(12)]

            # get the initial asset
            push.{ASSET_ID} call.mock_account::get_initial_asset
            # => [INITIAL_ASSET, pad(12)]

            push.{INITIAL_ASSET_VALUE}
            assert_eqw.err="initial asset is incorrect"

            # truncate the stack
            exec.sys::truncate_stack
        end
    "#,
        ASSET_ID = fungible_asset_for_note_existing.to_id_word(),
        REMOVED_ASSET_VALUE = fungible_asset_for_note_existing.to_value_word(),
        INITIAL_ASSET_VALUE = fungible_asset_for_account.to_value_word(),
        FINAL_ASSET = final_asset.to_value_word(),
    );

    let tx_script = CodeBuilder::with_mock_packages().compile_tx_script(remove_existing_source)?;

    mock_chain
        .build_transaction(account.id())
        .tx_script(tx_script)
        .expected_output_note(RawOutputNote::Full(expected_output_note))
        .build()?
        .execute()
        .await?;

    Ok(())
}

// PROCEDURE AUTHENTICATION TESTS
// ================================================================================================

#[tokio::test]
async fn test_authenticate_procedure() -> anyhow::Result<()> {
    let mock_component = MockAccountComponent::with_empty_slots();

    let components: Vec<AccountComponent> =
        Auth::IncrNonce.into_iter().chain([mock_component.into()]).collect();
    let account_code = AccountCode::from_components(&components).unwrap();

    let tc_0 = *account_code.procedures()[1].mast_root();
    let tc_1 = *account_code.procedures()[2].mast_root();
    let tc_2 = *account_code.procedures()[3].mast_root();

    let test_cases =
        vec![(tc_0, true), (tc_1, true), (tc_2, true), (Word::from([1, 0, 1, 0u32]), false)];

    for (root, valid) in test_cases.into_iter() {
        let mock_tx = TestTransactionBuilder::with_existing_mock_account().build().unwrap();

        let code = format!(
            "
            use miden::tx_kernel_core::account
            use miden::tx_kernel_core::prologue

            begin
                exec.prologue::prepare_transaction

                # authenticate procedure
                push.{root}
                exec.account::authenticate_procedure

                # truncate the stack
                dropw
            end
            ",
            root = root,
        );

        // Execution of this code will return an EventError(UnknownAccountProcedure) for procs
        // that are not in the advice provider.
        let exec_output = mock_tx.execute_code(&code).await;

        match valid {
            true => {
                assert!(exec_output.is_ok(), "A valid procedure must successfully authenticate")
            },
            false => {
                assert!(exec_output.is_err(), "An invalid procedure should fail to authenticate")
            },
        }
    }

    Ok(())
}

// PROCEDURE INTROSPECTION TESTS
// ================================================================================================

#[tokio::test]
async fn test_was_procedure_called() -> anyhow::Result<()> {
    // Create a standard account using the mock component
    let mock_component = MockAccountComponent::with_slots(AccountStorage::mock_storage_slots());
    let account = AccountBuilder::new(rand::random())
        .with_components(Auth::IncrNonce)
        .with_component(mock_component)
        .build_existing()
        .unwrap();
    let mock_value_slot1 = &*MOCK_VALUE_SLOT1;

    // Verifies the call-tracking of an account procedure: `get_item` is tracked as called once it
    // has been invoked. `was_procedure_called` is itself account-context-gated, so it is invoked
    // through the account.
    let tx_script_code = format!(
        r#"
        use miden::core::sys
        use mock::account as mock_account

        const MOCK_VALUE_SLOT1 = word("{mock_value_slot1}")

        @transaction_script
        pub proc main
            # First check that get_item procedure hasn't been called yet
            procref.mock_account::get_item
            call.mock_account::was_procedure_called
            assertz.err="procedure should not have been called"

            # Call the procedure first time
            push.MOCK_VALUE_SLOT1[0..2]
            call.mock_account::get_item dropw

            procref.mock_account::get_item
            call.mock_account::was_procedure_called
            assert.err="procedure should have been called"

            # Call the procedure second time
            push.MOCK_VALUE_SLOT1[0..2]
            call.mock_account::get_item dropw

            procref.mock_account::get_item
            call.mock_account::was_procedure_called
            assert.err="2nd call should not change the was_called flag"

            exec.sys::truncate_stack
        end
        "#
    );

    // Compile the transaction script using the testing assembler with mock account
    let tx_script = CodeBuilder::with_mock_packages().compile_tx_script(tx_script_code)?;

    // Create mock transaction and execute
    let mock_tx = TestTransactionBuilder::new(account).tx_script(tx_script).build().unwrap();

    mock_tx.execute().await?;

    Ok(())
}

/// Tests that an account can call code in a custom package when loading that package into the
/// executor.
///
/// The call chain and dependency graph in this test is:
/// `tx script -> account code -> external package`
#[tokio::test]
async fn transaction_executor_account_code_using_custom_package() -> anyhow::Result<()> {
    let slot_value = Word::from([2, 3, 4, 5u32]);
    let external_package_code = format!(
        r#"
      use miden::protocol::native_account

      const MOCK_VALUE_SLOT0 = word("{mock_value_slot0}")

      pub proc external_setter
        push.{slot_value}
        push.MOCK_VALUE_SLOT0[0..2]
        exec.native_account::set_item
        dropw dropw
      end"#,
        mock_value_slot0 = *MOCK_VALUE_SLOT0,
    );

    const ACCOUNT_COMPONENT_CODE: &str = "
      use external_package::external_module

      @account_procedure
      pub proc custom_setter
        exec.external_module::external_setter
      end";

    let source_manager = Arc::new(DefaultSourceManager::default());
    let mut parser = ModuleParser::new(Some(ModuleKind::Library));
    let external_package_root = parser
        .parse_str(
            Some(Path::new("external_package::external_module")),
            &external_package_code,
            source_manager.clone(),
        )
        .map_err(|err| {
            anyhow::anyhow!("failed to parse package: {}", PrintDiagnostic::new(&err))
        })?;
    let external_package = TransactionKernel::assembler_with_source_manager(source_manager.clone())
        .assemble_library("external-library", external_package_root, None::<&str>)
        .map_err(|err| {
            anyhow::anyhow!("failed to assemble package: {}", PrintDiagnostic::new(&err))
        })?;

    let assembler: miden_protocol::assembly::Assembler =
        CodeBuilder::with_mock_packages_with_source_manager(source_manager.clone()).into();
    let assembler =
        assembler
            .with_package(Arc::from(external_package), Linkage::Static)
            .map_err(|err| {
                anyhow::anyhow!("failed to link static package: {}", PrintDiagnostic::new(&err))
            })?;

    let account_component_root = parser
        .parse_str(
            Some(Path::new("account_component::account_module")),
            ACCOUNT_COMPONENT_CODE,
            source_manager,
        )
        .map_err(|err| {
            anyhow::anyhow!("failed to parse account component: {}", PrintDiagnostic::new(&err))
        })?;
    let account_component_lib = *assembler
        .clone()
        .assemble_library("account-component", account_component_root, None::<&str>)
        .unwrap();

    let tx_script_src = "\
          use account_component::account_module

          @transaction_script
          pub proc main
            call.account_module::custom_setter
          end";

    let account_component = AccountComponent::new(
        account_component_lib.clone(),
        AccountStorage::mock_storage_slots(),
        AccountComponentMetadata::mock("account_module"),
    )?;

    // Build an existing account with nonce 1.
    let native_account = AccountBuilder::new(rand::random())
        .with_components(Auth::IncrNonce)
        .with_component(account_component)
        .build_existing()?;

    let tx_script = CodeBuilder::default()
        .with_dynamically_linked_package(&account_component_lib)?
        .compile_tx_script(tx_script_src)?;

    let mock_tx = TestTransactionBuilder::new(native_account.clone())
        .tx_script(tx_script)
        .build()
        .unwrap();

    let executed_tx = mock_tx.execute().await?;

    // Account's initial nonce of 1 should have been incremented by 1.
    assert_eq!(
        executed_tx.account_patch().final_nonce(),
        Some(native_account.nonce() + Felt::ONE)
    );

    // Make sure that account storage has been updated as per the tx script call.
    assert_eq!(executed_tx.account_patch().storage().values().count(), 1);
    assert_eq!(
        executed_tx.account_patch().storage().get(&MOCK_VALUE_SLOT0).unwrap(),
        &StorageSlotPatch::Value(StorageValuePatch::Update { value: slot_value }),
    );
    Ok(())
}

/// Tests that incrementing the account nonce twice fails.
#[tokio::test]
async fn incrementing_nonce_twice_fails() -> anyhow::Result<()> {
    let source_code = "
        use miden::protocol::native_account

        @auth_script
        pub proc auth_incr_nonce_twice
            exec.native_account::incr_nonce drop
            exec.native_account::incr_nonce drop
        end
    ";

    let faulty_auth_code =
        CodeBuilder::default().compile_component_code("test::faulty_auth", source_code)?;
    let faulty_auth_component = AccountComponent::new(
        faulty_auth_code,
        vec![],
        AccountComponentMetadata::mock("test::faulty_auth"),
    )?;
    let account = AccountBuilder::new([5; 32])
        .with_component(faulty_auth_component)
        .with_component(MockAccountComponent::with_empty_slots())
        .build()
        .context("failed to build account")?;

    let result = TestTransactionBuilder::new(account).build()?.execute().await;

    assert_transaction_executor_error!(result, ERR_ACCOUNT_NONCE_CAN_ONLY_BE_INCREMENTED_ONCE);

    Ok(())
}

#[tokio::test]
async fn test_has_procedure() -> anyhow::Result<()> {
    // Create a standard account using the mock component
    let mock_component = MockAccountComponent::with_slots(AccountStorage::mock_storage_slots());
    let account = AccountBuilder::new(rand::random())
        .with_components(Auth::IncrNonce)
        .with_component(mock_component)
        .build_existing()
        .unwrap();

    let tx_script_code = r#"
        use miden::core::sys
        use mock::account as mock_account

        @transaction_script
        pub proc main
            # check that get_item procedure is available on the mock account
            procref.mock_account::get_item
            # => [GET_ITEM_ROOT]

            call.mock_account::has_procedure
            # => [is_procedure_available, pad(15)]

            # assert that the get_item is exposed
            assert.err="get_item procedure should be exposed by the mock account"

            # get some random word and assert that it is not exposed
            push.5.3.15.686

            call.mock_account::has_procedure
            # => [is_procedure_available, pad(15)]

            # assert that the procedure with some random root is not exposed
            assertz.err="procedure with some random root should not be exposed by the mock account"

            exec.sys::truncate_stack
        end
        "#;

    // Compile the transaction script using the testing assembler with mock account
    let tx_script = CodeBuilder::with_mock_packages()
        .compile_tx_script(tx_script_code)
        .map_err(|err| anyhow::anyhow!("{err}"))?;

    // Create mock transaction and execute
    let mock_tx = TestTransactionBuilder::new(account).tx_script(tx_script).build().unwrap();

    mock_tx.execute().await?;

    Ok(())
}

#[tokio::test]
async fn test_has_storage_slot() -> anyhow::Result<()> {
    let existing_slot_name = format!("{}", AccountStorage::mock_value_slot0().name());

    // (slot name, whether a slot with that name is expected to exist on the account)
    let test_cases = [(existing_slot_name.as_str(), true), ("unknown::slot::name", false)];

    for (slot_name, expected_to_exist) in test_cases {
        let mock_tx = TestTransactionBuilder::with_existing_mock_account().build().unwrap();

        let assertion = if expected_to_exist {
            r#"assert.err="installed storage slot should be reported as present""#
        } else {
            r#"assertz.err="unknown storage slot should be reported as absent""#
        };

        let code = format!(
            r#"
            use miden::core::sys

            use miden::tx_kernel_core::prologue
            use mock::account as mock_account

            const SLOT_NAME = word("{slot_name}")

            begin
                exec.prologue::prepare_transaction

                # pad the stack for the call
                push.SLOT_NAME[0..2]
                repeat.14 push.0 movdn.2 end
                # => [slot_id_suffix, slot_id_prefix, pad(14)]

                call.mock_account::has_storage_slot
                # => [has_slot, pad(15)]

                {assertion}

                exec.sys::truncate_stack
            end
            "#,
        );

        mock_tx.execute_code(&code).await?;
    }

    Ok(())
}

// ACCOUNT INITIAL STORAGE TESTS
// ================================================================================================

#[tokio::test]
async fn test_get_initial_item() -> anyhow::Result<()> {
    let mock_tx = TestTransactionBuilder::with_existing_mock_account().build().unwrap();

    // Test that get_initial_item returns the initial value before any changes
    let code = format!(
        r#"
        use miden::tx_kernel_core::account
        use miden::tx_kernel_core::prologue
        use mock::account as mock_account

        const MOCK_VALUE_SLOT0 = word("{mock_value_slot0}")

        begin
            exec.prologue::prepare_transaction

            # get initial value of the storage slot
            push.MOCK_VALUE_SLOT0[0..2]
            exec.account::get_initial_item

            push.{expected_initial_value}
            assert_eqw.err="initial value should match expected"

            # modify the storage slot
            push.9.10.11.12
            push.MOCK_VALUE_SLOT0[0..2]
            call.mock_account::set_item dropw drop drop

            # get_item should return the new value
            push.MOCK_VALUE_SLOT0[0..2]
            exec.account::get_item
            push.9.10.11.12
            assert_eqw.err="current value should be updated"

            # get_initial_item should still return the initial value
            push.MOCK_VALUE_SLOT0[0..2]
            exec.account::get_initial_item
            push.{expected_initial_value}
            assert_eqw.err="initial value should remain unchanged"
        end
        "#,
        mock_value_slot0 = *MOCK_VALUE_SLOT0,
        expected_initial_value = AccountStorage::mock_value_slot0().content().value(),
    );

    mock_tx.execute_code(&code).await?;

    Ok(())
}

#[tokio::test]
async fn test_get_initial_map_item() -> anyhow::Result<()> {
    let map_slot = AccountStorage::mock_map_slot();
    let account = AccountBuilder::new(rand::random())
        .with_components(Auth::IncrNonce)
        .with_component(MockAccountComponent::with_slots(vec![map_slot.clone()]))
        .build_existing()
        .unwrap();

    let mock_tx = TestTransactionBuilder::new(account).build().unwrap();

    // Use the first key-value pair from the mock storage
    let StorageSlotContent::Map(map) = map_slot.content() else {
        panic!("expected map");
    };

    let (initial_key, initial_value) = map.entries().next().unwrap();
    let new_key = Word::from([201, 202, 203, 204u32]);
    let new_value = Word::from([301, 302, 303, 304u32]);
    let mock_map_slot = map_slot.name();

    let code = format!(
        r#"
        use miden::tx_kernel_core::prologue
        use mock::account as mock_account

        const MOCK_MAP_SLOT = word("{mock_map_slot}")

        begin
            exec.prologue::prepare_transaction

            # get initial value from map
            push.{initial_key}
            push.MOCK_MAP_SLOT[0..2]
            call.mock_account::get_initial_map_item
            push.{initial_value}
            assert_eqw.err="initial map value should match expected"

            # add a new key-value pair to the map
            push.{new_value}
            push.{new_key}
            push.MOCK_MAP_SLOT[0..2]
            call.mock_account::set_map_item dropw dropw

            # get_map_item should return the new value
            push.{new_key}
            push.MOCK_MAP_SLOT[0..2]
            call.mock_account::get_map_item
            push.{new_value}
            assert_eqw.err="current map value should be updated"

            # get_initial_map_item should still return the initial value for the initial key
            push.{initial_key}
            push.MOCK_MAP_SLOT[0..2]
            call.mock_account::get_initial_map_item
            push.{initial_value}
            assert_eqw.err="initial map value should remain unchanged"

            # get_initial_map_item for the new key should return empty word (default)
            push.{new_key}
            push.MOCK_MAP_SLOT[0..2]
            call.mock_account::get_initial_map_item
            padw
            assert_eqw.err="new key should have empty initial value"

            dropw dropw dropw
        end
        "#,
        initial_key = initial_key,
        initial_value = initial_value,
        new_key = new_key,
        new_value = new_value,
    );

    mock_tx.execute_code(&code).await.unwrap();

    Ok(())
}

/// Tests that `get_initial_item` returns the original slot values and `get_item` returns updated
/// values after modification, for all possible storage slot indices.
#[tokio::test]
async fn test_get_item_and_get_initial_item_for_all_slots() -> anyhow::Result<()> {
    // Build storage slots for all valid indices.
    let slots: Vec<StorageSlot> = (0..AccountStorage::MAX_NUM_STORAGE_SLOTS as u32)
        .map(|index| {
            StorageSlot::with_value(
                StorageSlotName::mock(index as usize),
                Word::from([0, 0, 0, index]),
            )
        })
        .collect();

    let account = AccountBuilder::new(rand::random())
        .with_components(Auth::IncrNonce)
        .with_component(MockAccountComponent::with_slots(slots.clone()))
        .build_existing()
        .unwrap();

    let mock_tx = TestTransactionBuilder::new(account).build().unwrap();

    // Build MASM code that, for each slot:
    // 1. Sets a new value [index, 0, 0, 0]
    // 2. Asserts get_initial_item returns the original value [0, 0, 0, index]
    // 3. Asserts get_item returns the new value [index, 0, 0, 0]
    let mut slot_constants = String::new();
    let mut slot_operations = String::new();

    for (index, slot) in slots.iter().enumerate() {
        let slot_name = slot.name();
        let initial_value = slot.value();
        // Use a different format than the initial value (index at word position 0).
        let new_value = Word::from([index as u32, 0, 0, 0]);
        let const_name = format!("SLOT_{index}");

        slot_constants.push_str(&format!("const {const_name} = word(\"{slot_name}\")\n"));

        slot_operations.push_str(&format!(
            r#"
                # slot {index}: set new value
                push.{new_value}
                push.{const_name}[0..2]
                call.mock_account::set_item dropw drop drop

                # slot {index}: assert get_initial_item returns original value
                push.{const_name}[0..2]
                exec.account::get_initial_item
                push.{initial_value}
                assert_eqw.err="slot {index}: initial value mismatch"

                # slot {index}: assert get_item returns the new value
                push.{const_name}[0..2]
                exec.account::get_item
                push.{new_value}
                assert_eqw.err="slot {index}: current value mismatch"
            "#,
        ));
    }

    let code = format!(
        r#"
        use miden::tx_kernel_core::account
        use miden::tx_kernel_core::prologue
        use mock::account as mock_account

        {slot_constants}

        begin
            exec.prologue::prepare_transaction
            {slot_operations}
        end
        "#,
    );

    mock_tx.execute_code(&code).await?;

    Ok(())
}

/// Tests that incrementing the account nonce fails if it would overflow the field.
#[tokio::test]
async fn incrementing_nonce_overflow_fails() -> anyhow::Result<()> {
    let mut account = AccountBuilder::new([42; 32])
        .with_components(Auth::IncrNonce)
        .with_component(MockAccountComponent::with_empty_slots())
        .build_existing()
        .context("failed to build account")?;
    // Increment the nonce to the maximum felt value. The nonce is already 1, so we increment by
    // modulus - 2.
    account.increment_nonce(Felt::new_unchecked(Felt::ORDER_U64 - 2))?;

    let result = TestTransactionBuilder::new(account).build()?.execute().await;

    assert_transaction_executor_error!(result, ERR_ACCOUNT_NONCE_AT_MAX);

    Ok(())
}

/// Tests that merging two components that have a procedure with the same mast root
/// (`get_slot_content`) works.
///
/// Asserts that the procedure is callable via both names.
#[tokio::test]
async fn merging_components_with_same_mast_root_succeeds() -> anyhow::Result<()> {
    static TEST_SLOT_NAME: LazyLock<StorageSlotName> = LazyLock::new(|| {
        StorageSlotName::new("miden::slot::test").expect("storage slot name should be valid")
    });

    static COMPONENT_1_PACKAGE: LazyLock<Package> = LazyLock::new(|| {
        let code = format!(
            r#"
              use miden::protocol::active_account

              const TEST_SLOT_NAME = word("{test_slot_name}")

              @account_procedure
              pub proc get_slot_content
                  push.TEST_SLOT_NAME[0..2]
                  exec.active_account::get_item
                  swapw dropw
              end
            "#,
            test_slot_name = *TEST_SLOT_NAME
        );

        let source_manager = Arc::new(DefaultSourceManager::default());
        let root = ModuleParser::new(Some(ModuleKind::Library))
            .parse_str(Some(Path::new("component1::interface")), &code, source_manager.clone())
            .expect("mock account code should parse");
        *TransactionKernel::assembler_with_source_manager(source_manager)
            .assemble_library("component1", root, None::<&str>)
            .expect("mock account code should be valid")
    });

    static COMPONENT_2_PACKAGE: LazyLock<Package> = LazyLock::new(|| {
        let code = format!(
            r#"
              use miden::protocol::active_account
              use miden::protocol::native_account

              const TEST_SLOT_NAME = word("{test_slot_name}")

              @account_procedure
              pub proc get_slot_content
                  push.TEST_SLOT_NAME[0..2]
                  exec.active_account::get_item
                  swapw dropw
              end

              @account_procedure
              pub proc set_slot_content
                  push.[5,6,7,8]
                  push.TEST_SLOT_NAME[0..2]
                  exec.native_account::set_item
                  swapw dropw
              end
            "#,
            test_slot_name = *TEST_SLOT_NAME
        );

        let source_manager = Arc::new(DefaultSourceManager::default());
        let root = ModuleParser::new(Some(ModuleKind::Library))
            .parse_str(Some(Path::new("component2::interface")), &code, source_manager.clone())
            .expect("mock account code should parse");
        *TransactionKernel::assembler_with_source_manager(source_manager)
            .assemble_library("component2", root, None::<&str>)
            .expect("mock account code should be valid")
    });

    struct CustomComponent1 {
        slot: StorageSlot,
    }

    impl From<CustomComponent1> for AccountComponent {
        fn from(component: CustomComponent1) -> AccountComponent {
            AccountComponent::new(
                COMPONENT_1_PACKAGE.clone(),
                vec![component.slot],
                AccountComponentMetadata::mock("component1::interface"),
            )
            .expect("should be valid")
        }
    }

    struct CustomComponent2;

    impl From<CustomComponent2> for AccountComponent {
        fn from(_component: CustomComponent2) -> AccountComponent {
            AccountComponent::new(
                COMPONENT_2_PACKAGE.clone(),
                vec![],
                AccountComponentMetadata::mock("component2::interface"),
            )
            .expect("should be valid")
        }
    }

    let slot_content1 = Word::from([1, 2, 3, 4u32]);
    let slot = StorageSlot::with_value(TEST_SLOT_NAME.clone(), slot_content1);

    let account = AccountBuilder::new([42; 32])
        .with_components(Auth::IncrNonce)
        .with_component(CustomComponent1 { slot: slot.clone() })
        .with_component(CustomComponent2)
        .build()
        .context("failed to build account")?;

    let tx_script = format!(
        r#"
      use component1::interface as comp1_interface
      use component2::interface as comp2_interface

      @transaction_script
      pub proc main
          call.comp1_interface::get_slot_content
          push.{slot_content1}
          assert_eqw.err="failed to get slot content1"

          call.comp2_interface::set_slot_content

          call.comp2_interface::get_slot_content
          push.[5,6,7,8]
          assert_eqw.err="failed to get slot content2"
      end
    "#
    );

    let tx_script = CodeBuilder::default()
        .with_dynamically_linked_package(COMPONENT_1_PACKAGE.clone())?
        .with_dynamically_linked_package(COMPONENT_2_PACKAGE.clone())?
        .compile_tx_script(tx_script)?;

    TestTransactionBuilder::new(account)
        .tx_script(tx_script)
        .build()?
        .execute()
        .await?;

    Ok(())
}