zisk-sdk 1.1.0-alpha

use super::EmbeddedClient;
use crate::embedded::{validate_embedded_request, EmbeddedProver, HintsKind, StdinKind};
use crate::hints::HintsSource;
use crate::input_source::InputSource;
use crate::job_handle::{fire_event, fire_result_event, JobHandle, SubscriberList};
use crate::prove::ProveResult;
use crate::{ExecutorKind, JobEvent, Result, SdkError};
use std::sync::Arc;
use std::time::Duration;
use zisk_common::io::StreamSource;
use zisk_common::ProofKind;
use zisk_prover_backend::GuestProgram;

impl EmbeddedClient {
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn do_prove(
        &self,
        program: &GuestProgram,
        stdin: InputSource,
        hints: Option<HintsSource>,
        executor: ExecutorKind,
        proof_kind: ProofKind,
        timeout: Option<Duration>,
        subs: SubscriberList,
    ) -> Result<JobHandle<ProveResult>> {
        let program = program.clone();
        let subs_cloned = Arc::clone(&subs);
        let prover = self.prover.clone();

        let handle = tokio::task::spawn_blocking(move || {
            fire_event(&subs_cloned, JobEvent::Started);

            let result = Self::do_prove_inner(prover, &program, stdin, hints, executor, proof_kind);

            fire_result_event(&subs_cloned, &result);

            result
        });

        Ok(JobHandle::new_embedded(handle, subs, timeout))
    }

    /// Run proof generation synchronously on the calling thread.
    ///
    /// Unlike [`do_prove`](Self::do_prove), this performs no `spawn_blocking`
    /// and returns the result directly, so it requires no async runtime.
    /// Registered event callbacks fire synchronously during the call.
    pub(crate) fn do_prove_sync(
        &self,
        program: &GuestProgram,
        stdin: InputSource,
        hints: Option<HintsSource>,
        executor: ExecutorKind,
        proof_kind: ProofKind,
        subs: SubscriberList,
    ) -> Result<ProveResult> {
        fire_event(&subs, JobEvent::Started);
        let result =
            Self::do_prove_inner(self.prover.clone(), program, stdin, hints, executor, proof_kind);
        fire_result_event(&subs, &result);
        result
    }

    fn do_prove_inner(
        prover: Arc<EmbeddedProver>,
        program: &GuestProgram,
        stdin: InputSource,
        hints: Option<HintsSource>,
        executor: ExecutorKind,
        proof_kind: ProofKind,
    ) -> Result<ProveResult> {
        macro_rules! apply_mode {
            ($builder:expr) => {
                match proof_kind {
                    ProofKind::VadcopFinal => $builder,
                    ProofKind::VadcopFinalMinimal => {
                        $builder.wrap_proof(ProofKind::VadcopFinalMinimal)
                    }
                    ProofKind::Plonk => $builder.wrap_proof(ProofKind::Plonk),
                }
            };
        }
        let (prover_is_asm, was_setup_with_hints) = match prover.as_ref() {
            EmbeddedProver::Asm(p) => (true, p.was_setup_with_hints()),
            EmbeddedProver::Emu(_) => (false, false),
        };
        validate_embedded_request(
            prover_is_asm,
            executor,
            hints.as_ref().map(HintsKind::of),
            StdinKind::of(&stdin),
            was_setup_with_hints,
        )?;

        // Inputs are validated above; the dispatch only routes valid requests.
        let result = match (prover.as_ref(), executor) {
            (EmbeddedProver::Emu(p), ExecutorKind::Emulator) => {
                let InputSource::Stdin(s) = stdin else { unreachable!() };
                apply_mode!(p.prove(program, s.into_inner())).run().map_err(SdkError::backend)?
            }
            (EmbeddedProver::Asm(p), ExecutorKind::Emulator) => {
                let InputSource::Stdin(s) = stdin else { unreachable!() };
                p.prove_emulator(program, s.into_inner(), proof_kind).map_err(SdkError::backend)?
            }
            (EmbeddedProver::Asm(p), ExecutorKind::Assembly) => {
                if let Some(hints) = hints {
                    match hints {
                        HintsSource::Hints(h) => {
                            p.register_hints_stream(h.into_inner()).map_err(SdkError::backend)?;
                        }
                        HintsSource::Stream(stream) => {
                            stream.start()?;
                            let uri = stream.uri().to_string();
                            let source = StreamSource::from_uri(&uri).map_err(SdkError::backend)?;
                            p.register_hints_stream(source).map_err(SdkError::backend)?;
                        }
                    }
                    apply_mode!(p.prove(program, zisk_common::io::ZiskStdin::new()))
                        .run()
                        .map_err(SdkError::backend)?
                } else {
                    match stdin {
                        InputSource::Stream(stream) => {
                            stream.start()?;
                            let uri = stream.uri().to_string();
                            let source = StreamSource::from_uri(&uri).map_err(SdkError::backend)?;
                            p.register_inputs_stream(source).map_err(SdkError::backend)?;
                            apply_mode!(p.prove(program, zisk_common::io::ZiskStdin::new()))
                                .run()
                                .map_err(SdkError::backend)?
                        }
                        InputSource::Stdin(s) => apply_mode!(p.prove(program, s.into_inner()))
                            .run()
                            .map_err(SdkError::backend)?,
                    }
                }
            }
            (EmbeddedProver::Emu(_), ExecutorKind::Assembly) => {
                unreachable!("rejected by validate_embedded_request")
            }
        };
        Ok(ProveResult::from(result))
    }
}