api.masm

use.std::collections::smt

use.miden::kernels::tx::account
use.miden::kernels::tx::asset_vault
use.miden::kernels::tx::constants
use.miden::kernels::tx::faucet
use.miden::kernels::tx::memory
use.miden::kernels::tx::note
use.miden::kernels::tx::tx

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

# For faucets the slot FAUCET_STORAGE_DATA_SLOT is reserved and can not be used with set_account_item
const.ERR_FAUCET_RESERVED_DATA_SLOT=0x00020000

# Procedure can only be called from faucet accounts
const.ERR_ACCT_MUST_BE_A_FAUCET=0x00020001

# Getting a map item on a non-map slot
const.ERR_READING_MAP_VALUE_FROM_NON_MAP_SLOT=0x00020049

# EVENTS
# =================================================================================================

# Event emitted before an asset is added to the account vault.
const.ACCOUNT_VAULT_BEFORE_ADD_ASSET_EVENT=131072
# Event emitted after an asset is added to the account vault.
const.ACCOUNT_VAULT_AFTER_ADD_ASSET_EVENT=131073

# Event emitted before an asset is removed from the account vault.
const.ACCOUNT_VAULT_BEFORE_REMOVE_ASSET_EVENT=131074
# Event emitted after an asset is removed from the account vault.
const.ACCOUNT_VAULT_AFTER_REMOVE_ASSET_EVENT=131075

# AUTHENTICATION
# =================================================================================================

#! Authenticates that the invocation of a kernel procedure originates from the account context.
#!
#! Panics:
#!   - if the invocation of the kernel procedure does not originate from the account context.
#!
#! Stack: [...]
#! Output: [...]
proc.authenticate_account_origin
    # get the hash of the caller
    padw caller
    # => [CALLER, ...]

    # assert that the caller is from the user context
    exec.account::authenticate_procedure
    # => [CALLER, ...]

    # drop the caller
    dropw
    # => [...]
end

# KERNEL PROCEDURES
# =================================================================================================

#! Returns the account id.
#!
#! Stack: [0]
#! Output: [acct_id]
#!
#! - acct_id is the account id.
export.get_account_id
    # get the account id
    exec.account::get_id
    # => [acct_id, 0]

    swap drop
    # => [acct_id]
end

#! Returns the account nonce.
#!
#! Stack: [0]
#! Output: [nonce]
#!
#! - nonce is the account nonce.
export.get_account_nonce
    # get the account nonce
    exec.account::get_nonce
    # => [0, nonce]

    swap drop
    # => [nonce]
end

#! Returns the initial account hash.
#!
#! Stack: [0, 0, 0, 0]
#! Output: [H]
#!
#! - H is the initial account hash.
export.get_initial_account_hash
    # get the initial account hash
    exec.account::get_initial_hash
    # => [H, 0, 0, 0, 0]

    swapw dropw
    # => [H]
end

#! Computes and returns the account hash from account data stored in memory.
#!
#! Stack: [0, 0, 0, 0]
#! Output: [ACCT_HASH]
#!
#! - ACCT_HASH is the hash of the account data.
export.get_current_account_hash
    # get the current account hash
    exec.account::get_current_hash
    # => [ACCT_HASH, 0, 0, 0, 0]

    swapw dropw
    # => [ACCT_HASH]
end

#! Increments the account nonce by the provided value.
#!
#! Stack: [value]
#! Output: [0]
#!
#! - value is the value to increment the nonce by. value can be at most 2^32 - 1 otherwise this
#!   procedure panics.
export.incr_account_nonce
    # authenticate that the procedure invocation originates from the account context
    exec.authenticate_account_origin
    # => [value]

    # arrange stack
    push.0 swap
    # => [value, 0]

    # increment the account nonce
    exec.account::incr_nonce
    # => [0]
end

#! Gets an item from the account storage. Panics if the index is out of bounds.
#!
#! Stack: [index, 0, 0, 0]
#! Output: [VALUE]
#!
#! - index is the index of the item to get.
#! - VALUE is the value of the item.
export.get_account_item
    # fetch the account storage item
    exec.account::get_item
    # => [VALUE, 0, 0, 0]

    # organize the stack for return
    movup.4 drop movup.4 drop movup.4 drop
    # => [VALUE]
end

