use.kernel::account
use.kernel::asset_vault
use.kernel::constants
use.kernel::memory
use.kernel::note
use.kernel::tx

use.std::crypto::hashes::rpo

# ERRORS
# =================================================================================================

# Account nonce did not increase after a state changing transaction
const.ERR_ACCOUNT_NONCE_DID_NOT_INCREASE_AFTER_STATE_CHANGE=0x00020028

# Total number of assets in the account and all involved notes must stay the same
const.ERR_EPILOGUE_TOTAL_NUMBER_OF_ASSETS_MUST_STAY_THE_SAME=0x00020029

# OUTPUT NOTES PROCEDURES
# =================================================================================================

#! Copies output note data to the advice map. If no notes were created by a transaction, nothing
#! is copied to the advice map.
#!
#! Stack: [OUTPUT_NOTES_COMMITMENT, ...]
#! Output: [OUTPUT_NOTES_COMMITMENT, ...]
proc.copy_output_notes_to_advice_map
    # get the number of notes created by the transaction
    exec.memory::get_num_output_notes
    # => [num_notes, OUTPUT_NOTES_COMMITMENT, ...]

    # if there are output notes, add them to the advice map
    dup eq.0
    if.true
        # drop num_notes
        drop
    else
        # compute the end boundary of the output notes section
        exec.memory::get_output_note_ptr movdn.4
        # => [OUTPUT_NOTES_COMMITMENT, output_notes_end_ptr, ...]

        # compute the start boundary of the output notes section
        exec.memory::get_output_note_data_offset movdn.4
        # => [OUTPUT_NOTES_COMMITMENT, output_note_ptr, output_notes_end_ptr, ...]

        # insert created data into the advice map
        adv.insert_mem push.15511 drop          # TODO: remove line, see miden-vm/#1122
        # => [OUTPUT_NOTES_COMMITMENT, output_note_ptr, output_notes_end_ptr, ...]

        # drop output note pointers
        movup.4 drop movup.4 drop
    end
    # => [OUTPUT_NOTES_COMMITMENT, ...]
end

# BUILD OUTPUT VAULT
# =================================================================================================

#! Builds the output vault which is combination of the assets in the account vault at the end of
#! the transaction and all the assets in the output notes.
#!
#! The output vault is built as follows:
#! - we first copy the account vault root to the output vault root.
#! - we then loop over the output notes and insert their assets into the output vault.
#!
#! Stack: []
#! Output: []
proc.build_output_vault
    # copy final account vault root to output account vault root
    exec.memory::get_acct_vault_root exec.memory::set_output_vault_root dropw
    # => []

    # get the number of output notes from memory
    exec.memory::get_num_output_notes
    # => [num_output_notes]

    # calculate the address at which we should stop looping
    exec.memory::get_output_note_ptr
    # => [output_notes_end_ptr]

    # compute pointer for the first output note
    push.0 exec.memory::get_output_note_ptr
    # => [output_note_ptr, output_notes_end_ptr]

    # check if the number of output notes is greater then 0. Conditional for the while loop.
    dup.1 dup.1 neq
    # => [should_loop, output_note_ptr, output_notes_end_ptr]

    # loop over output notes and add assets to output vault
    while.true
        # get the number of assets for the output note from memory
        dup exec.memory::get_output_note_num_assets
        # => [num_assets, note_data_ptr, output_notes_end_ptr]

        # prepare stack for reading output note assets
        exec.memory::get_output_vault_root_ptr dup.2 exec.memory::get_output_note_asset_data_ptr dup
        # => [assets_start_ptr, assets_start_ptr, output_vault_root_ptr, num_assets, note_data_ptr,
        #     output_notes_end_ptr]

        # compute the end pointer for output note asset looping
        dup.3 add swap
        # => [assets_start_ptr, assets_end_ptr, output_vault_root_ptr, num_assets, note_data_ptr,
        #     output_notes_end_ptr]

        # assess if we should loop
        dup.1 dup.1 neq
        # => [should_loop, assets_start_ptr, assets_end_ptr, output_vault_root_ptr, num_assets,
        #     note_data_ptr, output_notes_end_ptr]

        # loop over output note assets and insert them into the output vault
        while.true
            # duplicate output_vault_root_ptr
            dup.2
            # => [output_vault_root_ptr, assets_start_ptr, assets_end_ptr, output_vault_root_ptr,
            #     num_assets, note_data_ptr, output_notes_end_ptr]

            # read the output note asset from memory
            padw dup.5 mem_loadw
            # => [ASSET, output_vault_root_ptr, assets_start_ptr, assets_end_ptr, output_vault_root_ptr,
            #     num_assets, note_data_ptr, output_notes_end_ptr]

            # insert output note asset into output vault
            exec.asset_vault::add_asset dropw
            # => [assets_start_ptr, assets_end_ptr, output_vault_root_ptr, num_assets, note_data_ptr,
            #     output_notes_end_ptr]

            # increment assets_start_ptr and asses if we should loop again
            add.1 dup.1 dup.1 neq
            # => [should_loop, assets_start_ptr, assets_end_ptr, output_vault_root_ptr, num_assets,
            #     note_data_ptr, output_notes_end_ptr]
        end

        # clean stack
        drop drop drop drop
        # => [note_data_ptr, output_note_end_ptr]

        # increment output note pointer and check if we should loop again
        exec.constants::get_note_mem_size add dup.1 dup.1 neq
        # => [should_loop, output_note_ptr, output_notes_end_ptr]
    end

    # clean stack
    drop drop
    # => []
