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

use.kernel::constants
use.kernel::memory

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

# Attempted to access note sender from incorrect context
const.ERR_NOTE_ATTEMPT_TO_ACCESS_NOTE_SENDER_FROM_INCORRECT_CONTEXT=0x00020031

# Attempted to access note assets from incorrect context
const.ERR_NOTE_ATTEMPT_TO_ACCESS_NOTE_ASSETS_FROM_INCORRECT_CONTEXT=0x00020032

# Attempted to access note inputs from incorrect context
const.ERR_NOTE_ATTEMPT_TO_ACCESS_NOTE_INPUTS_FROM_INCORRECT_CONTEXT=0x00020033

# Number of assets in a note exceed 255
const.ERR_NOTE_NUM_OF_ASSETS_EXCEED_LIMIT=0x0002002F

# CONSTANTS
# =================================================================================================

# The diff between the memory address after first mem_stream operation and the next target when
# generating the output notes commitment. Must be NOTE_MEM_SIZE - 2;
const.OUTPUT_NOTE_HASHING_MEM_DIFF=510

# INPUT NOTE PROCEDURES
# =================================================================================================

#! Returns the sender of the note currently being processed. Panics if a note is not being
#! processed.
#!
#! Inputs: []
#! Outputs: [sender]
#!
#! - sender is the sender of the note currently being processed.
export.get_sender
    # get the current input note pointer
    exec.memory::get_current_input_note_ptr
    # => [ptr]

    # assert the pointer is not zero - this would suggest the procedure has been called from an
    # incorrect context
    dup neq.0 assert.err=ERR_NOTE_ATTEMPT_TO_ACCESS_NOTE_SENDER_FROM_INCORRECT_CONTEXT
    # => [ptr]

    # get the sender from the note pointer
    exec.memory::get_input_note_sender
    # => [sender]
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: []
#! Outputs: [ASSETS_HASH, num_assets]
#!
#! - num_assets is the number of assets in the note currently being processed.
#! - ASSETS_HASH is the vault hash of the note currently being processed.
export.get_assets_info
    # get the current input note pointer
    exec.memory::get_current_input_note_ptr
    # => [ptr]

    # assert the pointer is not zero - this would suggest the procedure has been called from an
    # incorrect context
    dup neq.0 assert.err=ERR_NOTE_ATTEMPT_TO_ACCESS_NOTE_ASSETS_FROM_INCORRECT_CONTEXT
    # => [ptr]

    # get the number of assets in the note
    dup exec.memory::get_input_note_num_assets
    # => [num_assets, ptr]

    # get the assets hash from the note pointer
    swap exec.memory::get_input_note_assets_hash
    # => [ASSETS_HASH, num_assets]
end

#! Returns the commitment to the note's inputs.
#!
#! Panics if a note is not being processed.
#!
#! Inputs: []
#! Outputs: [NOTE_INPUTS_HASH]
#!
#! Where:
#! - NOTE_INPUTS_HASH is the note inputs hash of the note currently being processed.
export.get_note_inputs_hash
    exec.memory::get_current_input_note_ptr
    # => [ptr]

    # The kernel memory is initialized by prologue::process_input_notes_data, and reset by
    # note_processing_teardown before running the tx_script. If the value is `0` it is likely this
    # procedure is being called outside of the kernel context.
    dup neq.0 assert.err=ERR_NOTE_ATTEMPT_TO_ACCESS_NOTE_INPUTS_FROM_INCORRECT_CONTEXT
    # => [ptr]

    exec.memory::get_input_note_inputs_hash
    # => [NOTE_INPUTS_HASH]
end

#! Move the current input note pointer to the next note and returns the pointer value.
#!
#! Inputs: []
#! Outputs: [current_input_note_ptr]
#!
#! Where:
#! - current_input_note_ptr is the pointer to the next note to be processed.
export.increment_current_input_note_ptr
    # get the current input note pointer
    exec.memory::get_current_input_note_ptr
    # => [orig_input_note_ptr]

    # increment the pointer
    exec.constants::get_note_mem_size add
    # => [current_input_note_ptr]

    # set the current input note pointer to the incremented value
    dup exec.memory::set_current_input_note_ptr
    # => [current_input_note_ptr]
end

#! Sets the current input note pointer to 0. This should be called after all input notes have
#! been processed.
#!
#! Inputs: []
#! Outputs: []
export.note_processing_teardown
    # set the current input note pointer to 0
    push.0 exec.memory::set_current_input_note_ptr
    # => []
end

#! Prepares a note for execution.
#!
#! Loads the note's script root address and args onto the stack.
#!
#! Stack: []
#! Output: [note_script_root_ptr, NOTE_ARGS]
#!
#! Where:
#! - note_script_root_ptr, the memory address where note's script root is stored.
#! - NOTE_ARGS, the note's arguments.
export.prepare_note
    exec.memory::get_current_input_note_ptr
    # => [note_ptr]

    dup exec.memory::get_input_note_args movup.4
    # => [note_ptr, NOTE_ARGS]

    exec.memory::get_input_note_script_root_ptr
    # => [note_script_root_ptr, NOTE_ARGS]
end

# OUTPUT NOTE PROCEDURES
# =================================================================================================

