# The MASM code of the Multi-Signature Falcon 512 RPO Authentication Component.
#
# See the `AuthFalcon512RpoMultisig` Rust type's documentation for more details.

use miden::protocol::active_account
use miden::protocol::native_account
use miden::standards::auth

type BeWord = struct @bigendian { a: felt, b: felt, c: felt, d: felt }

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

# Auth Request Constants

# The event emitted when a signature is not found for a required signer.
const AUTH_UNAUTHORIZED_EVENT = event("miden::auth::unauthorized")

# Storage Layout Constants
#
# ┌───────────────────────────────┬──────────┬──────────────┬───────────────────┐
# │ THRESHOLD & APPROVERS CONFIG  │ PUB KEYS │ EXECUTED TXS │  PROC THRESHOLDS  │
# │           (slot)              │   (map)  │    (map)     │       (map)       │
# ├───────────────────────────────┼──────────┼──────────────┼───────────────────┤
# │              0                │    1     │      2       │         3         │
# └───────────────────────────────┴──────────┴──────────────┴───────────────────┘

# The slot in this component's storage layout where the default signature threshold and
# number of approvers are stored as:
# [default_threshold, num_approvers, 0, 0].
# The threshold is guaranteed to be less than or equal to num_approvers.
const THRESHOLD_CONFIG_SLOT = word("miden::standards::auth::falcon512_rpo_multisig::threshold_config")

# The slot in this component's storage layout where the public keys map is stored.
# Map entries: [key_index, 0, 0, 0] => APPROVER_PUBLIC_KEY
const APPROVER_PUBLIC_KEYS_SLOT = word("miden::standards::auth::falcon512_rpo_multisig::approver_public_keys")

# The slot in this component's storage layout where executed transactions are stored.
# Map entries: transaction_message => [is_executed, 0, 0, 0]
const EXECUTED_TXS_SLOT = word("miden::standards::auth::falcon512_rpo_multisig::executed_transactions")

# The slot in this component's storage layout where procedure thresholds are stored.
# Map entries: PROC_ROOT => [proc_threshold, 0, 0, 0]
const PROC_THRESHOLD_ROOTS_SLOT = word("miden::standards::auth::falcon512_rpo_multisig::procedure_thresholds")

# Executed Transaction Flag Constant
const IS_EXECUTED_FLAG = [1, 0, 0, 0]

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

const ERR_TX_ALREADY_EXECUTED = "failed to approve multisig transaction as it was already executed"

const ERR_MALFORMED_MULTISIG_CONFIG = "number of approvers must be equal to or greater than threshold"

const ERR_ZERO_IN_MULTISIG_CONFIG = "number of approvers or threshold must not be zero"

#! Check if transaction has already been executed and add it to executed transactions for replay protection.
#!
#! Inputs:  [MSG]
#! Outputs: []
#!
#! Panics if:
#! - the same transaction has already been executed
proc assert_new_tx(msg: BeWord)
    push.IS_EXECUTED_FLAG
    # => [[0, 0, 0, is_executed], MSG]

    swapw
    # => [MSG, IS_EXECUTED_FLAG]

    push.EXECUTED_TXS_SLOT[0..2]
    # => [txs_slot_prefix, txs_slot_suffix, MSG, IS_EXECUTED_FLAG]

    # Set the key value pair in the map to mark transaction as executed
    exec.native_account::set_map_item
    # => [[0, 0, 0, is_executed]]

    drop drop drop
    # => [is_executed]

    assertz.err=ERR_TX_ALREADY_EXECUTED
    # => []
end