#! Sets an item in the account storage. Panics if the index is out of bounds.
#!
#! Stack: [index, V', 0, 0, 0]
#! Output: [R', V]
#!
#! - index is the index of the item to set.
#! - V' is the value to set.
#! - V is the previous value of the item.
#! - R' is the new storage root.
export.set_account_item
    # if the transaction is being executed against a faucet account then assert
    # index != FAUCET_STORAGE_DATA_SLOT (reserved slot)
    dup exec.account::get_faucet_storage_data_slot eq
    exec.account::get_id exec.account::is_faucet
    and assertz.err=ERR_FAUCET_RESERVED_DATA_SLOT
    # => [index, V', 0, 0, 0]

    # authenticate that the procedure invocation originates from the account context
    exec.authenticate_account_origin
    # => [index, V', 0, 0, 0]

    # set the account storage item
    exec.account::set_item
    # => [R', V, 0, 0, 0]

    # organize the stack for return
    movup.8 drop movup.8 drop movup.8 drop
    # => [R', V]
end

#! Returns VALUE located under specified KEY in map in specified account storage slot.
#! Panics if
#! - the index is out of bounds (>255).
#! - the requested storage slot type is not map
#!
#! Stack: [index, KEY, ...]
#! Output: [VALUE, 0]
#!
#! - index is the index of the item to get.
#! - VALUE is the value of the item.
export.get_account_map_item
    # check if storage type is map
    dup exec.account::get_storage_slot_type_info drop
    # => [slot_type, index, KEY, ...]

    # fails if slot_type is not 1 = map
    exec.constants::get_storage_slot_type_map eq assert.err=ERR_READING_MAP_VALUE_FROM_NON_MAP_SLOT
    # => [index, KEY, ...]

    # fetch the account storage item, which is ROOT of the map
    exec.account::get_item swapw
    # => [KEY, ROOT ...]

    # fetch the VALUE located under KEY in the tree
    exec.smt::get
    # => [VALUE, ROOT, ...]

    # prepare the stack for return
    swapw dropw
    # => [VALUE, 0]
end

#! Inserts specified NEW_VALUE under specified KEY in map in specified account storage slot.
#! Panics if
#! - the index is out of bounds (>255).
#! - the requested storage slot type is not map
#! - the procedure is called from a non-account context
#!
#! Stack: [index, KEY, NEW_VALUE, ...]
#! Output: [OLD_MAP_ROOT, OLD_MAP_VALUE, 0]
#!
#! - index is the index of the item to get.
#! - NEW_VALUE is the value of the new map item for the respective KEY.
#! - OLD_VALUE is the value of the old map item for the respective KEY.
#! - KEY is the key of the new item.
#! - OLD_MAP_ROOT is the root of the old map before insertion
#! - NEW_MAP_ROOT is the root of the new map after insertion.
export.set_account_map_item.1
    # authenticate that the procedure invocation originates from the account context
    exec.authenticate_account_origin
    # => [index, KEY, NEW_VALUE, ...]

    # store index for later
    dup loc_store.0
    # => [index, KEY, NEW_VALUE, ...]

    # fetch the account storage item, which is ROOT of the map
    exec.account::get_item movdnw.2
    # => [KEY, NEW_VALUE, OLD_MAP_ROOT, ...]

    # set the new map item
    loc_load.0 exec.account::set_map_item
    # => [OLD_MAP_ROOT, OLD_VALUE, ...]

    # organize the stack for return (16 elements)
    movupw.2 dropw
    # => [OLD_MAP_ROOT, OLD_MAP_VALUE, 0, ...]
end

#! Sets the code of the account the transaction is being executed against. This procedure can only
#! executed on regular accounts with updatable code. Otherwise, this procedure fails.
#!
#! Stack: [CODE_ROOT]
#! Output: [0, 0, 0, 0]
#!
#! - CODE_ROOT is the hash of the code to set.
export.set_account_code
    # authenticate that the procedure invocation originates from the account context
    exec.authenticate_account_origin
    # => [CODE_ROOT]

    # arrange stack
    padw swapw
    # => [CODE_ROOT, 0, 0, 0, 0]

    # set the account code
    exec.account::set_code
    # => [0, 0, 0, 0]
end