end

# ACCOUNT CODE UPDATE
# =================================================================================================

#! Updates the account code commitment if the account code has changed. `NEW_ACCT_CODE_COMMITMENT`
#! is set to the initial account code commitment in the prologue and as such this procedure will
#! not result in a change to the account code commitment if the `account::set_code` procedure has
#! not been invoked in this transaction.
#!
#! Stack: []
#! Output: []
proc.update_account_code
    # check if the account code commitment has been updated
    exec.memory::get_new_acct_code_commitment
    # => [NEW_ACCT_CODE_COMMITMENT]

    # set the account code commitment to the new account code commitment (may not have changed)
    exec.memory::set_acct_code_commitment dropw
    # => []
end


# ACCOUNT STORAGE UPDATE
# =================================================================================================

#! Updates the account storage commitment by hashing the storage slots and setting the storage
#! commitment to the newly computed storage hash.
#!
#! Stack: []
#! Output: []
proc.update_account_storage_commitment
    # get number of storage slots
    exec.memory::get_num_storage_slots
    # => [num_storage_slots]

    # check if there are storage slots
    dup push.0 neq
    # => [storage_slots_is_not_empty, num_storage_slots]

    # only loop and hash over storage if there are storage slots
    if.true
        # setup start and end ptr
        mul.2 exec.memory::get_acct_storage_slots_section_ptr dup movdn.2 add swap
        # => [start_ptr, end_ptr]

        # pad stack to read and hash from memory
        padw padw padw
        # => [PAD, PAD, PAD, start_ptr, end_ptr]

        # hash elements from memory
        exec.rpo::absorb_double_words_from_memory
        # => [PERM, PERM, PERM, start_ptr, end_ptr]

        # extract the digest
        exec.rpo::squeeze_digest
        # => [DIGEST, end_ptr, end_ptr]

        # clean stack
        movup.4 drop movup.4 drop
        # => [DIGEST]

        # set new account storage commitment
        exec.memory::set_acct_storage_commitment dropw
        # => []
    else
        # drop number of storage slots
        drop
        # => []
    end
end

# TRANSACTION EPILOGUE PROCEDURE
# =================================================================================================