#! Remove old approver public keys from the approver public key mapping.
#!
#! This procedure cleans up the storage by removing public keys of approvers that are no longer
#! part of the multisig configuration. This procedure assumes that init_num_of_approvers and
#! new_num_of_approvers are u32 values.
#!
#! Inputs: [init_num_of_approvers, new_num_of_approvers]
#! Outputs: []
#!
#! Where:
#! - init_num_of_approvers is the original number of approvers before the update
#! - new_num_of_approvers is the new number of approvers after the update
proc cleanup_pubkey_mapping(init_num_of_approvers: u32, new_num_of_approvers: u32)
    dup.1 dup.1
    u32assert2 u32lt
    # => [should_loop, i = init_num_of_approvers, new_num_of_approvers]

    while.true
        # => [i, new_num_of_approvers]

        sub.1
        # => [i-1, new_num_of_approvers]

        dup
        # => [i-1, i-1, new_num_of_approvers]

        push.0.0.0
        # => [[0, 0, 0, i-1], i-1, new_num_of_approvers]

        padw swapw
        # => [[0, 0, 0, i-1], EMPTY_WORD, i-1, new_num_of_approvers]

        push.APPROVER_PUBLIC_KEYS_SLOT[0..2]
        # => [pub_key_slot_prefix, pub_key_slot_suffix, [0, 0, 0, i-1], EMPTY_WORD, i-1, new_num_of_approvers]

        exec.native_account::set_map_item
        # => [OLD_VALUE, i-1, new_num_of_approvers]

        dropw
        # => [i-1, new_num_of_approvers]

        dup.1 dup.1
        u32lt
        # => [should_loop, i-1, new_num_of_approvers]
    end

    drop drop
    # => []
end

#! Update threshold config and add / remove approvers
#!
#! Inputs:
#!   Operand stack: [MULTISIG_CONFIG_HASH, pad(12)]
#!   Advice map: {
#!     MULTISIG_CONFIG_HASH => [CONFIG, PUB_KEY_N, PUB_KEY_N-1, ..., PUB_KEY_0]
#!   }
#! Outputs:
#!   Operand stack: []
#!
#! Where:
#! - MULTISIG_CONFIG_HASH is the hash of the threshold and new public key vector
#! - MULTISIG_CONFIG is [threshold, num_approvers, 0, 0]
#! - PUB_KEY_i is the public key of the i-th signer
#!
#! Locals:
#! 0: new_num_of_approvers
#! 1: init_num_of_approvers
@locals(2)
pub proc update_signers_and_threshold(multisig_config_hash: BeWord)
    adv.push_mapval
    # => [MULTISIG_CONFIG_HASH, pad(12)]

    adv_loadw
    # => [MULTISIG_CONFIG, pad(12)]

    # store new_num_of_approvers for later
    dup.2 loc_store.0
    # => [MULTISIG_CONFIG, pad(12)]

    dup.3 dup.3
    # => [num_approvers, threshold, MULTISIG_CONFIG, pad(12)]

    # make sure that the threshold is smaller than the number of approvers
    u32assert2.err=ERR_MALFORMED_MULTISIG_CONFIG
    u32gt assertz.err=ERR_MALFORMED_MULTISIG_CONFIG
    # => [MULTISIG_CONFIG, pad(12)]

    dup.3 dup.3
    # => [num_approvers, threshold, MULTISIG_CONFIG, pad(12)]

    # make sure that threshold or num_approvers are not zero
    eq.0 assertz.err=ERR_ZERO_IN_MULTISIG_CONFIG
    eq.0 assertz.err=ERR_ZERO_IN_MULTISIG_CONFIG
    # => [MULTISIG_CONFIG, pad(12)]

    push.THRESHOLD_CONFIG_SLOT[0..2]
    # => [config_slot_prefix, config_slot_suffix, MULTISIG_CONFIG, pad(12)]

    exec.native_account::set_item
    # => [OLD_THRESHOLD_CONFIG, pad(12)]

    # store init_num_of_approvers for later
    drop drop loc_store.1 drop
    # => [pad(12)]

    loc_load.0
    # => [num_approvers]

    dup neq.0
    while.true
        sub.1
        # => [i-1, pad(12)]

        dup push.0.0.0
        # => [[0, 0, 0, i-1], i-1, pad(12)]

        padw adv_loadw
        # => [PUB_KEY, [0, 0, 0, i-1], i-1, pad(12)]

        swapw
        # => [[0, 0, 0, i-1], PUB_KEY, i-1, pad(12)]

        push.APPROVER_PUBLIC_KEYS_SLOT[0..2]
        # => [pub_key_slot_prefix, pub_key_slot_suffix, [0, 0, 0, i-1], PUB_KEY, i-1, pad(12)]

        exec.native_account::set_map_item
        # => [OLD_VALUE, i-1, pad(12)]

        dropw
        # => [i-1, pad(12)]

        dup neq.0
        # => [is_non_zero, i-1, pad(12)]
    end
    # => [pad(13)]

    drop
    # => [pad(12)]

    # compare initial vs current multisig config

    # load init_num_of_approvers & new_num_of_approvers
    loc_load.0 loc_load.1
    # => [init_num_of_approvers, new_num_of_approvers, pad(12)]

    exec.cleanup_pubkey_mapping
    # => [pad(12)]
