miden-protocol 0.16.0-alpha.4

use miden::protocol_utils::mem
use {INPUT_NOTE_GET_METADATA_OFFSET, INPUT_NOTE_GET_RECIPIENT_OFFSET, INPUT_NOTE_GET_SCRIPT_ROOT_OFFSET, INPUT_NOTE_GET_SERIAL_NUMBER_OFFSET, INPUT_NOTE_GET_STORAGE_INFO_OFFSET}
    from miden::protocol::kernel_proc_offsets
use miden::protocol::note
use miden::protocol::input_note
use {NOTE_TYPE_PRIVATE, NOTE_TYPE_PUBLIC} from miden::protocol::note
use {AccountId, Bool, MemoryAddress, NoteMetadata, NoteRecipient, NoteScriptRoot}
    from miden::protocol::types

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

const ERR_NOTE_DATA_DOES_NOT_MATCH_COMMITMENT = "note data does not match the commitment"

const ERR_NOTE_INVALID_NUMBER_OF_STORAGE_ITEMS =
    "the specified number of note storage items does not match the actual number"

# ACTIVE NOTE PROCEDURES
# =================================================================================================
#
# The "active note" refers to the note that is currently being processed by the transaction kernel.

#! Writes the initial assets of the active note into memory starting at the specified address.
#!
#! The written assets are the ones the note was created with; they are unaffected by asset removal
#! during transaction execution.
#!
#! Attention: memory starting from the `dest_ptr` should have enough space to store all the assets
#! the note was created with. Make sure that at least `8 * num_assets` memory elements are
#! available, or if the number of assets is not yet known, at least `8 * MAX_ASSETS_PER_NOTE`.
#!
#! The memory layout after the execution of this procedure will look like so:
#! [ASSET_ID_0, ASSET_VALUE_0, ASSET_ID_1, ASSET_VALUE_1, ..., ASSET_ID_N, ASSET_VALUE_N], where
#! each asset occupies two words. For more detailed information about the layout of each asset see
#! the description of the `Asset` Rust type.
#!
#! Inputs:  [dest_ptr]
#! Outputs: [num_assets]
#!
#! Where:
#! - dest_ptr is the memory address to write the assets.
#! - num_assets is the number of assets the active note was created with.
#!
#! Panics if:
#! - no note is currently active.
#!
#! Invocation: exec
pub proc get_initial_assets(dest_ptr: MemoryAddress) -> u8
    # push a placeholder note_index and the active note flag
    push.0.1
    # => [is_active_note = 1, note_index = 0, dest_ptr]

    exec.input_note::get_initial_assets_raw
    # => [num_assets]
end

#! Returns the initial assets information of the active note.
#!
#! The returned commitment and number of assets are the ones the note was created with; they do
#! not change when assets are removed from the note during transaction execution.
#!
#! Inputs:  []
#! Outputs: [ASSETS_COMMITMENT, num_assets]
#!
#! Where:
#! - num_assets is the number of assets the active note was created with.
#! - ASSETS_COMMITMENT is a sequential hash of the assets the active note was created with.
#!
#! Panics if:
#! - no note is currently active.
#!
#! Invocation: exec
pub proc get_initial_assets_info
    # push a placeholder note_index and the active note flag
    push.0.1
    # => [is_active_note = 1, note_index = 0]

    exec.input_note::get_initial_assets_info_raw
    # => [ASSETS_COMMITMENT, num_assets]
end

#! Returns the number of assets the active note was created with.
#!
#! The returned number does not change when assets are removed from the note during transaction
#! execution.
#!
#! Inputs:  []
#! Outputs: [num_assets]
#!
#! Where:
#! - num_assets is the number of assets the active note was created with.
#!
#! Panics if:
#! - no note is currently active.
#!
#! Invocation: exec
pub proc get_initial_num_assets
    exec.get_initial_assets_info
    # => [ASSETS_COMMITMENT, num_assets]

    dropw
    # => [num_assets]
end

