miden-vm 0.24.0

use std::{path::PathBuf, time::Instant};

use clap::Parser;
use miden_assembly::diagnostics::{IntoDiagnostic, Report, WrapErr};
use miden_core_lib::CoreLibrary;
use miden_processor::{DefaultHost, ExecutionOptions, FastProcessor};
use miden_vm::{
    HashFunction, ProvingOptions, TraceProvingInputs, internal::InputFile, prove_from_trace_sync,
};

use super::{
    data::{Libraries, OutputFile, ProofFile},
    utils::{get_masm_program, get_masp_program},
};

#[derive(Debug, Clone, Parser)]
#[command(about = "Prove a Miden program")]
pub struct ProveCmd {
    /// Path to a .masm assembly file or a .masp package file
    #[arg(value_parser)]
    program_file: PathBuf,

    /// Number of cycles the program is expected to consume
    #[arg(short = 'e', long = "exp-cycles", default_value = "64")]
    expected_cycles: u32,

    /// Path to input file
    #[arg(short = 'i', long = "input", value_parser)]
    input_file: Option<PathBuf>,

    /// Paths to .masp library files
    #[arg(short = 'l', long = "libraries", value_parser)]
    library_paths: Vec<PathBuf>,

    /// Maximum number of cycles a program is allowed to consume
    #[arg(short = 'm', long = "max-cycles", default_value_t = ExecutionOptions::MAX_CYCLES)]
    max_cycles: u32,

    /// Number of outputs
    #[arg(short = 'n', long = "num-outputs", default_value = "16")]
    num_outputs: usize,

    /// Path to output file
    #[arg(short = 'o', long = "output", value_parser)]
    output_file: Option<PathBuf>,

    /// Path to proof file
    #[arg(short = 'p', long = "proof", value_parser)]
    proof_file: Option<PathBuf>,

    /// Specifies the hash function to be used
    /// Valid options: blake3-256, rpo, rpx, poseidon2
    #[arg(long = "hasher", default_value = "blake3-256")]
    hasher: String,

    /// Security level for execution proofs generated by the VM
    #[arg(short = 's', long = "security", default_value = "96bits")]
    security: String,

    /// Path to a file (.masm or .masp) containing the kernel to be loaded with the program
    #[arg(long = "kernel", value_parser)]
    kernel_file: Option<PathBuf>,
}

impl ProveCmd {
    pub fn get_execution_options(&self) -> Result<ExecutionOptions, Report> {
        ExecutionOptions::new(
            Some(self.max_cycles),
            self.expected_cycles,
            ExecutionOptions::DEFAULT_CORE_TRACE_FRAGMENT_SIZE,
        )
        .map_err(|err| Report::msg(format!("{err}")))
    }

    pub fn get_proof_options(&self) -> Result<ProvingOptions, Report> {
        let hash_fn = HashFunction::try_from(self.hasher.as_str())
            .map_err(|err| Report::msg(format!("{err}")))?;
        let proving_options = match self.security.as_str() {
            "96bits" => ProvingOptions::with_96_bit_security(hash_fn),
            other => {
                return Err(Report::msg(format!(
                    "{other} is not a valid security setting. Currently only '96bits' is supported."
                )));
            },
        };
        Ok(proving_options)
    }
    pub fn execute(&self) -> Result<(), Report> {
        println!("===============================================================================");
        println!("Prove program: {}", self.program_file.display());
        println!("-------------------------------------------------------------------------------");

        // load libraries from files
        let libraries = Libraries::new(&self.library_paths)?;

        // validate kernel file if provided
        if let Some(ref kernel_path) = self.kernel_file
            && !kernel_path.is_file()
        {
            return Err(Report::msg(format!(
                "Kernel file `{}` must be a file.",
                kernel_path.display()
            )));
        }

        // determine file type based on extension
        let ext = self
            .program_file
            .extension()
            .and_then(|s| s.to_str())
            .unwrap_or("")
            .to_lowercase();

        let input_data = InputFile::read(&self.input_file, &self.program_file)?;

        let host = DefaultHost::default().with_library(&CoreLibrary::default())?;
        // Use a single match expression to load the program.
        let (program, package_debug_info, entrypoint_source_node, mut host) = match ext.as_str() {
            "masp" => (get_masp_program(&self.program_file)?, None, None, host),
            "masm" => {
                let (program, package_debug_info, entrypoint_source_node, source_manager) =
                    get_masm_program(&self.program_file, &libraries, self.kernel_file.as_deref())?;
                let mut host = host.with_source_manager(source_manager);
                for library in libraries.libraries.iter().cloned() {
                    host.load_library(library)
                        .into_diagnostic()
                        .wrap_err("Failed to load library")?;
                }
                (program, package_debug_info, entrypoint_source_node, host)
            },
            _ => return Err(Report::msg("The provided file must have a .masm or .masp extension")),
        };

        let program_hash: [u8; 32] = program.hash().into();
        println!("Proving program with hash {}...", hex::encode(program_hash));
        let now = Instant::now();

        // fetch the stack and program inputs from the arguments
        let stack_inputs = input_data.parse_stack_inputs().map_err(Report::msg)?;
        let advice_inputs = input_data.parse_advice_inputs().map_err(Report::msg)?;

        let execution_options = self.get_execution_options()?;
        let proving_options = self.get_proof_options()?;

        // execute program and generate proof
        let processor =
            FastProcessor::new_with_options(stack_inputs, advice_inputs, execution_options)
                .map_err(|err| Report::msg(format!("{err}")))?;
        let trace_inputs = match (package_debug_info.as_ref(), entrypoint_source_node) {
            (Some(debug_info), Some(entrypoint_source_node_id)) => processor
                .execute_trace_inputs_with_package_debug_info_at_source_node_sync(
                    &program,
                    debug_info,
                    entrypoint_source_node_id,
                    &mut host,
                )
                .wrap_err("Failed to execute program")?,
            (Some(debug_info), None) => processor
                .execute_trace_inputs_with_package_debug_info_sync(&program, debug_info, &mut host)
                .wrap_err("Failed to execute program")?,
            (None, _) => processor
                .execute_trace_inputs_sync(&program, &mut host)
                .wrap_err("Failed to execute program")?,
        };
        let (stack_outputs, proof) =
            prove_from_trace_sync(TraceProvingInputs::new(trace_inputs, proving_options))
                .wrap_err("Failed to prove program")?;

        println!("Program proved in {} ms", now.elapsed().as_millis());

        // write proof to file
        ProofFile::write(proof, &self.proof_file, &self.program_file).map_err(Report::msg)?;

        // provide outputs
        if let Some(output_path) = &self.output_file {
            // write all outputs to specified file.
            OutputFile::write(&stack_outputs, output_path).map_err(Report::msg)?;
        } else {
            // if no output path was provided, get the stack outputs for printing to the screen.
            let stack = stack_outputs.get_num_elements(self.num_outputs).to_vec();

            // write all outputs to default location if none was provided
            let default_output_path = self.program_file.with_extension("outputs");
            OutputFile::write(&stack_outputs, &default_output_path).map_err(Report::msg)?;

            // print stack outputs to screen.
            println!("Output: {stack:?}");
        }

        Ok(())
    }
}