end

# Computes the effective transaction threshold based on called procedures and per-procedure
# overrides stored in PROC_THRESHOLD_ROOTS_SLOT. Falls back to default_threshold if no
# overrides apply.
#
#! Inputs:  [default_threshold]
#! Outputs: [transaction_threshold]
@locals(1)
proc compute_transaction_threshold(default_threshold: u32) -> u32
    # 1. initialize transaction_threshold = 0
    # 2. iterate through all account procedures
    #   a. check if the procedure was called during the transaction
    #   b. if called, get the override threshold of that procedure from the config map
    #   c. if proc_threshold > transaction_threshold, set transaction_threshold = proc_threshold
    # 3. if transaction_threshold == 0 at the end, revert to using default_threshold

    # store default_threshold for later
    loc_store.0
    # => []

    # 1. initialize transaction_threshold = 0
    push.0
    # => [transaction_threshold]

    # get the number of account procedures
    exec.active_account::get_num_procedures
    # => [num_procedures, transaction_threshold]

    # 2. iterate through all account procedures
    dup neq.0
    # => [should_continue, num_procedures, transaction_threshold]
    while.true
        sub.1 dup
        # => [num_procedures-1, num_procedures-1, transaction_threshold]

        # get procedure root of the procedure with index i
        exec.active_account::get_procedure_root dupw
        # => [PROC_ROOT, PROC_ROOT, num_procedures-1, transaction_threshold]

        # 2a. check if this procedure has been called in the transaction
        exec.native_account::was_procedure_called
        # => [was_called, PROC_ROOT, num_procedures-1, transaction_threshold]

        # if it has been called, get the override threshold of that procedure
        if.true
            # => [PROC_ROOT, num_procedures-1, transaction_threshold]

            push.PROC_THRESHOLD_ROOTS_SLOT[0..2]
            # => [proc_roots_slot_prefix, proc_roots_slot_suffix, PROC_ROOT, num_procedures-1, transaction_threshold]

            # 2b. get the override proc_threshold of that procedure
            # if the procedure has no override threshold, the returned map item will be [0, 0, 0, 0]
            exec.active_account::get_initial_map_item
            # => [[0, 0, 0, proc_threshold], num_procedures-1, transaction_threshold]

            drop drop drop dup dup.3
            # => [transaction_threshold, proc_threshold, proc_threshold, num_procedures-1, transaction_threshold]

            u32assert2.err="transaction threshold or procedure threshold are not u32"
            u32gt
            # => [is_gt, proc_threshold, num_procedures-1, transaction_threshold]
            # 2c. if proc_threshold > transaction_threshold, update transaction_threshold
            movup.2 movdn.3
            # => [is_gt, proc_threshold, transaction_threshold, num_procedures-1]
            cdrop
            # => [updated_transaction_threshold, num_procedures-1]
            swap
            # => [num_procedures-1, updated_transaction_threshold]
        # if it has not been called during this transaction, nothing to do, move to the next procedure
        else
            dropw
            # => [num_procedures-1, transaction_threshold]
        end

        dup neq.0
        # => [should_continue, num_procedures-1, transaction_threshold]
    end

    drop
    # => [transaction_threshold]

    loc_load.0
    # => [default_threshold, transaction_threshold]

    # 3. if transaction_threshold == 0 at the end, revert to using default_threshold
    dup.1 eq.0
    # => [is_zero, default_threshold, transaction_threshold]

    cdrop
    # => [effective_transaction_threshold]