#! Returns the balance of a fungible asset associated with a faucet_id.
#! Panics if the asset is not a fungible asset.
#!
#! Stack: [faucet_id]
#! Output: [balance]
#!
#! - faucet_id is the faucet id of the fungible asset of interest.
#! - balance is the vault balance of the fungible asset.
export.account_vault_get_balance
    # get the vault root
    exec.memory::get_acct_vault_root_ptr swap
    # => [faucet_id, acct_vault_root_ptr]

    # get the asset balance
    exec.asset_vault::get_balance
    # => [balance]
end

#! Returns a boolean indicating whether the non-fungible asset is present in the vault.
#! Panics if the ASSET is a fungible asset.
#!
#! Stack: [ASSET]
#! Output: [has_asset, 0, 0, 0]
#!
#! - ASSET is the non-fungible asset of interest
#! - has_asset is a boolean indicating whether the account vault has the asset of interest
export.account_vault_has_non_fungible_asset
    # arrange stack and get the vault root
    push.0 movdn.4 push.0 movdn.4 push.0 movdn.4 exec.memory::get_acct_vault_root_ptr movdn.4
    # => [ASSET, 0, 0, 0]

    # check if the account vault has the non-fungible asset
    exec.asset_vault::has_non_fungible_asset
    # => [has_asset, 0, 0, 0]
end

#! Add the specified asset to the vault.
#!
#! Panics:
#! - If the asset is not valid.
#! - If the total value of two fungible assets is greater than or equal to 2^63.
#! - If the vault already contains the same non-fungible asset.
#!
#! Stack: [ASSET]
#! Output: [ASSET']
#!
#! - ASSET is the asset to add to the vault.
#! - ASSET' final asset in the account vault defined as follows:
#!   - If ASSET is a non-fungible asset, then ASSET' is the same as ASSET.
#!   - If ASSET is a fungible asset, then ASSET' is the total fungible asset in the account vault
#!     after ASSET was added to it.
export.account_vault_add_asset
    # authenticate that the procedure invocation originates from the account context
    exec.authenticate_account_origin
    # => [ASSET]

    push.0 drop                                             # TODO: remove line, see miden-vm/#1122
    emit.ACCOUNT_VAULT_BEFORE_ADD_ASSET_EVENT
    # => [ASSET]

    # duplicate the ASSET to be able to emit an event after an asset is being added
    dupw
    # => [ASSET, ASSET]

    # fetch the vault root
    exec.memory::get_acct_vault_root_ptr movdn.4
    # => [ASSET, acct_vault_root_ptr, ASSET]

    # add the asset to the account vault
    exec.asset_vault::add_asset
    # => [ASSET', ASSET]

    # emit event to signal that an asset is being added to the account vault
    swapw emit.ACCOUNT_VAULT_AFTER_ADD_ASSET_EVENT dropw
    # => [ASSET']
end

#! Remove the specified asset from the vault.
#!
#! Panics:
#! - The fungible asset is not found in the vault.
#! - The amount of the fungible asset in the vault is less than the amount to be removed.
#! - The non-fungible asset is not found in the vault.
#!
#! Stack: [ASSET]
#! Output: [ASSET]
#!
#! - ASSET is the asset to remove from the vault.
export.account_vault_remove_asset
    # authenticate that the procedure invocation originates from the account context
    exec.authenticate_account_origin
    # => [ASSET]

    push.0 drop                                             # TODO: remove line, see miden-vm/#1122
    emit.ACCOUNT_VAULT_BEFORE_REMOVE_ASSET_EVENT
    # => [ASSET]

    # fetch the vault root
    exec.memory::get_acct_vault_root_ptr movdn.4
    # => [ASSET, acct_vault_root_ptr]

    # remove the asset from the account vault
    exec.asset_vault::remove_asset
    # => [ASSET]

    push.0 drop                                             # TODO: remove line, see miden-vm/#1122
    # emit event to signal that an asset is being removed from the account vault
    emit.ACCOUNT_VAULT_AFTER_REMOVE_ASSET_EVENT
    # => [ASSET]
end

