miden_protocol/transaction/kernel/

advice_inputs.rs

use alloc::vec::Vec;

use miden_processor::advice::AdviceMutation;

use crate::account::PartialAccount;
use crate::block::account_tree::{AccountIdKey, AccountWitness};
use crate::crypto::SequentialCommit;
use crate::crypto::merkle::InnerNodeInfo;
use crate::protocol_config::ProtocolConfig;
use crate::transaction::{AccountInputs, InputNote, PartialBlockchain, TransactionInputs};
use crate::vm::AdviceInputs;
use crate::{EMPTY_WORD, Felt, Word, ZERO};

// TRANSACTION ADVICE INPUTS
// ================================================================================================

/// Advice inputs wrapper for inputs that are meant to be used exclusively in the transaction
/// kernel.
#[derive(Debug, Clone, Default)]
pub struct TransactionAdviceInputs(AdviceInputs);

impl TransactionAdviceInputs {
    /// Creates a [`TransactionAdviceInputs`].
    ///
    /// The created advice inputs will be populated with the data required for executing a
    /// transaction with the specified transaction inputs.
    pub fn new(tx_inputs: &TransactionInputs) -> Self {
        let mut inputs = TransactionAdviceInputs(tx_inputs.advice_inputs().clone());

        inputs.build_stack(tx_inputs);
        inputs.add_protocol_config(tx_inputs.protocol_config());
        inputs.add_partial_blockchain(tx_inputs.blockchain());
        inputs.add_input_notes(tx_inputs);

        // Add the script's MAST forest's advice inputs.
        if let Some(tx_script) = tx_inputs.tx_args().tx_script() {
            inputs.extend_map(
                tx_script
                    .mast()
                    .advice_map()
                    .iter()
                    .map(|(key, values)| (*key, values.to_vec())),
            );
        }

        // Inject native account.
        let partial_native_acc = tx_inputs.account();
        inputs.add_account(partial_native_acc);

        // If a seed was provided, extend the map appropriately.
        if let Some(seed) = tx_inputs.account().seed() {
            // ACCOUNT_ID |-> ACCOUNT_SEED
            let account_id_key = AccountIdKey::from(partial_native_acc.id());
            inputs.add_map_entry(account_id_key.as_word(), seed.to_vec());
        }

        // if the account is new, insert the storage map entries into the advice provider.
        if partial_native_acc.is_new() {
            for storage_map in partial_native_acc.storage().maps() {
                let map_entries = storage_map
                    .entries()
                    .flat_map(|(key, value)| {
                        value.as_elements().iter().chain(key.as_elements().iter()).copied()
                    })
                    .collect();
                inputs.add_map_entry(storage_map.root(), map_entries);
            }
        }

        // The host reads the new code from this entry when the kernel initializes the upgrade.
        if let Some(code_upgrade) = tx_inputs.tx_args().account_code_upgrade() {
            let (key, elements) = code_upgrade.to_advice_map_entry();
            inputs.add_map_entry(key, elements);
        }

        // Extend with extra user-supplied advice.
        inputs.extend(tx_inputs.tx_args().advice_inputs().clone());

        inputs
    }

    /// Returns a reference to the underlying advice inputs.
    pub fn as_advice_inputs(&self) -> &AdviceInputs {
        &self.0
    }

    /// Converts these transaction advice inputs into the underlying advice inputs.
    pub fn into_advice_inputs(self) -> AdviceInputs {
        self.0
    }

    /// Consumes self and returns an iterator of [`AdviceMutation`]s in arbitrary order.
    pub fn into_advice_mutations(self) -> impl Iterator<Item = AdviceMutation> {
        let (stack, map, store) = self.0.into_parts();
        [
            AdviceMutation::extend_map(map),
            AdviceMutation::extend_merkle_store(store.inner_nodes()),
            AdviceMutation::extend_advice_stack(stack),
        ]
        .into_iter()
    }

    // PUBLIC UTILITIES
    // --------------------------------------------------------------------------------------------

    // MUTATORS
    // --------------------------------------------------------------------------------------------

    /// Extends these advice inputs with the provided advice inputs.
    pub fn extend(&mut self, adv_inputs: AdviceInputs) {
        self.0.extend(adv_inputs);
    }

