miden-tx 0.17.0

use alloc::collections::BTreeMap;
use alloc::vec::Vec;

use miden_processor::ExecutionError;
use miden_processor::advice::AdviceInputs;
use miden_protocol::account::AccountId;
use miden_protocol::block::BlockNumber;
use miden_protocol::note::Note;
use miden_protocol::transaction::{
    InputNote,
    InputNotes,
    TransactionArgs,
    TransactionInputs,
    TransactionKernel,
};
use miden_standards::note::{NoteConsumptionStatus, StandardNote};

use super::{ProgramExecutor, TransactionExecutor};
use crate::auth::TransactionAuthenticator;
use crate::errors::TransactionCheckerError;
use crate::executor::map_execution_error;
use crate::{DataStore, NoteCheckerError, TransactionExecutorError};

mod checker_utils;

pub use checker_utils::{
    FailedNote,
    MAX_NUM_CHECKER_NOTES,
    NoteConsumptionInfo,
    NoteFailure,
    SuccessfulNote,
};
use checker_utils::{NoteBundle, handle_epilogue_error};

// NOTE CONSUMPTION CHECKER
// ================================================================================================

/// This struct performs input notes check against provided target account.
///
/// The check is performed using the [NoteConsumptionChecker::check_notes_consumability] procedure.
/// Essentially runs the transaction to make sure that provided input notes could be consumed by the
/// account.
pub struct NoteConsumptionChecker<'a, STORE, AUTH, EXEC: ProgramExecutor>(
    &'a TransactionExecutor<'a, 'a, STORE, AUTH, EXEC>,
);

