miden-testing 0.17.0-rc.5

use miden_protocol::Word;
use miden_protocol::account::auth::AuthScheme;
use miden_protocol::account::component::AccountComponentMetadata;
use miden_protocol::account::{Account, AccountComponent, AccountType};
use miden_protocol::asset::{Asset, FungibleAsset};
use miden_protocol::errors::tx_kernel::ERR_EPILOGUE_TOTAL_NUMBER_OF_ASSETS_MUST_STAY_THE_SAME;
use miden_protocol::note::{Note, NoteRecipient, NoteStorage, NoteType};
use miden_protocol::transaction::RawOutputNote;
use miden_standards::code_builder::CodeBuilder;
use miden_standards::note::TxFeeNote;
use miden_testing::{AccountState, Auth, MockChain, assert_transaction_executor_error};

/// The TX_FEE script leaves its assets in the note: a wallet that consumes one without moving the
/// assets out fails the kernel's asset conservation check. This pins the division of labor:
/// collecting the fee is the consuming account's job.
#[tokio::test]
async fn tx_fee_note_leaves_assets_in_the_note() -> anyhow::Result<()> {
    let fee_asset: Asset = FungibleAsset::mock(123);

    let mut builder = MockChain::builder();

    // Create accounts; the consumer is completely unrelated to the note
    let sender_account = builder.create_new_wallet(Auth::BasicAuth {
        auth_scheme: AuthScheme::Falcon512Poseidon2,
    })?;
    let consumer_account = builder.add_existing_wallet(Auth::BasicAuth {
        auth_scheme: AuthScheme::Falcon512Poseidon2,
    })?;

    // Create the note
    let note = builder.add_tx_fee_note(sender_account.id(), &[fee_asset])?;

    assert_eq!(note.metadata().tag(), TxFeeNote::TAG);
    assert_eq!(note.metadata().note_type(), NoteType::Public);

    let mock_chain = builder.build()?;

    // a basic wallet has nothing that moves the assets out of the note, so they stay in it and
    // the transaction fails asset conservation
    let result = mock_chain
        .build_transaction(consumer_account.id())
        .authenticated_input_note(note.id())
        .build()?
        .execute()
        .await;

    assert_transaction_executor_error!(
        result,
        ERR_EPILOGUE_TOTAL_NUMBER_OF_ASSETS_MUST_STAY_THE_SAME
    );

    Ok(())
}

/// An account whose own code removes the assets from the note collects them; the TX_FEE note
/// imposes no target restriction, so an account unrelated to the note can do so.
#[tokio::test]
async fn tx_fee_note_assets_are_collected_by_account_code() -> anyhow::Result<()> {
    let fee_asset = FungibleAsset::mock(123);

    let mut builder = MockChain::builder();
    let sender_account = builder.create_new_wallet(Auth::BasicAuth {
        auth_scheme: AuthScheme::Falcon512Poseidon2,
    })?;

    // a component that moves a given asset from an input note into the account's vault
    let collector_component = AccountComponent::new(
        CodeBuilder::default().compile_component_code(
            "fee_collector",
            "
            use miden::protocol::asset
            use miden::protocol::input_note
            use miden::protocol::native_account

            #! Removes the given asset from the input note at `note_index` and adds it to this
            #! account's vault.
            #!
            #! Inputs:  [ASSET_ID, ASSET_VALUE, note_index, pad(7)]
            #! Outputs: [pad(16)]
            @account_procedure
            @locals(8)
            pub proc collect_note_asset
                # keep a copy of the asset so it can be added to the vault after removal
                dupw.1 dupw.1 locaddr.0 exec.asset::store
                # => [ASSET_ID, ASSET_VALUE, note_index, pad(7)]

                exec.input_note::remove_asset dropw
                # => [pad(16)]

                locaddr.0 exec.asset::load exec.native_account::add_asset dropw
                # => [pad(16)]
            end
            ",
        )?,
        Vec::new(),
        AccountComponentMetadata::mock("fee_collector"),
    )?;

    let mut consumer_account = builder.add_account_from_builder(
        Auth::BasicAuth {
            auth_scheme: AuthScheme::Falcon512Poseidon2,
        },
        Account::builder([9; 32])
            .account_type(AccountType::Public)
            .with_component(collector_component.clone()),
        AccountState::Exists,
    )?;

    let note = builder.add_tx_fee_note(sender_account.id(), &[fee_asset])?;
    let mock_chain = builder.build()?;

    // the transaction script collects the asset of input note 0 through the account's procedure
    let tx_script_src = format!(
        r#"
        use miden::core::sys

        @transaction_script
        pub proc main
            push.0.0.0.0 push.0.0.0
            push.0
            push.{asset_value}
            push.{asset_id}
            call.::fee_collector::collect_note_asset
            exec.sys::truncate_stack
        end
        "#,
        asset_value = fee_asset.to_value_word(),
        asset_id = fee_asset.to_id_word(),
    );
    let tx_script = CodeBuilder::with_mock_packages()
        .with_dynamically_linked_package(collector_component.component_code())?
        .compile_tx_script(tx_script_src)?;

    let executed = mock_chain
        .build_transaction(consumer_account.id())
        .authenticated_input_note(note.id())
        .tx_script(tx_script)
        .build()?
        .execute()
        .await?;

    consumer_account.apply_patch(executed.account_patch())?;
    assert_eq!(consumer_account.vault().get(fee_asset.id()), Some(fee_asset));

    Ok(())
}

