miden-prover 0.24.2

use std::sync::Arc;

use miden_assembly::Assembler;
use miden_debug_types::{Location, SourceFile, SourceSpan};
use miden_processor::{
    BaseHost, DefaultHost, ExecutionOptions, Felt, FutureMaybeSend, Host, LoadedMastForest,
    ProcessorState, Word,
    advice::AdviceMutation,
    event::{EventError, EventName},
};
use miden_prover::{AdviceInputs, ProvingOptions, StackInputs, prove, prove_sync};

struct YieldingAsyncHost {
    event_calls: usize,
}

impl YieldingAsyncHost {
    fn new() -> Self {
        Self { event_calls: 0 }
    }
}

impl BaseHost for YieldingAsyncHost {
    fn get_label_and_source_file(
        &self,
        _location: &Location,
    ) -> (SourceSpan, Option<Arc<SourceFile>>) {
        (SourceSpan::UNKNOWN, None)
    }
}

impl Host for YieldingAsyncHost {
    fn get_mast_forest(
        &self,
        _node_digest: &Word,
    ) -> impl FutureMaybeSend<Option<LoadedMastForest>> {
        async { None }
    }

    fn on_event(
        &mut self,
        _process: &ProcessorState<'_>,
    ) -> impl FutureMaybeSend<Result<Vec<AdviceMutation>, EventError>> {
        self.event_calls += 1;
        async {
            tokio::task::yield_now().await;
            Ok(Vec::new())
        }
    }
}

fn simple_program() -> miden_processor::Program {
    Assembler::default()
        .assemble_program(
            "program",
            r#"
            begin
                repeat.64
                    swap dup.1 add
                end
            end
            "#,
        )
        .expect("program should compile")
        .unwrap_program()
}

#[tokio::test(flavor = "current_thread")]
async fn prove_async_matches_prove() {
    let program = simple_program();
    let stack_inputs = StackInputs::new(&[Felt::new_unchecked(0), Felt::new_unchecked(1)]).unwrap();
    let advice_inputs = AdviceInputs::default();
    let execution_options = ExecutionOptions::default();
    let options = ProvingOptions::default();

    let mut sync_host = DefaultHost::default();
    let (sync_outputs, sync_proof) = prove_sync(
        &program,
        stack_inputs,
        advice_inputs.clone(),
        &mut sync_host,
        execution_options,
        options.clone(),
    )
    .unwrap();

    let mut async_host = DefaultHost::default();
    let (async_outputs, async_proof) = prove(
        &program,
        stack_inputs,
        advice_inputs,
        &mut async_host,
        execution_options,
        options,
    )
    .await
    .unwrap();

    assert_eq!(sync_outputs, async_outputs);
    assert_eq!(sync_proof.hash_fn(), async_proof.hash_fn());
    assert!(!sync_proof.stark_proof().is_empty());
    assert!(!async_proof.stark_proof().is_empty());
}

#[tokio::test(flavor = "current_thread")]
async fn prove_async_supports_async_only_host_events() {
    let event_name = EventName::new("test::async::prove");
    let event_id = event_name.to_event_id().as_u64();
    let program = Assembler::default()
        .assemble_program("program", format!("begin push.{event_id} emit drop end"))
        .expect("program should compile")
        .unwrap_program();

    let mut host = YieldingAsyncHost::new();
    let (_outputs, proof) = prove(
        &program,
        StackInputs::default(),
        AdviceInputs::default(),
        &mut host,
        ExecutionOptions::default(),
        ProvingOptions::default(),
    )
    .await
    .expect("async proving should succeed");

    assert_eq!(host.event_calls, 1);
    assert!(!proof.stark_proof().is_empty());
}