miden-prover 0.30.0

use alloc::string::ToString;

use miden_core::proof::{ExecutionProof, HashFunction, PrecompileProof, VmProof};
use miden_processor::{
    ExecutionError, ExecutionOptions, ExecutionWitness, FastProcessor, PrecompileWitness, Program,
    StackInputs, StackOutputs, SyncHost, VmWitness,
    advice::AdviceInputs,
    trace::{self, VmTrace, build_trace_with_budget},
};

use crate::{config, prove_stark};

/// A synchronous, configurable prover for post-execution Miden VM witnesses.
///
/// This type does not execute programs. It owns proof-generation policy — including the memory
/// budget for materializing a VM execution trace — and consumes witnesses produced by the
/// processor.
#[derive(Debug, Clone, Eq, PartialEq)]
pub struct Prover {
    hash_fn: HashFunction,
    max_prover_memory_bytes: u64,
}

impl Prover {
    /// Default maximum memory, in bytes, this prover is permitted to allocate for a proof over a
    /// VM execution trace.
    ///
    /// This bounds only the lifted-STARK Miden VM proof modelled by `miden_air::memory`; it does
    /// not cover the precompile prover's memory footprint.
    pub const DEFAULT_MAX_PROVER_MEMORY_BYTES: u64 = trace::DEFAULT_MAX_PROVER_MEMORY_BYTES;

    /// Creates a prover with the canonical proof-generation configuration.
    pub const fn new() -> Self {
        Self {
            hash_fn: HashFunction::Blake3_256,
            max_prover_memory_bytes: Self::DEFAULT_MAX_PROVER_MEMORY_BYTES,
        }
    }

    /// Sets the hash function used for proofs generated by this prover.
    #[must_use]
    pub const fn with_hash_fn(mut self, hash_fn: HashFunction) -> Self {
        self.hash_fn = hash_fn;
        self
    }

    /// Sets the maximum memory, in bytes, this prover is permitted to allocate for a proof over a
    /// VM execution trace.
    #[must_use]
    pub const fn with_max_prover_memory_bytes(mut self, max_prover_memory_bytes: u64) -> Self {
        self.max_prover_memory_bytes = max_prover_memory_bytes;
        self
    }

    /// Returns the maximum memory, in bytes, this prover is permitted to allocate for a proof
    /// over a VM execution trace.
    pub const fn max_prover_memory_bytes(&self) -> u64 {
        self.max_prover_memory_bytes
    }

    /// Proves only the VM portion of an execution witness.
    ///
    /// If the execution authenticated deferred precompile work, the returned proof carries its
    /// passive singleton wire for later hydration and proving. Otherwise, it is complete.
    pub fn prove(&self, witness: ExecutionWitness) -> Result<ExecutionProof, ProverError> {
        let (vm_witness, precompile_witness) = witness.into_parts();
        let vm = self.prove_vm(vm_witness)?;
        let Some(precompile_witness) = precompile_witness else {
            return Ok(ExecutionProof::Complete { vm, precompile: None });
        };
        let precompile = precompile_witness
            .state()
            .to_wire()
            .expect("execution witness state must have canonical deferred wire");
        Ok(ExecutionProof::Deferred { vm, precompile })
    }

    /// Proves a complete execution witness entirely in memory.
    ///
    /// Both VM and precompile proving consume the hydrated witness directly; the witness is not
    /// serialized on this local path.
    pub fn prove_full(&self, witness: ExecutionWitness) -> Result<ExecutionProof, ProverError> {
        let (vm_witness, precompile_witness) = witness.into_parts();
        let vm = self.prove_vm(vm_witness)?;
        let precompile = precompile_witness
            .as_ref()
            .map(|witness| self.prove_precompile(witness))
            .transpose()?;
        Ok(ExecutionProof::Complete { vm, precompile })
    }