/// Tests the TX_FEE `create_output_note` MASM constructor procedure.
/// This test verifies that calling `tx_fee::create_output_note` from a transaction script
/// creates an output note with the same recipient as the Rust `TxFeeNote` implementation would
/// create, and that the MASM note tag constant matches [`TxFeeNote::TAG`].
#[tokio::test]
async fn test_tx_fee_create_output_note_constructor() -> anyhow::Result<()> {
    let mut builder = MockChain::builder();

    let sender_account = builder.add_existing_wallet_with_assets(
        Auth::BasicAuth {
            auth_scheme: AuthScheme::Falcon512Poseidon2,
        },
        [FungibleAsset::mock(100)],
    )?;

    let mock_chain = builder.build()?;

    // Create a serial number for the note
    let serial_num = Word::from([1u32, 2u32, 3u32, 4u32]);

    // Build the expected recipient using the Rust implementation
    let expected_recipient =
        NoteRecipient::new(serial_num, TxFeeNote::script(), NoteStorage::default());

    // Build a transaction script that uses tx_fee::create_output_note to create a note
    let tx_script_src = format!(
        r#"
        use miden::standards::notes::tx_fee

        @transaction_script
        pub proc main
            # Push inputs for tx_fee::create_output_note
            push.{serial_num}
            # => [SERIAL_NUM]

            exec.tx_fee::create_output_note
            # => [note_idx]

            # Add an asset to the created note
            push.{ASSET_VALUE}
            push.{ASSET_ID}
            call.::miden::standards::wallets::basic::move_asset_to_note

            # Clean up stack
            dropw dropw dropw dropw
        end
        "#,
        serial_num = serial_num,
        ASSET_ID = FungibleAsset::mock(50).to_id_word(),
        ASSET_VALUE = FungibleAsset::mock(50).to_value_word(),
    );

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

    // Build expected output note
    let expected_output_note: Note = TxFeeNote::builder()
        .sender(sender_account.id())
        .asset(FungibleAsset::mock(50))
        .serial_number(serial_num)
        .build()?
        .into();

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

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

    // Verify that one note was created
    assert_eq!(executed_transaction.output_notes().num_notes(), 1);

    // Get the created note's recipient and verify it matches
    let output_note = executed_transaction.output_notes().get_note(0);
    let created_recipient = output_note.recipient().expect("output note should have recipient");

    // Verify the recipient matches what we expected
    assert_eq!(
        created_recipient.digest(),
        expected_recipient.digest(),
        "The recipient created by tx_fee::create_output_note should match the Rust TxFeeNote \
         implementation"
    );

    // Verify the MASM TX_FEE_NOTE_TAG constant matches the Rust TxFeeNote::TAG and the note
    // is public
    assert_eq!(output_note.metadata().tag(), TxFeeNote::TAG);
    assert_eq!(output_note.metadata().note_type(), NoteType::Public);

    Ok(())
}