#! Returns the asset at the specified index in the active note.
#!
#! The asset is returned as it currently is: if it was removed from the note, the returned
#! ASSET_ID and ASSET_VALUE are both the EMPTY_WORD.
#!
#! Inputs:  [asset_index]
#! Outputs: [ASSET_ID, ASSET_VALUE]
#!
#! Where:
#! - asset_index is the index of the asset to return.
#! - ASSET_ID is the asset ID of the asset at the specified index.
#! - ASSET_VALUE is the value of the asset at the specified index.
#!
#! Panics if:
#! - no note is currently active.
#! - the asset index is greater or equal to the number of assets in the note.
#!
#! Invocation: exec
pub proc get_asset
    # push a placeholder note_index below asset_index and the active note flag on top
    push.0 swap push.1
    # => [is_active_note = 1, asset_index, note_index = 0]

    exec.input_note::get_asset_raw
    # => [ASSET_ID, ASSET_VALUE]
end

#! Removes an asset from the active note.
#!
#! For instance, if the note holds a fungible asset with amount 100 under ASSET_ID, and
#! ASSET_VALUE has amount 30, then the value remaining in the note has amount 70.
#!
#! If the entire asset value is removed, the asset's slot in the note is cleared, i.e. the asset
#! at the corresponding index becomes (EMPTY_WORD, EMPTY_WORD). This is always the case for
#! non-composable assets.
#!
#! Inputs:  [ASSET_ID, ASSET_VALUE]
#! Outputs: [FINAL_ASSET_VALUE]
#!
#! Where:
#! - ASSET_ID is the asset ID of the asset to remove from the note.
#! - ASSET_VALUE is the value of the asset to remove from the note.
#! - FINAL_ASSET_VALUE is the value of the asset remaining in the note after removal, which is
#!   the EMPTY_WORD if the entire asset was removed.
#!
#! Panics if:
#! - no note is currently active.
#! - the asset's composition is Custom (not yet supported).
#! - the asset is not present in the note.
#! - the asset is not composable and is not present in the note with the exact value.
#! - the amount of the fungible asset in the note is less than the amount to be removed.
#!
#! Invocation: exec
pub proc remove_asset
    # push a placeholder note_index below the asset and the active note flag on top
    push.0 movdn.8 push.1
    # => [is_active_note = 1, ASSET_ID, ASSET_VALUE, note_index = 0]

    exec.input_note::remove_asset_raw
    # => [FINAL_ASSET_VALUE]
end

#! Removes all remaining assets from the active note and writes them into memory starting at the
#! specified address.
#!
#! The written assets are the ones that remained in the note at the time of the call, i.e.
#! excluding assets that were already removed. After this procedure returns, the note holds no
#! more assets; the written memory is a snapshot of the removed assets. The note's initial assets
#! info (see get_initial_assets_info) is not affected.
#!
#! Attention: memory starting from the `dest_ptr` should have enough space to store all the assets
#! in the note. Make sure that at least `8 * num_assets` memory elements are available, or if the
#! number of assets is not yet known, at least `8 * MAX_ASSETS_PER_NOTE`.
#!
#! The memory layout after the execution of this procedure will look like so:
#! [ASSET_0, ASSET_1, ..., ASSET_N], where each asset occupies two words. For more detailed
#! information about the layout of each asset see the description of the `Asset` Rust type.
#!
#! Inputs:  [dest_ptr]
#! Outputs: [num_assets]
#!
#! Where:
#! - dest_ptr is the memory address to write the assets.
#! - num_assets is the number of assets removed by this procedure.
#!
#! Panics if:
#! - no note is currently active.
#!
#! Invocation: exec
pub proc remove_all_assets
    # push a placeholder note_index and the active note flag
    push.0.1
    # => [is_active_note = 1, note_index = 0, dest_ptr]

    exec.input_note::remove_all_assets_raw
    # => [num_assets]
end