    /// Materializes and proves the VM trace represented by `witness`.
    fn prove_vm(&self, witness: VmWitness) -> Result<VmProof, ProverError> {
        let trace = {
            let _span = tracing::info_span!("build_miden_vm_trace").entered();
            build_trace_with_budget(witness, self.max_prover_memory_bytes)
                .map_err(ProverError::TraceGeneration)?
        };

        self.prove_vm_trace(trace)
    }

    /// Proves one singleton or merged precompile witness without consuming its hydrated DAG.
    pub fn prove_precompile(
        &self,
        witness: &PrecompileWitness,
    ) -> Result<PrecompileProof, ProverError> {
        let proof = miden_precompiles_prover::prove_deferred_state(witness.state(), self.hash_fn)
            .map_err(ProverError::PrecompileProofGeneration)?;
        Ok(PrecompileProof { proof, roots: witness.roots().to_vec() })
    }

    #[cfg(feature = "std")]
    fn prove_full_trace(
        &self,
        trace: VmTrace,
        precompile: Option<&PrecompileWitness>,
    ) -> Result<ExecutionProof, ProverError> {
        let vm = self.prove_vm_trace(trace)?;
        let precompile = precompile.map(|witness| self.prove_precompile(witness)).transpose()?;
        Ok(ExecutionProof::Complete { vm, precompile })
    }

    /// Proves a fully materialized VM trace.
    ///
    /// Buffered and overlapped trace construction share this private implementation so STARK
    /// generation and VM proof packaging cannot diverge.
    #[tracing::instrument(name = "miden_vm", skip_all)]
    fn prove_vm_trace(&self, trace: VmTrace) -> Result<VmProof, ProverError> {
        let trace_len_summary = trace.trace_len_summary();
        let params = config::pcs_params();
        tracing::event!(
            tracing::Level::INFO,
            "Generated execution traces: core={}, range={}, chiplets={}, poseidon2={}, padded={}, \
             estimated_prover_memory_bytes={:?}",
            trace_len_summary.core_trace_len(),
            trace_len_summary.range_trace_len(),
            trace_len_summary.chiplets_trace_len().trace_len(),
            trace_len_summary.poseidon2_permutation_trace_len(),
            trace_len_summary.padded_trace_len(),
            trace_len_summary.prover_memory_bytes(&params)
        );

        let precompile_root = trace.precompile_root();
        let (public_values, aux_inputs) = trace.public_inputs().to_air_inputs();
        let (core_matrix, chiplets_matrix, poseidon2_matrix) = trace.into_air_matrices();

        let proof_bytes = match self.hash_fn {
            HashFunction::Blake3_256 => {
                let config = config::blake3_256_config(params, config::RELATION_DIGEST);
                prove_stark(
                    &config,
                    core_matrix,
                    chiplets_matrix,
                    poseidon2_matrix,
                    &public_values,
                    &aux_inputs,
                )
            },
            HashFunction::Keccak => {
                let config = config::keccak_config(params, config::RELATION_DIGEST);
                prove_stark(
                    &config,
                    core_matrix,
                    chiplets_matrix,
                    poseidon2_matrix,
                    &public_values,
                    &aux_inputs,
                )
            },
            HashFunction::Rpo256 => {
                let config = config::rpo_config(params, config::RELATION_DIGEST);
                prove_stark(
                    &config,
                    core_matrix,
                    chiplets_matrix,
                    poseidon2_matrix,
                    &public_values,
                    &aux_inputs,
                )
            },
            HashFunction::Poseidon2 => {
                let config = config::poseidon2_config(params, config::RELATION_DIGEST);
                prove_stark(
                    &config,
                    core_matrix,
                    chiplets_matrix,
                    poseidon2_matrix,
                    &public_values,
                    &aux_inputs,
                )
            },
            HashFunction::Rpx256 => {
                let config = config::rpx_config(params, config::RELATION_DIGEST);
                prove_stark(
                    &config,
                    core_matrix,
                    chiplets_matrix,
                    poseidon2_matrix,
                    &public_values,
                    &aux_inputs,
                )
            },
        }
        .map_err(ProverError::VmProofGeneration)?;

        let proof = miden_core::proof::StarkProof::new(proof_bytes, self.hash_fn);
        Ok(VmProof { proof, precompile_root })
    }
}