impl<'a, STORE, AUTH, EXEC> NoteConsumptionChecker<'a, STORE, AUTH, EXEC>
where
    STORE: DataStore + Sync,
    AUTH: TransactionAuthenticator + Sync,
    EXEC: ProgramExecutor,
{
    /// Creates a new [`NoteConsumptionChecker`] instance with the given transaction executor.
    pub fn new(tx_executor: &'a TransactionExecutor<'a, 'a, STORE, AUTH, EXEC>) -> Self {
        NoteConsumptionChecker(tx_executor)
    }

    /// Checks whether some set of the provided input notes could be consumed by the provided
    /// account by executing the transaction with varying combination of notes.
    ///
    /// This function attempts to find the maximum set of notes that can be successfully executed
    /// together by the target account.
    ///
    /// Because of the runtime complexity involved in this function, a limited range of
    /// [`MAX_NUM_CHECKER_NOTES`] input notes is allowed.
    ///
    /// If some notes succeed and others fail, the failed notes are removed from the candidate set
    /// and the remaining notes (successful + unattempted) are retried in the next iteration. This
    /// process continues until either all remaining notes succeed or no notes can be successfully
    /// executed
    ///
    /// For example, given notes A, B, C, D, E, the execution flow would be as follows:
    /// - Try [A, B, C, D, E] → A, B succeed, C fails → Remove C, try again.
    /// - Try [A, B, D, E] → A, B, D succeed, E fails → Remove E, try again.
    /// - Try [A, B, D] → All succeed → Return successful=[A, B, D], failed=[C, E].
    ///
    /// If a failure occurs at the epilogue phase of the transaction execution, the relevant set of
    /// otherwise-successful notes are retried in various combinations in an attempt to find a
    /// combination that passes the epilogue phase successfully. Notes that are only consumable
    /// together, such as a feature note and the FEE_SPONSORSHIP notes bound to it, are grouped and
    /// retried as a unit.
    ///
    /// Returns a list of successfully consumed notes and a list of failed notes.
    pub async fn check_notes_consumability(
        &self,
        target_account_id: AccountId,
        block_ref: BlockNumber,
        mut notes: Vec<Note>,
        tx_args: TransactionArgs,
    ) -> Result<NoteConsumptionInfo, NoteCheckerError> {
        let num_notes = notes.len();
        if num_notes == 0 || num_notes > MAX_NUM_CHECKER_NOTES {
            return Err(NoteCheckerError::InputNoteCountOutOfRange(num_notes));
        }
        // Ensure standard notes are ordered first.
        notes.sort_unstable_by_key(|note| {
            StandardNote::from_script_root(note.script().root()).is_none()
        });

        let notes = InputNotes::from(notes);
        let tx_inputs = self
            .0
            .prepare_tx_inputs(target_account_id, block_ref, notes, tx_args)
            .await
            .map_err(NoteCheckerError::TransactionPreparation)?;

        // Attempt to find an executable set of notes.
        self.find_executable_notes_by_elimination(tx_inputs).await
    }

    /// Checks whether the provided input note could be consumed by the provided account by
    /// executing a transaction at the specified block height.
    ///
    /// This function takes into account the possibility that the signatures may not be loaded into
    /// the transaction context and returns the [`NoteConsumptionStatus`] result accordingly.
    ///
    /// This function first applies the static analysis of the provided note, and if it doesn't
    /// reveal any errors next it tries to execute the transaction. Based on the execution result,
    /// it either returns a [`NoteCheckerError`] or the [`NoteConsumptionStatus`]: depending on
    /// whether the execution succeeded, failed in the prologue, during the note execution process
    /// or in the epilogue.
    pub async fn can_consume(
        &self,
        target_account_id: AccountId,
        block_ref: BlockNumber,
        note: InputNote,
        tx_args: TransactionArgs,
    ) -> Result<NoteConsumptionStatus, NoteCheckerError> {
        // Return the consumption status if we manage to determine it from the standard note
        if let Some(standard_note) = StandardNote::from_script_root(note.note().script().root())
            && let Some(consumption_status) =
                standard_note.is_consumable(note.note(), target_account_id, block_ref)
        {
            return Ok(consumption_status);
        }

        // Prepare transaction inputs.
        let mut tx_inputs = self
            .0
            .prepare_tx_inputs(
                target_account_id,
                block_ref,
                InputNotes::new_unchecked(vec![note]),
                tx_args,
            )
            .await
            .map_err(NoteCheckerError::TransactionPreparation)?;

        // try to consume the provided note
        match self.try_execute_notes(&mut tx_inputs).await {
            // execution succeeded
            Ok(_cycle_counts) => Ok(NoteConsumptionStatus::Consumable),
            Err(tx_checker_error) => {
                match tx_checker_error {
                    // execution failed on the preparation stage, before we actually executed the tx
                    TransactionCheckerError::TransactionPreparation(e) => {
                        Err(NoteCheckerError::TransactionPreparation(e))
                    },
                    // execution failed during the prologue
                    TransactionCheckerError::PrologueExecution(e) => {
                        Err(NoteCheckerError::PrologueExecution(e))
                    },
                    // execution failed during the note processing
                    TransactionCheckerError::NoteExecution { .. } => {
                        Ok(NoteConsumptionStatus::UnconsumableConditions)
                    },
                    // execution failed during the epilogue
                    TransactionCheckerError::EpilogueExecution {
                        error: epilogue_error, ..
                    } => Ok(handle_epilogue_error(epilogue_error)),
                }
            },
        }
    }

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

    /// Finds a set of executable notes and eliminates failed notes from the list in the process.
    ///
    /// The result contains some combination of the input notes partitioned by whether they
    /// succeeded or failed to execute.
    async fn find_executable_notes_by_elimination(
        &self,
        mut tx_inputs: TransactionInputs,
    ) -> Result<NoteConsumptionInfo, NoteCheckerError> {
        let mut candidate_notes = tx_inputs
            .input_notes()
            .iter()
            .map(|note| note.clone().into_note())
            .collect::<Vec<_>>();
        let mut failed_notes = Vec::new();

        // Attempt to execute notes in a loop. Reduce the set of notes based on failures until
        // either a set of notes executes without failure or the set of notes cannot be
        // further reduced.
        loop {
            // Execute the candidate notes.
            tx_inputs.set_input_notes(candidate_notes.clone());
            match self.try_execute_notes(&mut tx_inputs).await {
                Ok(cycle_counts) => {
                    // A full set of successful notes has been found.
                    let successful = candidate_notes
                        .into_iter()
                        .zip(cycle_counts)
                        .map(|(note, num_cycles)| SuccessfulNote::new(note, num_cycles))
                        .collect();
                    return Ok(NoteConsumptionInfo::new(successful, failed_notes));
                },
                Err(TransactionCheckerError::NoteExecution {
                    failed_note_index,
                    error,
                    failed_note_cycle_count,
                    ..
                }) => {
                    // SAFETY: Failed note index is in bounds of the candidate notes.
                    let failed_note = candidate_notes.remove(failed_note_index);
                    failed_notes.push(FailedNote::new(
                        failed_note,
                        NoteFailure::Blamed {
                            error,
                            num_cycles: failed_note_cycle_count,
                        },
                    ));

                    // All possible candidate combinations have been attempted.
                    if candidate_notes.is_empty() {
                        return Ok(NoteConsumptionInfo::new(Vec::new(), failed_notes));
                    }
                    // Continue and process the next set of candidates.
                },
                Err(TransactionCheckerError::EpilogueExecution { .. }) => {
                    let consumption_info = self
                        .find_largest_executable_combination(
                            candidate_notes,
                            failed_notes,
                            tx_inputs,
                        )
                        .await;
                    return Ok(consumption_info);
                },
                Err(TransactionCheckerError::PrologueExecution(err)) => {
                    return Err(NoteCheckerError::PrologueExecution(err));
                },
                Err(TransactionCheckerError::TransactionPreparation(err)) => {
                    return Err(NoteCheckerError::TransactionPreparation(err));
                },
            }
        }
    }

    /// Attempts to find the largest possible combination of notes that can execute successfully
    /// together.
    ///
    /// The notes are first grouped into [`NoteBundle`]s, and the search grows a known-good set one
    /// bundle at a time: each round appends every remaining bundle to the accepted set in turn and
    /// keeps the first bundle that lets the whole set pass, until a round adds nothing.
    async fn find_largest_executable_combination(
        &self,
        remaining_notes: Vec<Note>,
        mut failed_notes: Vec<FailedNote>,
        mut tx_inputs: TransactionInputs,
    ) -> NoteConsumptionInfo {
        let mut remaining_bundles = NoteBundle::group(remaining_notes);
        let mut successful_notes: Vec<Note> = Vec::new();
        let mut successful_cycle_counts = Vec::new();
        let mut failed_note_index = BTreeMap::new();

        // Grow the accepted set until a full pass over the remaining bundles adds nothing, at which
        // point no bundle can extend it and the set is as large as this search can make it.
        loop {
            let mut extended = false;

            for idx in 0..remaining_bundles.len() {
                let bundle_notes = remaining_bundles[idx].notes().to_vec();
                let candidate_notes: Vec<Note> =
                    successful_notes.iter().chain(&bundle_notes).cloned().collect();

                tx_inputs.set_input_notes(candidate_notes.clone());
                match self.try_execute_notes(&mut tx_inputs).await {
                    Ok(cycle_counts) => {
                        // The notes just added might have failed earlier, either on their own or
                        // as part of another candidate set. Remove them from the failed list.
                        for note in bundle_notes {
                            failed_note_index.remove(&note.id());
                        }
                        // Store the cycle counts from the latest successful execution.
                        successful_cycle_counts = cycle_counts;
                        // This combination succeeded; commit it and drop the bundle from the
                        // remaining set.
                        successful_notes = candidate_notes;
                        remaining_bundles.remove(idx);
                        extended = true;
                        break;
                    },
                    Err(error) => {
                        // This combination failed, so the whole bundle is rejected. Blame the note
                        // the executor pointed at, when it pointed at one of the bundle's notes;
                        // an epilogue failure blames no particular note, so it falls to the note
                        // heading the bundle, the one the rest of the bundle is bound to.
                        let (blamed_idx, num_cycles) = match &error {
                            TransactionCheckerError::NoteExecution {
                                failed_note_index,
                                failed_note_cycle_count,
                                ..
                            } => (
                                failed_note_index
                                    .checked_sub(successful_notes.len())
                                    .filter(|idx| *idx < bundle_notes.len())
                                    .unwrap_or(0),
                                *failed_note_cycle_count,
                            ),
                            _ => (0, None),
                        };

                        let blamed_note = bundle_notes[blamed_idx].clone();
                        let blamed_id = blamed_note.id();

                        // Record every note of the bundle (overwriting previous failures for the
                        // relevant notes), so the reported notes always account for all inputs.
                        // Only the blamed note owns the error; the rest failed with its bundle.
                        failed_note_index.insert(
                            blamed_id,
                            FailedNote::new(
                                blamed_note,
                                NoteFailure::Blamed { error: error.into(), num_cycles },
                            ),
                        );
                        for (note_idx, note) in bundle_notes.iter().enumerate() {
                            if note_idx != blamed_idx {
                                failed_note_index.insert(
                                    note.id(),
                                    FailedNote::new(
                                        note.clone(),
                                        NoteFailure::Collateral { blamed_by: blamed_id },
                                    ),
                                );
                            }
                        }
                    },
                }
            }

            if !extended {
                break;
            }
        }

        // Pair successful notes with their cycle counts from the last successful execution.
        let successful = successful_notes
            .into_iter()
            .zip(successful_cycle_counts)
            .map(|(note, num_cycles)| SuccessfulNote::new(note, num_cycles))
            .collect();

        // Append failed notes to the list of failed notes provided as input.
        failed_notes.extend(failed_note_index.into_values());
        NoteConsumptionInfo::new(successful, failed_notes)
    }

    /// Attempts to execute a transaction with the provided input notes.
    ///
    /// This method executes the full transaction pipeline including prologue, note execution,
    /// and epilogue phases. It returns `Ok(cycle_counts)` if all notes are successfully consumed
    /// (where `cycle_counts` contains the number of cycles for each note), or a specific
    /// [`TransactionCheckerError`] indicating where and why the execution failed. The order of the
    /// returned `cycle_counts` is guaranteed to match the order of the input notes.
    async fn try_execute_notes(
        &self,
        tx_inputs: &mut TransactionInputs,
    ) -> Result<Vec<usize>, TransactionCheckerError> {
        if tx_inputs.input_notes().is_empty() {
            return Ok(Vec::new());
        }

        let (mut host, stack_inputs, advice_inputs) =
            self.0
                .prepare_transaction(tx_inputs)
                .await
                .map_err(TransactionCheckerError::TransactionPreparation)?;

        let program = TransactionKernel::main();
        let kernel_debug_info = TransactionKernel::main_debug_info();
        let executor = EXEC::new(stack_inputs, advice_inputs, self.0.exec_options)
            .map_err(ExecutionError::advice_error_no_context)
            .map_err(map_execution_error)
            .map_err(TransactionCheckerError::PrologueExecution)?;
        let result = executor
            .with_debug_info(kernel_debug_info.as_deref().cloned().unwrap_or_default())
            .with_entrypoint_source_node(TransactionKernel::main_entrypoint_source_node())
            .execute(&program, &mut host)
            .await
            .map_err(map_execution_error);

        match result {
            Ok(execution_output) => {
                let cycle_counts = host
                    .tx_progress()
                    .note_execution()
                    .iter()
                    .map(|(_, interval)| interval.len())
                    .collect();

                // Set the advice inputs from the successful execution as advice inputs for
                // reexecution. This avoids calls to the data store (to load data lazily) that have
                // already been done as part of this execution.
                let (_, advice_map, merkle_store) = execution_output.advice.into_parts();
                let advice_inputs = AdviceInputs::from(advice_map).with_merkle_store(merkle_store);
                tx_inputs.set_advice_inputs(advice_inputs);
                Ok(cycle_counts)
            },
            Err(error) => {
                let notes = host.tx_progress().note_execution();

                // Empty notes vector means that we didn't process the notes, so an error
                // occurred.
                if notes.is_empty() {
                    return Err(TransactionCheckerError::PrologueExecution(error));
                }

                let ((_, last_note_interval), success_notes) =
                    notes.split_last().expect("notes vector is not empty because of earlier check");

                // If the interval end of the last note is specified, then an error occurred after
                // notes processing. All notes executed successfully in this case.
                if last_note_interval.end().is_some() {
                    let successful_notes_cycle_counts =
                        notes.iter().map(|(_, interval)| interval.len()).collect();
                    Err(TransactionCheckerError::EpilogueExecution {
                        error,
                        successful_notes_cycle_counts,
                    })
                } else {
                    // Return the index of the failed note.
                    let failed_note_index = success_notes.len();
                    let successful_notes_cycle_counts =
                        success_notes.iter().map(|(_, interval)| interval.len()).collect();

                    // Compute the failed note's cycle count when the failure was due to
                    // exceeding the cycle limit. In this case, the note's interval has a
                    // start but no end, and the total cycles consumed equals the max allowed.
                    let failed_note_cycle_count = match &error {
                        TransactionExecutorError::TransactionProgramExecutionFailed(
                            ExecutionError::CycleLimitExceeded(max_cycles),
                        ) => last_note_interval
                            .start()
                            .map(|start| *max_cycles as usize - usize::from(start)),
                        _ => None,
                    };

                    Err(TransactionCheckerError::NoteExecution {
                        failed_note_index,
                        error,
                        successful_notes_cycle_counts,
                        failed_note_cycle_count,
                    })
                }
            },
        }
    }
}