end

#! Authenticate a transaction using the Falcon signature scheme with multi-signature support.
#!
#! This procedure implements multi-signature authentication by:
#! 1. Computing the transaction summary message that needs to be signed
#! 2. Verifying signatures from multiple required signers against their public keys
#! 3. Ensuring the minimum threshold of valid signatures is met
#! 4. Implementing replay protection by tracking executed transactions
#!
#! Inputs:
#!   Operand stack: [SALT]
#!   Advice map: {
#!     h(SIG_0, MSG): SIG_0,
#!     h(SIG_1, MSG): SIG_1,
#!     h(SIG_n, MSG): SIG_n
#!   }
#! Outputs:
#!   Operand stack: []
#!
#! Where:
#! - SALT is a cryptographically random nonce that enables multiple concurrent
#!   multisig transactions while maintaining replay protection. Each transaction
#!   must use a unique SALT value to ensure transaction uniqueness.
#! - SIG_i is the signature from the i-th signer.
#! - MSG is the transaction message being signed.
#! - h(SIG_i, MSG) is the hash of the signature and message used as the advice map key.
#!
#! Panics if:
#! - insufficient number of valid signatures (below threshold).
#! - the same transaction has already been executed (replay protection).
#!
#! Invocation: call
@locals(1)
pub proc auth_tx_falcon512_rpo_multisig(salt: BeWord)
    exec.native_account::incr_nonce drop
    # => [SALT]

    # ------ Computing transaction summary ------

    exec.auth::create_tx_summary
    # => [SALT, OUTPUT_NOTES_COMMITMENT, INPUT_NOTES_COMMITMENT, ACCOUNT_DELTA_COMMITMENT]

    # to build a tx_summary in the host, we need these four words in the advice provider
    exec.auth::adv_insert_hqword
    # => [SALT, OUTPUT_NOTES_COMMITMENT, INPUT_NOTES_COMMITMENT, ACCOUNT_DELTA_COMMITMENT]

    # the commitment to the tx summary is the message that is signed
    exec.auth::hash_tx_summary
    # => [TX_SUMMARY_COMMITMENT]

    # ------ Verifying approver signatures ------

    push.THRESHOLD_CONFIG_SLOT[0..2]
    # => [config_slot_prefix, config_slot_suffix, TX_SUMMARY_COMMITMENT]

    exec.active_account::get_initial_item
    # => [0, 0, num_of_approvers, default_threshold, TX_SUMMARY_COMMITMENT]

    drop drop
    # => [num_of_approvers, default_threshold, TX_SUMMARY_COMMITMENT]

    swap movdn.5
    # => [num_of_approvers, TX_SUMMARY_COMMITMENT, default_threshold]

    push.APPROVER_PUBLIC_KEYS_SLOT[0..2]
    # => [pub_key_slot_prefix, pub_key_slot_suffix, num_of_approvers, TX_SUMMARY_COMMITMENT, default_threshold]

    exec.::miden::standards::auth::falcon512_rpo::verify_signatures
    # => [num_verified_signatures, TX_SUMMARY_COMMITMENT, default_threshold]

    # ------ Checking threshold is >= num_verified_signatures ------

    movup.5
    # => [default_threshold, num_verified_signatures, TX_SUMMARY_COMMITMENT]

    exec.compute_transaction_threshold
    # => [transaction_threshold, num_verified_signatures, TX_SUMMARY_COMMITMENT]

    u32assert2 u32lt
    # => [is_unauthorized, TX_SUMMARY_COMMITMENT]

    # If signatures are non-existent the tx will fail here.
    if.true
        emit.AUTH_UNAUTHORIZED_EVENT
        push.0 assert.err="insufficient number of signatures"
    end

    # ------ Writing executed transaction MSG to map ------

    exec.assert_new_tx
    # => [TX_SUMMARY_COMMITMENT]
end
