miden-vm 0.22.1

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

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

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 .masl 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,

    /// Enable tracing to monitor execution of the VM
    #[arg(short = 't', long = "trace")]
    trace: bool,

    /// Disable debug instructions (release mode)
    #[arg(long = "release")]
    release: bool,

    /// 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_proof_options(&self) -> Result<ProvingOptions, Report> {
        let exec_options = ExecutionOptions::new(
            Some(self.max_cycles),
            self.expected_cycles,
            ExecutionOptions::DEFAULT_CORE_TRACE_FRAGMENT_SIZE,
            self.trace,
            !self.release,
        )
        .map_err(|err| Report::msg(format!("{err}")))?;

        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.with_execution_options(exec_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, mut host) = match ext.as_str() {
            "masp" => (get_masp_program(&self.program_file)?, host),
            "masm" => {
                let (program, source_manager) = get_masm_program(
                    &self.program_file,
                    &libraries,
                    !self.release,
                    self.kernel_file.as_deref(),
                )?;
                (program, host.with_source_manager(source_manager))
            },
            _ => 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 proving_options = self.get_proof_options()?;

        // execute program and generate proof
        let (stack_outputs, proof) = miden_prover::prove_sync(
            &program,
            stack_inputs,
            advice_inputs,
            &mut host,
            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(())
    }
}