#! Returns the number of assets and the assets hash of the note currently being processed. Panics
#! if a note is not being processed.
#!
#! Inputs: [0, 0, 0, 0, 0]
#! Outputs: [ASSETS_HASH, num_assets]
#!
#! - num_assets is the number of assets in the note currently being processed.
#! - ASSETS_HASH is the assets hash of the note currently being processed.
export.get_note_assets_info
    # get the assets info
    exec.note::get_assets_info
    # => [ASSETS_HASH, num_assets, 0, 0, 0, 0, 0]

    # organize the stack for return
    movup.5 drop movup.5 drop movup.5 drop movup.5 drop movup.5 drop
    # => [ASSETS_HASH, num_assets]
end

#! Returns the current note's inputs hash.
#!
#! Inputs: [EMPTY_WORD]
#! Outputs: [NOTE_INPUTS_HASH]
#!
#! Where:
#! - NOTE_INPUTS_HASH, is the current note's inputs hash.
export.get_note_inputs_hash
    exec.note::get_note_inputs_hash
    # => [NOTE_INPUTS_HASH, EMPTY_WORD]

    swapw dropw
    # => [NOTE_INPUTS_HASH]
end

#! Returns the sender of the note currently being processed. Panics if a note is not being
#! processed.
#!
#! Inputs: [0]
#! Outputs: [sender]
#!
#! Where:
#! - sender is the sender of the note currently being processed.
export.get_note_sender
    exec.note::get_sender swap drop
    # => [sender]
end

#! Returns the block number of the last known block at the time of transaction execution.
#!
#! Inputs: [0]
#! Outputs: [num]
#!
#! num is the last known block number.
export.get_block_number
    # get the block number
    exec.tx::get_block_number
    # => [num, 0]

    # organize the stack for return
    swap drop
    # => [num]
end

#! Returns the block hash of the reference block.
#!
#! Stack: [EMPTY_WORD]
#! Output: [BLOCK_HASH]
#!
#! Where:
#! - BLOCK_HASH, reference block for the transaction execution.
export.get_block_hash
    dropw exec.tx::get_block_hash
    # => [BLOCK_HASH]
end

#! Returns the input notes commitment.
#!
#! This is computed as a sequential hash of `(NULLIFIER, EMPTY_WORD_OR_NOTE_HASH)` over all input notes. The
#! data `EMPTY_WORD_OR_NOTE_HASH` functions as a flag, if the value is set to zero, then the notes are
#! authenticated by the transaction kernel. If the value is non-zero, then note authentication will
#! be delayed to the batch/block kernel. The delayed authentication allows a transaction to consume a
#! public note that is not yet included to a block.
#!
#! Inputs: [0, 0, 0, 0]
#! Outputs: [INPUT_NOTES_COMMITMENT]
#!
#! Where:
#! - INPUT_NOTES_COMMITMENT is the input notes commitment hash.
export.get_input_notes_commitment
    exec.tx::get_input_notes_commitment
    # => [COM, 0, 0, 0, 0]

    swapw dropw
    # => [COM]
end

#! Returns the output notes hash. This is computed as a sequential hash of (note_id, note_metadata)
#! tuples over all output notes.
#!
#! Inputs: [0, 0, 0, 0]
#! Outputs: [COM]
#!
#! COM is the output notes hash.
export.get_output_notes_hash
    # get the output notes hash
    exec.tx::get_output_notes_hash
    # => [COM, 0, 0, 0, 0]

    # organize the stack for return
    swapw dropw
    # => [COM]
end

#! Creates a new note and returns the index of the note.
#!
#! Inputs: [tag, aux, note_type, RECIPIENT]
#! Outputs: [note_idx, 0, 0, 0, 0, 0]
#!
#! tag is the tag to be included in the note.
#! aux is the auxiliary metadata to be included in the note.
#! note_type is the note storage type
#! RECIPIENT is the recipient of the note.
#! note_idx is the index of the crated note.
export.create_note
    # authenticate that the procedure invocation originates from the account context
    exec.authenticate_account_origin
    # => [tag, aux, note_type, RECIPIENT]

    exec.tx::create_note
    # => [note_idx]

    # prepare stack for return. Note: when create_note is called, the stack looks
    # like [tag, aux, note_type, RECIPIENT, x, X, X], with 16 elements and X might be important data
    # for the user. Without padding the kernel returns [note_idx, x, X, X, 0, 0, EMPTY_WORD] adding 0's.
    # To keep the data in position we move 0's to the left between note_idx and the potentially
    # important first element x.
    movupw.3 movup.15 movup.15 movup.6
    # => [note_idx, 0, 0, 0, 0, 0, 0]