#! Returns the recipient of the active note.
#!
#! Inputs:  []
#! Outputs: [RECIPIENT]
#!
#! Where:
#! - RECIPIENT is the commitment to the active note's script, storage, the serial number.
#!
#! Panics if:
#! - no note is currently active.
#!
#! Invocation: exec
pub proc get_recipient() -> NoteRecipient
    # pad the stack
    padw padw padw push.0.0
    # => [pad(14)]

    # push the flag indicating that we want to request recipient from the active note
    push.1
    # => [is_active_note = 1, pad(14)]

    push.INPUT_NOTE_GET_RECIPIENT_OFFSET
    # => [offset, is_active_note = 1, pad(14)]

    syscall.exec_kernel_proc
    # => [RECIPIENT, pad(12)]

    # clean the stack
    swapdw dropw dropw swapw dropw
    # => [RECIPIENT]
end

#! Writes the active note's storage to memory starting at the specified address.
#!
#! Inputs:
#!   Stack: [dest_ptr]
#!   Advice Map: { NOTE_STORAGE_COMMITMENT: [STORAGE] }
#! Outputs:
#!   Stack: [num_storage_items]
#!
#! Where:
#! - dest_ptr is the memory address to write the note storage.
#! - NOTE_STORAGE_COMMITMENT is the commitment to the note's storage.
#! - STORAGE is the data corresponding to the note's storage.
#!
#! Panics if:
#! - no note is currently active.
#!
#! Invocation: exec
pub proc get_storage(dest_ptr: MemoryAddress) -> u16
    # pad the stack
    padw padw padw push.0.0
    # => [pad(14), dest_ptr]

    # push the flag indicating that we want to request inputs info from the active note
    push.1
    # => [is_active_note = 1, pad(14), dest_ptr]

    push.INPUT_NOTE_GET_STORAGE_INFO_OFFSET
    # => [offset, is_active_note = 1, pad(14), dest_ptr]

    syscall.exec_kernel_proc
    # => [NOTE_STORAGE_COMMITMENT, num_storage_items, pad(11), dest_ptr]

    # clean the stack
    swapdw dropw dropw
    movup.5 drop movup.5 drop movup.5 drop
    # => [NOTE_STORAGE_COMMITMENT, num_storage_items, dest_ptr]

    # save num_storage_items for the return value
    dup.4 movdn.6
    # => [NOTE_STORAGE_COMMITMENT, num_storage_items, dest_ptr, num_storage_items]

    # write the inputs to the provided destination pointer
    exec.write_storage_to_memory
    # => [num_storage_items]
end

#! Returns the metadata of the active note.
#!
#! Inputs:  []
#! Outputs: [METADATA]
#!
#! Where:
#! - METADATA is the metadata of the active note.
#!
#! Panics if:
#! - no note is currently active.
#!
#! Invocation: exec
pub proc get_metadata() -> NoteMetadata
    # pad the stack
    padw padw padw push.0.0
    # => [pad(14)]

    # push the flag indicating that we want to request metadata from the active note
    push.1
    # => [is_active_note = 1, pad(14)]

    push.INPUT_NOTE_GET_METADATA_OFFSET
    # => [offset, is_active_note = 1, pad(14)]

    syscall.exec_kernel_proc
    # => [METADATA, pad(12)]

    # clean the stack
    swapdw dropw dropw swapw dropw
    # => [METADATA]
end

#! Returns whether the active note is public.
#!
#! Inputs:  []
#! Outputs: [is_public]
#!
#! Where:
#! - is_public is 1 if the active note is public, 0 otherwise.
#!
#! Panics if:
#! - no note is currently active.
#!
#! Invocation: exec
pub proc is_public() -> Bool
    exec.get_metadata
    # => [METADATA]

    exec.note::metadata_into_note_type
    # => [note_type]

    eq.NOTE_TYPE_PUBLIC
    # => [is_public]
end

#! Returns whether the active note is private.
#!
#! Inputs:  []
#! Outputs: [is_private]
#!
#! Where:
#! - is_private is 1 if the active note is private, 0 otherwise.
#!
#! Panics if:
#! - no note is currently active.
#!
#! Invocation: exec
pub proc is_private() -> Bool
    exec.get_metadata
    # => [METADATA]

    exec.note::metadata_into_note_type
    # => [note_type]

    eq.NOTE_TYPE_PRIVATE
    # => [is_private]