#! Finalizes the transaction by performing the following steps:
#! - computes the final account hash
#! - if the account has changed, assert that the final account nonce is greater than the initial
#!   account nonce
#! - computes the output notes commitment
#! - asserts that the input and output vault roots are equal
#!
#! Stack: []
#! Output: [OUTPUT_NOTES_COMMITMENT, FINAL_ACCOUNT_HASH, tx_expiration_block_num]
#!
#! - OUTPUT_NOTES_COMMITMENT is the commitment of the output notes
#! - FINAL_ACCOUNT_HASH is the final account hash
export.finalize_transaction
    # make sure that the context was switched back to the native account
    exec.memory::assert_native_account

    # update account code
    exec.update_account_code
    # => []

    # update account storage commitment
    exec.update_account_storage_commitment
    # => []

    # get the initial account hash
    exec.memory::get_init_acct_hash
    # => [INIT_ACCT_HASH]

    # get the offset for the end of the account data section
    exec.memory::get_core_acct_data_end_ptr
    # => [acct_data_end_ptr, INIT_ACCT_HASH]

    # get the offset for the start of the account data section
    exec.memory::get_current_account_data_ptr
    # => [acct_data_ptr, acct_data_end_ptr, INIT_ACCT_HASH]

    # compute the final account hash
    exec.account::get_current_hash
    # => [FINAL_ACCOUNT_HASH, acct_data_ptr, acct_data_end_ptr, INIT_ACCT_HASH]

    # insert final account data into the advice map
    adv.insert_mem push.15619 drop              # TODO: remove line, see miden-vm/#1122
    # => [FINAL_ACCOUNT_HASH, acct_data_ptr, acct_data_end_ptr, INIT_ACCT_HASH]

    # drop account data section pointers
    movup.4 drop movup.4 drop
    # => [FINAL_ACCOUNT_HASH, INIT_ACCT_HASH]

    # check if the account has changed
    eqw not
    # => [has_acct_changed, FINAL_ACCOUNT_HASH, INIT_ACCT_HASH]

    # if the account has changed then apply nonce checks
    if.true
        # get initial nonce from memory
        exec.memory::get_init_nonce
        # => [init_nonce, FINAL_ACCOUNT_HASH, INIT_ACCT_HASH]

        # get current nonce from memory
        exec.memory::get_acct_nonce
        # => [current_nonce, init_nonce, FINAL_ACCOUNT_HASH, INIT_ACCT_HASH]

        push.1005 drop                              # TODO: remove line, see miden-vm/#1122

        # assert that initial nonce is less than current nonce
        lt assert.err=ERR_ACCOUNT_NONCE_DID_NOT_INCREASE_AFTER_STATE_CHANGE
        # => [FINAL_ACCOUNT_HASH, INIT_ACCT_HASH]
    end

    # drop initial account hash
    swapw dropw
    # => [FINAL_ACCOUNT_HASH]

    # build the output vault
    exec.build_output_vault
    # => [FINAL_ACCOUNT_HASH]

    # compute output notes commitment
    exec.note::compute_output_notes_commitment
    # => [OUTPUT_NOTES_COMMITMENT, FINAL_ACCOUNT_HASH]

    # copy output note data to the advice map
    exec.copy_output_notes_to_advice_map
    # => [OUTPUT_NOTES_COMMITMENT, FINAL_ACCOUNT_HASH]

    # truncate stack
    swapdw dropw dropw
    # => [OUTPUT_NOTES_COMMITMENT, FINAL_ACCOUNT_HASH]

    # assert no net creation or destruction of assets over the transaction
    exec.memory::get_input_vault_root exec.memory::get_output_vault_root assert_eqw.err=ERR_EPILOGUE_TOTAL_NUMBER_OF_ASSETS_MUST_STAY_THE_SAME
    # => [OUTPUT_NOTES_COMMITMENT, FINAL_ACCOUNT_HASH]

    swapdw
    exec.memory::get_expiration_block_num
    swap drop swapdw
    # => [OUTPUT_NOTES_COMMITMENT, FINAL_ACCOUNT_HASH, tx_expiration_block_num]
end