end

#! Adds the ASSET to the note specified by the index.
#!
#! Inputs: [note_idx, ASSET]
#! Outputs: [note_idx, 0, 0, 0, 0]
#!
#! note_idx is the index of the the note to which the asset is added.
#! ASSET can be a fungible or non-fungible asset.
export.add_asset_to_note
    # authenticate that the procedure invocation originates from the account context
    exec.authenticate_account_origin
    # => [note_idx, ASSET]

    exec.tx::add_asset_to_note
    # => [note_idx]

    # prepare stack for return. Note: when add_asset_to_note is called, the stack looks
    # like [idx, ASSET, X, X, x, x, x], with 16 elements and X might be important data
    # for the user. Without padding the kernel returns [idx, X, X, x, x, x, EMPTY_WORD] adding 0's.
    # To keep the data in position we insert 0's between idx and the potentially important
    # first element x.
    movupw.3 movup.4
    # => [note_idx, 0, 0, 0, 0]
end

#! Returns a commitment to the account vault the transaction is being executed against.
#!
#! Stack: [0, 0, 0, 0]
#! Outputs: [COM]
#!
#! - COM is the commitment to the account vault.
export.get_account_vault_commitment
    # fetch the account vault root
    exec.memory::get_acct_vault_root
    # => [COM, 0, 0, 0, 0]

    # organize the stack for return
    swapw dropw
    # => [COM]
end

#! Mint an asset from the faucet the transaction is being executed against.
#!
#! Panics:
#! - If the transaction is not being executed against a faucet.
#! - If the asset being minted is not associated with the faucet the transaction is being executed
#!   against.
#! - If the asset is not well formed.
#! - For fungible faucets if the total issuance after minting is greater than the maximum amount
#!   allowed.
#! - For non-fungible faucets if the non-fungible asset being minted already exists.
#!
#! Stack: [ASSET]
#! Outputs: [ASSET]
#!
#! - ASSET is the asset that was minted.
export.mint_asset
    # authenticate that the procedure invocation originates from the account context
    exec.authenticate_account_origin
    # => [ASSET]

    # mint the asset
    exec.faucet::mint
    # => [ASSET]
end

#! Burn an asset from the faucet the transaction is being executed against.
#!
#! Panics:
#! - If the transaction is not being executed against a faucet.
#! - If the asset being burned is not associated with the faucet the transaction is being executed
#!   against.
#! - If the asset is not well formed.
#! - For fungible faucets if the amount being burned is greater than the total input to the
#!   transaction.
#! - For non-fungible faucets if the non-fungible asset being burned does not exist or was not
#!   provided as input to the transaction via a note or the accounts vault.
#!
#! Stack: [ASSET]
#! Outputs: [ASSET]
#!
#! - ASSET is the asset that was burned.
export.burn_asset
    # authenticate that the procedure invocation originates from the account context
    exec.authenticate_account_origin
    # => [ASSET]

    # burn the asset
    exec.faucet::burn
    # => [ASSET]
end

#! Returns the total issuance of the fungible faucet the transaction is being executed against.
#!
#! Panics:
#! - If the transaction is not being executed against a fungible faucet.
#!
#! Stack: [0]
#! Outputs: [total_issuance]
#!
#! - total_issuance is the total issuance of the fungible faucet the transaction is being executed
#!   against.
export.get_fungible_faucet_total_issuance
    # assert that we are executing a transaction against a fungible faucet (access checks)
    exec.account::get_id exec.account::is_fungible_faucet assert.err=ERR_ACCT_MUST_BE_A_FAUCET
    # => [0]

    # get the total issuance
    exec.faucet::get_total_issuance
    # => [total_issuance]

    # drop the padding
    swap drop
    # => []
end

#! Returns the serial number of the note currently being processed.
#! Panics if no note is not being processed.
#!
#! Inputs: []
#! Outputs: [SERIAL_NUMBER]
#!
#! - SERIAL_NUMBER is the serial number of the note currently being processed.
export.get_note_serial_number
    exec.note::get_serial_number
    # => [SERIAL_NUMBER]

    # drop the padding
    swapw dropw
    # => [SERIAL_NUMBER]
end