end

#! Returns the sender of the active note.
#!
#! Inputs:  []
#! Outputs: [sender_id_suffix, sender_id_prefix]
#!
#! Where:
#! - sender_{suffix,prefix} are the suffix and prefix felts of the sender of the active note.
#!
#! Panics if:
#! - no note is currently active.
#!
#! Invocation: exec
pub proc get_sender() -> AccountId
    exec.get_metadata
    # => [METADATA]

    # extract the sender ID from the metadata
    exec.note::metadata_into_sender
    # => [sender_id_suffix, sender_id_prefix]
end

#! Returns the serial number of the active note.
#!
#! Inputs:  []
#! Outputs: [SERIAL_NUMBER]
#!
#! Where:
#! - SERIAL_NUMBER is the serial number of the active note.
#!
#! Panics if:
#! - no note is currently active.
#!
#! Invocation: exec
pub proc get_serial_number() -> word
    # pad the stack
    padw padw padw push.0.0
    # => [pad(14)]

    # push the flag indicating that we want to request serial number from the active note
    push.1
    # => [is_active_note = 1, pad(14)]

    push.INPUT_NOTE_GET_SERIAL_NUMBER_OFFSET
    # => [offset, is_active_note = 1, pad(14)]

    syscall.exec_kernel_proc
    # => [SERIAL_NUMBER, pad(12)]

    # clean the stack
    swapdw dropw dropw swapw dropw
    # => [SERIAL_NUMBER]
end

#! Returns the script root of the active note.
#!
#! Inputs:  []
#! Outputs: [SCRIPT_ROOT]
#!
#! Where:
#! - SCRIPT_ROOT is the script root of the active note.
#!
#! Panics if:
#! - no note is currently active.
#!
#! Invocation: exec
pub proc get_script_root() -> NoteScriptRoot
    # pad the stack
    padw padw padw push.0.0
    # => [pad(14)]

    # push the flag indicating that we want to request script root from the active note
    push.1
    # => [is_active_note = 1, pad(14)]

    push.INPUT_NOTE_GET_SCRIPT_ROOT_OFFSET
    # => [offset, is_active_note = 1, pad(14)]

    syscall.exec_kernel_proc
    # => [SCRIPT_ROOT, pad(12)]

    # clean the stack
    swapdw dropw dropw swapw dropw
    # => [SCRIPT_ROOT]
end

# HELPER PROCEDURES
# =================================================================================================

#! Writes the note storage stored in the advice map to the memory specified by the provided
#! destination pointer.
#!
#! Inputs:
#!   Operand stack: [NOTE_STORAGE_COMMITMENT, num_storage_items, dest_ptr]
#!   Advice map: {
#!     NOTE_STORAGE_COMMITMENT: [[INPUT_VALUES]]
#!   }
#! Outputs:
#!   Operand stack: []
proc write_storage_to_memory
    # load the inputs from the advice map to the advice stack, padded to the next multiple of 8 so
    # they can be written with `pipe_elements_preimage_to_memory`
    adv.push_mapvaln.8
    # OS => [NOTE_STORAGE_COMMITMENT, num_storage_items, dest_ptr]
    # AS => [advice_num_storage_items, [INPUT_VALUES]]

    # move the number of inputs obtained from advice map to the operand stack
    adv_push dup.5
    # OS => [num_storage_items, advice_num_storage_items, NOTE_STORAGE_COMMITMENT, num_storage_items, dest_ptr]
    # AS => [[INPUT_VALUES]]

    assert_eq.err=ERR_NOTE_INVALID_NUMBER_OF_STORAGE_ITEMS
    # OS => [NOTE_STORAGE_COMMITMENT, num_storage_items, dest_ptr]
    # AS => [[INPUT_VALUES]]

    # write the inputs to memory and compute their commitment
    movup.5 movup.5
    # OS => [num_storage_items, dest_ptr, NOTE_STORAGE_COMMITMENT]
    # AS => [[INPUT_VALUES]]

    exec.mem::pipe_elements_preimage_to_memory
    # OS => [COMPUTED_COMMITMENT, NOTE_STORAGE_COMMITMENT]

    # validate that the inputs written to memory match the inputs commitment
    assert_eqw.err=ERR_NOTE_DATA_DOES_NOT_MATCH_COMMITMENT
    # => []