    /// Adds the provided account inputs into the advice inputs.
    pub fn add_foreign_accounts<'inputs>(
        &mut self,
        foreign_account_inputs: impl IntoIterator<Item = &'inputs AccountInputs>,
    ) {
        for foreign_acc in foreign_account_inputs {
            self.add_account(foreign_acc.account());
            self.add_account_witness(foreign_acc.witness());

            // for foreign accounts, we need to insert the id to state mapping
            // NOTE: keep this in sync with the account::load_from_advice procedure
            let account_id_key = AccountIdKey::from(foreign_acc.id());

            // ACCOUNT_ID |-> [ACCOUNT_METADATA, VAULT_ROOT, STORAGE_COMMITMENT, CODE_COMMITMENT]
            self.add_map_entry(account_id_key.as_word(), foreign_acc.account().to_elements());
        }
    }

    /// Extend the advice stack with the transaction inputs.
    ///
    /// The following data is pushed to the advice stack (words shown in memory-order):
    ///
    /// [
    ///     [version, block_num, timestamp, 0],
    ///     PREV_BLOCK_COMMITMENT,
    ///     CHAIN_COMMITMENT,
    ///     ACCOUNT_ROOT,
    ///     NULLIFIER_ROOT,
    ///     TX_COMMITMENT,
    ///     PROTOCOL_CONFIG_COMMITMENT,
    ///     VALIDATOR_CONFIG_COMMITMENT,
    ///     NEXT_PROTOCOL_CONFIG_COMMITMENT,
    ///     [verification_base_fee, 0, 0, 0],
    ///     NOTE_ROOT,
    ///     [account_version, account_nonce, account_id_suffix, account_id_prefix],
    ///     ACCOUNT_VAULT_ROOT,
    ///     ACCOUNT_STORAGE_COMMITMENT,
    ///     ACCOUNT_CODE_COMMITMENT,
    ///     number_of_input_notes,
    ///     TX_SCRIPT_ROOT,
    ///     TX_SCRIPT_ARGS,
    ///     AUTH_ARGS,
    /// ]
    fn build_stack(&mut self, tx_inputs: &TransactionInputs) {
        // --- block header data (keep in sync with kernel's process_block_data) --
        self.extend_stack(tx_inputs.block_header().to_elements());

        // --- core account items (keep in sync with process_account_data) ----
        self.extend_stack(tx_inputs.account().to_elements());

        // --- number of notes, script root and args --------------------------
        self.extend_stack([Felt::from(tx_inputs.input_notes().num_notes())]);
        let tx_args = tx_inputs.tx_args();
        self.extend_stack(
            tx_args.tx_script().map_or(Word::empty(), |script| script.root().as_word()),
        );
        self.extend_stack(tx_args.tx_script_args());

        // --- auth procedure args --------------------------------------------
        self.extend_stack(tx_args.auth_args());
    }

    // BLOCKCHAIN INJECTIONS
    // --------------------------------------------------------------------------------------------

    /// Inserts the partial blockchain data into the provided advice inputs.
    ///
    /// Inserts the following items into the Merkle store:
    /// - Inner nodes of all authentication paths contained in the partial blockchain.
    ///
    /// Inserts the following data to the advice map:
    ///
    /// > {MMR_ROOT: [[num_blocks, 0, 0, 0], PEAK_1, ..., PEAK_N]}
    ///
    /// Where:
    /// - MMR_ROOT, is the sequential hash of the padded MMR peaks
    /// - num_blocks, is the number of blocks in the MMR.
    /// - PEAK_1 .. PEAK_N, are the MMR peaks.
    fn add_partial_blockchain(&mut self, mmr: &PartialBlockchain) {
        // NOTE: keep this code in sync with the `process_chain_data` kernel procedure
        // add authentication paths from the MMR to the Merkle store
        self.extend_merkle_store(mmr.inner_nodes());

        // insert MMR peaks info into the advice map
        let peaks = mmr.peaks();
        let num_leaves = Felt::try_from(peaks.num_leaves() as u64)
            .expect("number of blocks in chain should not exceed BlockNumber::MAX");
        let mut elements = vec![num_leaves, ZERO, ZERO, ZERO];
        elements.extend(peaks.flatten_and_pad_peaks());
        self.add_map_entry(peaks.hash_peaks(), elements);
    }

    // PROTOCOL CONFIG INJECTIONS
    // --------------------------------------------------------------------------------------------

    /// Inserts the protocol configuration into the advice map.
    ///
    /// The block header only commits to the configuration, so the kernel resolves these commitments
    /// to their preimage:
    /// - PROTOCOL_CONFIG_COMMITMENT |-> the protocol config elements.
    /// - TX_KERNEL_CONFIG_COMMITMENT |-> the transaction kernel config elements.
    /// - TX_KERNEL_PROCS_COMMITMENT |-> the array of the transaction kernel's procedure roots.
    ///
    /// NOTE: keep this in sync with the `process_protocol_config` and `process_kernel_data` kernel
    /// procedures.
    fn add_protocol_config(&mut self, protocol_config: &ProtocolConfig) {
        let tx_kernel = protocol_config.tx_kernel();

        self.add_map_entry(protocol_config.to_commitment(), protocol_config.to_elements());
        self.add_map_entry(tx_kernel.to_commitment(), tx_kernel.to_elements());
        self.add_map_entry(
            tx_kernel.kernel_procs_commitment(),
            tx_kernel.kernel_procs_elements().to_vec(),
        );
    }

    // ACCOUNT INJECTION
    // --------------------------------------------------------------------------------------------

    /// Inserts account data into the advice inputs.
    ///
    /// Inserts the following items into the Merkle store:
    /// - The Merkle nodes associated with the account vault tree.
    /// - If present, the Merkle nodes associated with the account storage maps.
    ///
    /// Inserts the following entries into the advice map:
    /// - The account storage commitment |-> storage slots and types vector.
    /// - The account code commitment |-> procedures vector.
    /// - The leaf hash |-> (key, value), for all leaves of the partial vault.
    /// - If present, the Merkle leaves associated with the account storage maps.
    fn add_account(&mut self, account: &PartialAccount) {
        // --- account code -------------------------------------------------------

        // CODE_COMMITMENT -> [[ACCOUNT_PROCEDURE_DATA]]
        let code = account.code();
        self.add_map_entry(code.commitment(), code.to_elements());

        // --- account storage ----------------------------------------------------

        // STORAGE_COMMITMENT |-> [[STORAGE_SLOT_DATA]]
        let storage_header = account.storage().header();
        self.add_map_entry(storage_header.to_commitment(), storage_header.to_elements());

        // populate Merkle store and advice map with nodes info needed to access storage map entries
        self.extend_merkle_store(account.storage().inner_nodes());
        self.extend_map(
            account
                .storage()
                .leaves()
                .map(|leaf| (leaf.hash(), leaf.to_elements().collect())),
        );

        // --- account vault ------------------------------------------------------

        // populate Merkle store and advice map with nodes info needed to access vault assets
        self.extend_merkle_store(account.vault().inner_nodes());
        self.extend_map(
            account.vault().leaves().map(|leaf| (leaf.hash(), leaf.to_elements().collect())),
        );
    }

    /// Adds an account witness to the advice inputs.
    ///
    /// This involves extending the map to include the leaf's hash mapped to its elements, as well
    /// as extending the merkle store with the nodes of the witness.
    fn add_account_witness(&mut self, witness: &AccountWitness) {
        // populate advice map with the account's leaf
        let leaf = witness.leaf();
        self.add_map_entry(leaf.hash(), leaf.to_elements().collect());

        // extend the merkle store and map with account witnesses merkle path
        self.extend_merkle_store(witness.authenticated_nodes());
    }

    // NOTE INJECTION
    // --------------------------------------------------------------------------------------------

    /// Populates the advice inputs for all input notes.
    ///
    /// The advice provider is populated with:
    ///
    /// - For each note:
    ///     - The note's private arguments.
    ///     - The note's details (serial number, script root, and its storage / assets commitment).
    ///     - The preimages of the note recipient's hash chain.
    ///     - The note's public metadata (sender account ID, note type, note tag, attachment
    ///       schemes).
    ///     - The note's storage (unpadded).
    ///     - The note's assets (key and value words).
    ///     - For authenticated notes (determined by the `is_authenticated` flag):
    ///         - The note's authentication path against its block's note tree.
    ///         - The block number, sub commitment, note root.
    ///         - The note's position in the note tree
    ///
    /// The data above is processed by `prologue::process_input_notes_data`.
    fn add_input_notes(&mut self, tx_inputs: &TransactionInputs) {
        if tx_inputs.input_notes().is_empty() {
            return;
        }

        let mut note_data = Vec::new();
        for input_note in tx_inputs.input_notes().iter() {
            let note = input_note.note();
            let assets = note.assets();
            let recipient = note.recipient();
            let note_arg = tx_inputs.tx_args().get_note_args(note.id()).unwrap_or(&EMPTY_WORD);

            // recipient chain entries
            self.extend_map(recipient.to_advice_map_entries());
            // assets commitments
            self.add_map_entry(assets.commitment(), assets.to_elements());

            // ATTACHMENTS_COMMITMENT |-> [[ATTACHMENT_COMMITMENTS]]
            self.add_map_entry(
                note.attachments().to_commitment(),
                note.attachments()
                    .commitments()
                    .iter()
                    .flat_map(Word::as_elements)
                    .copied()
                    .collect(),
            );

            // ATTACHMENT_COMMITMENT |-> [ATTACHMENT_ELEMENTS] for each attachment
            for attachment in note.attachments().iter() {
                let commitment = attachment.content().to_commitment();
                let elements = attachment.content().to_elements();
                self.add_map_entry(commitment, elements);
            }

            // note metadata / details
            note_data.extend(*note_arg);
            note_data.extend(recipient.serial_num());
            note_data.extend(Word::from(recipient.script().root()));
            note_data.extend(*recipient.storage().commitment());
            note_data.extend(*assets.commitment());
            note_data.extend(note.metadata().to_metadata_word());
            note_data.extend(note.attachments().to_commitment());
            note_data.push(Felt::from(recipient.storage().num_items()));
            note_data.push(Felt::from(assets.num_assets() as u32));
            note_data.extend(assets.to_elements());

            // authentication vs unauthenticated
            match input_note {
                InputNote::Authenticated { note, proof } => {
                    // Push the `is_authenticated` flag
                    note_data.push(Felt::ONE);

                    // Merkle path
                    self.extend_merkle_store(proof.authenticated_nodes(note.id()));

                    let block_num = proof.location().block_num();
                    let block_header = if block_num == tx_inputs.block_header().block_num() {
                        tx_inputs.block_header()
                    } else {
                        tx_inputs
                            .blockchain()
                            .get_block(block_num)
                            .expect("block not found in partial blockchain")
                    };

                    note_data.push(block_num.into());
                    note_data.extend(block_header.sub_commitment());
                    note_data.extend(block_header.note_root());
                    note_data.push(Felt::from(proof.location().block_note_tree_index()));
                },
                InputNote::Unauthenticated { .. } => {
                    // push the `is_authenticated` flag
                    note_data.push(Felt::ZERO)
                },
            }
        }

        self.add_map_entry(tx_inputs.input_notes().commitment(), note_data);
    }

    // HELPER METHODS
    // --------------------------------------------------------------------------------------------

    /// Extends the map of values with the given argument, replacing previously inserted items.
    fn extend_map(&mut self, iter: impl IntoIterator<Item = (Word, Vec<Felt>)>) {
        self.0.extend(AdviceInputs::default().with_map(iter));
    }

    fn add_map_entry(&mut self, key: Word, values: Vec<Felt>) {
        self.0.extend(AdviceInputs::default().with_map([(key, values)]));
    }

    /// Extends the stack with the given elements.
    fn extend_stack(&mut self, iter: impl IntoIterator<Item = Felt>) {
        // `AdviceInputs` exposes its stack only as a typed `AdviceStack`, so appending goes
        // through `extend`, which appends the other instance's stack elements to ours.
        self.0.extend(AdviceInputs::default().with_stack(iter.into_iter().collect()));
    }

    /// Extends the [`MerkleStore`](crate::crypto::merkle::MerkleStore) with the given
    /// nodes.
    fn extend_merkle_store(&mut self, iter: impl Iterator<Item = InnerNodeInfo>) {
        self.0.extend(AdviceInputs::default().with_merkle_store(iter.collect()));
    }
}

// CONVERSIONS
// ================================================================================================

impl From<TransactionAdviceInputs> for AdviceInputs {
    fn from(wrapper: TransactionAdviceInputs) -> Self {
        wrapper.0
    }
}

impl From<AdviceInputs> for TransactionAdviceInputs {
    fn from(inner: AdviceInputs) -> Self {
        Self(inner)
    }
}