miden-prover 0.35.0

use alloc::{string::ToString, vec, vec::Vec};

use miden_core::proof::{ExecutionProof, HashFunction, PrecompileProof, PrecompileStatus, 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,
    max_precompile_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`; the
    /// precompile prover's memory footprint is bounded separately by
    /// [`max_precompile_prover_memory_bytes`](Self::max_precompile_prover_memory_bytes).
    pub const DEFAULT_MAX_PROVER_MEMORY_BYTES: u64 = trace::DEFAULT_MAX_PROVER_MEMORY_BYTES;

    /// Default maximum memory, in bytes, this prover is permitted to allocate for a proof over a
    /// precompile witness.
    pub const DEFAULT_MAX_PRECOMPILE_PROVER_MEMORY_BYTES: u64 =
        miden_precompiles_prover::DEFAULT_MAX_PRECOMPILE_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,
            max_precompile_prover_memory_bytes: Self::DEFAULT_MAX_PRECOMPILE_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
    }

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

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

    /// Proves only the VM portion of an execution witness.
    ///
    /// If the execution authenticated deferred precompile work, the returned proof carries its
    /// portable singleton witness for later 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::new(vm, PrecompileStatus::Empty));
        };
        Ok(ExecutionProof::new(vm, PrecompileStatus::Deferred(precompile_witness)))
    }

    /// Proves a complete execution witness entirely in memory.
    ///
    /// VM replay and direct precompile import consume the in-memory witness without serialization.
    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
            .map(|witness| self.prove_precompiles(vec![witness]))
            .transpose()?;
        let precompile = match precompile {
            Some(precompile) => PrecompileStatus::Proven(precompile),
            None => PrecompileStatus::Empty,
        };
        Ok(ExecutionProof::new(vm, precompile))
    }

    /// Proves a VM witness that does not authenticate deferred precompile work.
    ///
    /// The returned execution proof has an empty precompile status. Use [`Self::prove`] or
    /// [`Self::prove_full`] when the original execution witness contains precompile work.
    pub fn prove_vm_witness(&self, witness: VmWitness) -> Result<ExecutionProof, ProverError> {
        if witness.has_precompiles() {
            return Err(ProverError::VmWitnessHasPrecompiles);
        }

        let vm = self.prove_vm(witness)?;
        Ok(ExecutionProof::new(vm, PrecompileStatus::Empty))
    }

    /// 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 an owned batch of singleton execution obligations in one STARK.
    ///
    /// The proof preserves the input roots in order, including repeated roots. An empty batch
    /// is rejected. Single-execution proving uses this same path with a one-element vector.
    ///
    /// Batch-wide limits are checked during import, before STARK generation. Individually valid
    /// witnesses may exceed these limits when combined. Input accounting counts each supplied
    /// occurrence before sharing computations, including repeated data across witnesses.
    pub fn prove_precompiles(
        &self,
        witnesses: Vec<PrecompileWitness>,
    ) -> Result<PrecompileProof, ProverError> {
        miden_precompiles_prover::prove_precompiles_with_budget(
            witnesses,
            self.hash_fn,
            self.max_precompile_prover_memory_bytes,
        )
        .map_err(ProverError::PrecompileProofGeneration)
    }

    #[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_precompiles(vec![witness])).transpose()?;
        let precompile = match precompile {
            Some(precompile) => PrecompileStatus::Proven(precompile),
            None => PrecompileStatus::Empty,
        };
        Ok(ExecutionProof::new(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)
            .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 VM witness authenticates deferred precompile work that this proving path cannot carry.
    #[error("VM witness contains deferred precompile work")]
    VmWitnessHasPrecompiles,
    /// 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::PrecompileProvingError),
}

impl ProverError {
    fn into_execution_error(self) -> ExecutionError {
        match self {
            Self::VmWitnessHasPrecompiles => ExecutionError::ProvingError(self.to_string()),
            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);
    }

    #[test]
    fn prover_uses_canonical_precompile_memory_budget_and_allows_override() {
        let prover = Prover::new();
        assert_eq!(
            prover.max_precompile_prover_memory_bytes(),
            Prover::DEFAULT_MAX_PRECOMPILE_PROVER_MEMORY_BYTES
        );

        let prover = prover.with_max_precompile_prover_memory_bytes(1 << 20);
        assert_eq!(prover.max_precompile_prover_memory_bytes(), 1 << 20);
    }

    /// A minimal non-`TRUE` precompile witness: no keccak/uint/EC ops, just the fixed session
    /// installs plus one trivial `AND` node. Its chiplet traces still have nonzero (fixed-minimum)
    /// heights, so it is enough to exercise the memory-budget check without a full precompile
    /// execution fixture.
    fn trivial_precompile_witness() -> PrecompileWitness {
        use alloc::sync::Arc;

        use miden_core::deferred::{DeferredState, Node, PrecompileRegistry, TRUE_DIGEST};

        let registry = Arc::new(PrecompileRegistry::new());
        let mut state =
            DeferredState::new(registry).expect("empty registry state should initialize");
        let statement = state
            .register(Node::and(TRUE_DIGEST, TRUE_DIGEST))
            .expect("trivial AND node should register");
        state
            .log_statement(statement)
            .expect("trivial statement should log into the deferred root");
        state
            .into_witness()
            .expect("trivial deferred state should export")
            .expect("non-TRUE root should export a witness")
    }

    #[test]
    fn prove_precompiles_applies_the_configured_memory_budget() {
        // A 1-byte budget must reject even the minimal chiplet trace shapes, and the returned
        // error carries the configured budget. Exact-boundary behavior is covered by
        // `miden-precompiles-prover`.
        let err = Prover::new()
            .with_max_precompile_prover_memory_bytes(1)
            .prove_precompiles(vec![trivial_precompile_witness()])
            .expect_err("a 1-byte budget must reject even the minimal chiplet trace shapes");
        assert!(
            matches!(
                err,
                ProverError::PrecompileProofGeneration(
                    miden_precompiles_prover::PrecompileProvingError::MemoryBudgetExceeded {
                        budget_bytes: 1,
                        ..
                    }
                )
            ),
            "expected MemoryBudgetExceeded {{ budget_bytes: 1, .. }}, got: {err:?}"
        );
    }
}