end

# ATTACHMENTS
# =================================================================================================

#! Returns the commitment over all attachments of the active note.
#!
#! Inputs:  []
#! Outputs: [ATTACHMENTS_COMMITMENT]
#!
#! Where:
#! - ATTACHMENTS_COMMITMENT is the commitment to all attachments of the note, or the EMPTY_WORD if
#!   the note does not have any attachments.
#!
#! Panics if:
#! - no note is currently active.
#!
#! Invocation: exec
pub proc get_attachments_commitment() -> word
    # push a placeholder note_index (ignored when is_active_note = 1) and the active note flag
    push.0.1
    # => [is_active_note = 1, note_index = 0]

    exec.input_note::get_attachments_commitment_raw
    # => [ATTACHMENTS_COMMITMENT]
end

#! Writes the attachment commitments of the active note to the provided destination pointer.
#!
#! Inputs:  [dest_ptr]
#! Outputs: [num_attachments]
#!
#! Where:
#! - dest_ptr is the memory address to which to write the attachment commitments.
#! - num_attachments is the number of attachments in the note.
#!
#! Panics if:
#! - no note is currently active.
#! - the sequential hash over the attachment commitments in the advice inputs does not match the
#!   attachments commitment.
#!
#! Invocation: exec
pub proc write_attachment_commitments_to_memory(dest_ptr: MemoryAddress) -> u8
    exec.get_attachments_commitment
    # => [ATTACHMENTS_COMMITMENT, dest_ptr]

    exec.note::write_attachment_commitments_to_memory
    # => [num_attachments]
end

#! Writes the attachment with the provided index from the active note the provided destination
#! pointer.
#!
#! Inputs:  [dest_ptr, attachment_idx]
#! Outputs: [num_words]
#!
#! Where:
#! - dest_ptr is the memory address to which to write the attachment data.
#! - attachment_idx is the index of the attachment to retrieve.
#! - num_words is the number of words in the attachment.
#!
#! Panics if:
#! - no note is currently active.
#! - the attachment index is greater or equal to the number of attachments.
#! - the sequential hash over the attachment data in the advice inputs does not match the
#!   attachment commitment.
#!
#! Invocation: exec
@locals(16)
pub proc write_attachment_to_memory(dest_ptr: MemoryAddress, attachment_idx: u8) -> u16
    # write the attachment commitments to local memory
    # we allocate 16 elements of memory (max num attachments * WORD_SIZE) to store all
    # four attachment commitments
    swap locaddr.0
    # => [attachment_commitments_ptr, attachment_idx, dest_ptr]

    exec.write_attachment_commitments_to_memory
    # => [num_attachments, attachment_idx, dest_ptr]

    locaddr.0 swap
    # => [num_attachments, attachment_commitments_ptr, attachment_idx, dest_ptr]

    exec.note::write_indexed_attachment_to_memory
    # => [num_words]
end

#! Searches the metadata of the active note for the specified attachment scheme and returns
#! the index of the first matching slot.
#!
#! Inputs:  [attachment_scheme]
#! Outputs: [is_found, attachment_idx]
#!
#! Where:
#! - attachment_scheme is the scheme of the attachment to find.
#! - is_found is 1 if the attachment with the provided scheme was found, 0 otherwise.
#! - attachment_idx is the index (0-3) of the first matching slot, or undefined if not found.
#!
#! Panics if:
#! - no note is currently active.
#!
#! Invocation: exec
pub proc find_attachment(attachment_scheme: u16) -> (Bool, u8)
    exec.get_metadata
    # => [METADATA, attachment_scheme]

    movup.4
    # => [attachment_scheme, METADATA]

    exec.note::find_attachment_idx
    # => [is_found, attachment_idx]
end