#! Computes the assets hash of the output note located at the specified memory address.
#!
#! The hash is computed as a sequential hash of the assets contained in the note. If there is an
#! odd number of assets, then for the final hashing permutation we pad the last word of the hasher
#! rate with zeros. If the note contains no assets, ASSET_HASH is set to EMPTY_WORD.
#!
#! Stack: [note_data_ptr]
#! Output: [ASSETS_HASH]
#!
#! - note_data_ptr is a pointer to the data section of the output note.
#! - ASSETS_HASH is the hash of the assets of the output note located at note_data_ptr.
proc.compute_output_note_assets_hash
    # duplicate note pointer and fetch num_assets
    dup dup exec.memory::get_output_note_num_assets
    # => [num_assets, note_data_ptr, note_data_ptr]

    # calculate the number of pairs of assets (takes ceiling if we have an odd number)
    add.1
    u32assert.err=ERR_NOTE_NUM_OF_ASSETS_EXCEED_LIMIT
    u32div.2
    # => [num_asset_pairs, note_data_ptr, note_data_ptr]

    # initiate counter for assets
    push.0
    # => [asset_counter, num_asset_pairs, note_data_ptr, note_data_ptr]

    # prepare address and stack for reading assets
    movup.2 exec.memory::get_output_note_asset_data_ptr padw padw padw
    # => [PAD, PAD, PAD, asset_data_ptr, asset_counter, num_asset_pairs, note_data_ptr]

    # check if we should loop
    dup.14 dup.14 neq
    # => [should_loop, PAD, PAD, PAD, asset_data_ptr, asset_counter, num_asset_pairs, note_data_ptr]

    # loop and read assets from memory
    while.true
        # read assets from memory.
        # if this is the last permutation of the loop and we have an odd number of assets then we
        # implicitly pad the last word of the hasher rate with zeros by reading from empty memory.
        mem_stream hperm
        # => [PERM, PERM, PERM, asset_data_ptr, asset_counter, num_asset_pairs, note_data_ptr]

        # check if we should loop again
        movup.13 add.1 dup movdn.14 dup.15 neq
        # => [should_loop, PERM, PERM, PERM, asset_data_ptr, asset_counter, num_asset_pairs,
        #     note_data_ptr]
    end

    # extract digest
    exec.rpo::squeeze_digest
    # => [ASSETS_HASH, asset_data_ptr, asset_counter, num_asset_pairs, note_data_ptr]

    # drop accessory variables from stack
    movup.4 drop
    movup.4 drop
    movup.4 drop
    # => [ASSETS_HASH, note_data_ptr]

    # save vault hash to memory
    movup.4 exec.memory::set_output_note_assets_hash
    # => []
end

#! Computes the ID of an output note located at the specified memory address.
#!
#! The note ID is computed as follows:
#! - we define, recipient =
#!       hash(hash(hash(serial_num, [0; 4]), script_hash), input_hash)
#! - we then compute the output note ID as:
#!       hash(recipient, assets_hash)
#!
#! Stack: [note_data_ptr]
#! Output: [NOTE_ID]
#!
#! - note_data_ptr is a pointer to the data section of the output note.
#! - NOTE_ID is the ID of the output note located at note_data_ptr.
proc.compute_output_note_id
    # pad capacity elements of hasher
    padw

    # insert output note recipient into the first four elements of the hasher rate
    dup.4 exec.memory::get_output_note_recipient

    # populate the last four elements of the hasher rate with the output note's asset hash
    dup.8 exec.compute_output_note_assets_hash

    # compute output note commitment and extract digest
    hperm exec.rpo::squeeze_digest

    # save the output notes commitment to memory
    movup.4 mem_storew
end

#! Computes a commitment to the output notes. This is computed as a sequential hash of
#! (note_id, note_metadata) tuples.
#!
#! Stack: []
#! Output: [OUTPUT_NOTES_COMMITMENT]
#!
#! - OUTPUT_NOTES_COMMITMENT is the commitment to the notes output by the transaction.
export.compute_output_notes_commitment
    # get the number of output notes from memory
    exec.memory::get_num_output_notes
    # => [num_notes, ...]

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

    # compute pointer for first address
    push.0 exec.memory::get_output_note_ptr
    # => [first_note_ptr, end_ptr, ...]

    # prepare stack for hashing
    padw padw padw
    # => [PERM, PERM, PERM, first_note_ptr, end_ptr, ...]

    # check if the number of output notes is greater then 0. Conditional for the while loop.
    dup.13 dup.13 neq
    # => [PERM, PERM, PERM, first_note_ptr, end_ptr, ...]

    # loop and hash output notes
    while.true
        # compute and save output note ID to memory (this also computes the note's asset hash)
        dup.12 exec.compute_output_note_id
        # => [NOTE_ID, PERM, PERM, PERM, note_ptr, end_ptr, ...]

        # drop output note ID from stack (it will be read from memory by the next instruction)
        dropw
        # => [PERM, PERM, PERM, note_ptr, end_ptr, ...]

        # permute over (note_id, note_metadata)
        mem_stream hperm
        # => [PERM, PERM, PERM, note_ptr + 2, end_ptr, ...]

        # increment output note pointer
        movup.12 push.OUTPUT_NOTE_HASHING_MEM_DIFF add
        # => [note_ptr + 512, PERM, PERM, PERM, end_ptr, ...]

        # check if we should loop again
        dup movdn.13 dup.14 neq
        # => [should_loop, PERM, PERM, PERM, note_ptr + 512, end_ptr, ...]
    end

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

    # drop accessory variables from stack
    movup.4 drop
    movup.4 drop
    # => [OUTPUT_NOTES_COMMITMENT, ...]
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_serial_number
    exec.memory::get_current_input_note_ptr
    # => [note_ptr, ...]

    exec.memory::get_input_note_serial_num
    # => [SERIAL_NUMBER, ...]
end