curvy-prover 0.1.0-rc.3

use std::env;
use std::error::Error;
use std::ffi::OsString;
use std::fs;
use std::path::PathBuf;
use std::time::{Duration, Instant};

use curvy_prover::CircuitProver;

const DEFAULT_THREADS: usize = 1;
const MAX_THREADS: usize = 64;

struct Arguments {
    zkey_path: PathBuf,
    zkey_sha256: String,
    graph_path: PathBuf,
    graph_sha256: String,
    input_path: PathBuf,
    proof_path: PathBuf,
    public_path: PathBuf,
}

fn main() -> Result<(), Box<dyn Error>> {
    let arguments = parse_arguments()?;
    let rayon_threads = configure_rayon()?;

    let load_started = Instant::now();
    let zkey = fs::read(&arguments.zkey_path)?;
    let graph = fs::read(&arguments.graph_path)?;
    let input_json = fs::read_to_string(&arguments.input_path)?;
    let artifact_load = load_started.elapsed();

    let initialization_started = Instant::now();
    let prover = CircuitProver::from_artifacts(
        &zkey,
        &arguments.zkey_sha256,
        &graph,
        &arguments.graph_sha256,
    )?;
    let artifact_initialization = initialization_started.elapsed();
    drop(zkey);
    drop(graph);

    let witness_started = Instant::now();
    let assignment = prover.calculate_witness_json(&input_json)?;
    let witness_calculation = witness_started.elapsed();
    drop(input_json);

    let proof_started = Instant::now();
    let bundle = prover.prove_assignment(&assignment)?;
    let proof_generation = proof_started.elapsed();

    fs::write(&arguments.proof_path, bundle.proof_json)?;
    fs::write(&arguments.public_path, bundle.public_signals_json)?;
    println!(
        concat!(
            "{{",
            "\"artifactLoadMs\":{:.3},",
            "\"artifactInitializationMs\":{:.3},",
            "\"witnessCalculationMs\":{:.3},",
            "\"proofGenerationMs\":{:.3},",
            "\"rayonThreads\":{}",
            "}}"
        ),
        elapsed_ms(artifact_load),
        elapsed_ms(artifact_initialization),
        elapsed_ms(witness_calculation),
        elapsed_ms(proof_generation),
        rayon_threads,
    );
    Ok(())
}

fn parse_arguments() -> Result<Arguments, Box<dyn Error>> {
    let mut arguments = env::args_os();
    let program = arguments
        .next()
        .unwrap_or_else(|| OsString::from("curvy-native-prover"));
    let usage = || {
        format!(
            concat!(
                "usage: {} <zkey> <zkey-sha256> <graph.bin> <graph-sha256> ",
                "<input.json> <proof.json> <public.json>"
            ),
            PathBuf::from(&program).display()
        )
    };
    let zkey_path = arguments.next().ok_or_else(usage)?;
    let zkey_sha256 = utf8_argument(arguments.next().ok_or_else(usage)?, "zkey SHA-256")?;
    let graph_path = arguments.next().ok_or_else(usage)?;
    let graph_sha256 = utf8_argument(arguments.next().ok_or_else(usage)?, "graph SHA-256")?;
    let input_path = arguments.next().ok_or_else(usage)?;
    let proof_path = arguments.next().ok_or_else(usage)?;
    let public_path = arguments.next().ok_or_else(usage)?;
    if arguments.next().is_some() {
        return Err(usage().into());
    }
    Ok(Arguments {
        zkey_path: zkey_path.into(),
        zkey_sha256,
        graph_path: graph_path.into(),
        graph_sha256,
        input_path: input_path.into(),
        proof_path: proof_path.into(),
        public_path: public_path.into(),
    })
}

fn utf8_argument(value: OsString, label: &str) -> Result<String, Box<dyn Error>> {
    value
        .into_string()
        .map_err(|_| format!("{label} must be valid UTF-8").into())
}

fn configured_threads() -> Result<usize, Box<dyn Error>> {
    let Some(value) = env::var_os("CURVY_PROVER_NUM_THREADS") else {
        return Ok(DEFAULT_THREADS);
    };
    let value = value
        .into_string()
        .map_err(|_| "CURVY_PROVER_NUM_THREADS must be valid UTF-8")?;
    let threads = value
        .parse::<usize>()
        .map_err(|_| "CURVY_PROVER_NUM_THREADS must be an integer")?;
    if !(1..=MAX_THREADS).contains(&threads) {
        return Err(format!("CURVY_PROVER_NUM_THREADS must be in 1..={MAX_THREADS}").into());
    }
    Ok(threads)
}

#[cfg(feature = "parallel")]
fn configure_rayon() -> Result<usize, Box<dyn Error>> {
    let threads = configured_threads()?;
    rayon::ThreadPoolBuilder::new()
        .num_threads(threads)
        .thread_name(|index| format!("curvy-prover-{index}"))
        .build_global()?;
    Ok(threads)
}

#[cfg(not(feature = "parallel"))]
fn configure_rayon() -> Result<usize, Box<dyn Error>> {
    let threads = configured_threads()?;
    if threads != 1 {
        return Err("this curvy-native-prover build does not include Rayon".into());
    }
    Ok(threads)
}

fn elapsed_ms(duration: Duration) -> f64 {
    duration.as_secs_f64() * 1_000.0
}