/// Executes and fully proves a program synchronously.
///
/// This FastProcessor-backed orchestration function preserves the optimized overlapped
/// execution/trace-building path. Proving policy belongs on [`Prover`].
///
/// When enabled in `execution_options`, the processor may build the hasher chiplet alongside
/// execution. A caller with no separate Rayon worker uses compact buffered replay. Both cases use
/// the same private VM STARK and complete-local packaging implementation.
#[tracing::instrument(name = "prove_program_sync", skip_all)]
pub fn prove_sync(
    prover: &Prover,
    program: &Program,
    stack_inputs: StackInputs,
    advice_inputs: AdviceInputs,
    host: &mut impl SyncHost,
    execution_options: ExecutionOptions,
) -> Result<(StackOutputs, ExecutionProof), ExecutionError> {
    #[cfg(feature = "std")]
    let overlapped_trace_build = execution_options.overlapped_trace_build();
    let processor = FastProcessor::new_with_options(stack_inputs, advice_inputs, execution_options)
        .map_err(ExecutionError::advice_error_no_context)?;

    #[cfg(feature = "std")]
    if overlapped_trace_build {
        let (trace, precompile) = {
            let _span = tracing::info_span!("execute_miden_vm").entered();
            processor.execute_and_build_trace_sync(
                program,
                host,
                prover.max_prover_memory_bytes(),
            )?
        };
        let stack_outputs = *trace.stack_outputs();
        let proof = prover
            .prove_full_trace(trace, precompile.as_ref())
            .map_err(ProverError::into_execution_error)?;
        return Ok((stack_outputs, proof));
    }

    let witness = {
        let _span = tracing::info_span!("execute_miden_vm").entered();
        processor.execute_for_proving_sync(program, host)?
    };
    let stack_outputs = *witness.claim().stack_outputs();
    let proof = prover.prove_full(witness).map_err(ProverError::into_execution_error)?;
    Ok((stack_outputs, proof))
}

impl Default for Prover {
    fn default() -> Self {
        Self::new()
    }
}

/// Errors produced while proving post-execution witnesses.
#[derive(Debug, thiserror::Error)]
#[non_exhaustive]
pub enum ProverError {
    /// The processor witness could not be materialized into a valid execution trace.
    #[error("failed to materialize VM execution trace: {0}")]
    TraceGeneration(#[source] ExecutionError),
    /// The materialized VM trace could not be proved.
    #[error("failed to prove VM execution trace: {0}")]
    VmProofGeneration(#[source] ExecutionError),
    /// The deferred precompile witness could not be proved.
    #[error("failed to prove precompile witness: {0}")]
    PrecompileProofGeneration(#[source] miden_precompiles_prover::ProveDeferredStateError),
}

impl ProverError {
    fn into_execution_error(self) -> ExecutionError {
        match self {
            Self::TraceGeneration(error) | Self::VmProofGeneration(error) => error,
            Self::PrecompileProofGeneration(error) => {
                ExecutionError::ProvingError(error.to_string())
            },
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn prover_uses_canonical_default_and_allows_hash_override() {
        let prover = Prover::new();
        assert_eq!(prover.hash_fn, HashFunction::Blake3_256);

        let prover = prover.with_hash_fn(HashFunction::Poseidon2);
        assert_eq!(prover.hash_fn, HashFunction::Poseidon2);
    }

    #[test]
    fn prover_uses_canonical_memory_budget_and_allows_override() {
        let prover = Prover::new();
        assert_eq!(prover.max_prover_memory_bytes(), Prover::DEFAULT_MAX_PROVER_MEMORY_BYTES);

        let prover = prover.with_max_prover_memory_bytes(1 << 20);
        assert_eq!(prover.max_prover_memory_bytes(), 1 << 20);
    }
}