aboutsummaryrefslogtreecommitdiff
path: root/zkvms/zkm/host/src/main.rs
blob: e71c1215fb2d639c647e841f4a342f474afcb989 (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
use anyhow::{bail, Result};
use std::{env, fs::read, time::Instant};
use zkm_sdk::{prover::ClientCfg, prover::{ProverInput, ProverResult} , ProverClient};

use zkvms_host_io::{read_args, benchmarkable, RunType::{ Execute, Prove, Verify }};

async fn get_proof(
    prover_client: &mut ProverClient,
    prover_input: &mut ProverInput,
) -> ProverResult {
    let proving_result = prover_client.prover.prove(&prover_input, None).await;

    if let Ok(Some(prover_result)) = proving_result {
        prover_result
    }
    else {
        panic!("Failed to generate proof!");
    }
}

async fn execute(
    prover_client: &mut ProverClient,
    prover_input: &mut ProverInput,
) {
    prover_input.execute_only = true;

    let prover_result = get_proof(prover_client, prover_input).await;

    prover_client
        .print_guest_execution_output(false, &prover_result)
        .expect("print guest program excution's output false.")
}

async fn prove(
    prover_client: &mut ProverClient,
    prover_input: &mut ProverInput,
    vk_path: &String,
    proof_results_path: &String,
) {
    prover_input.execute_only = false;

    match prover_client
        .setup_and_generate_sol_verifier("local", &vk_path, &prover_input)
        .await
    {
        Ok(()) => println!("Succussfully setup_and_generate_sol_verifier."),
        Err(e) => panic!("Error during setup_and_generate_sol_verifier: {}", e),
    }

    let prover_result = get_proof(prover_client, prover_input).await;

    prover_client
        .process_proof_results(
            &prover_result,
            &prover_input,
            &proof_results_path,
            "local",
        )
        .expect("process proof results error");
}

#[tokio::main]
async fn main() -> Result<()> {
    let run_info = read_args();
    if run_info.run_type == Verify {
        panic!("Off-chain verification is not supported!");
    }

    let seg_size: u32 = run_info
        .env_then_or(
            "SEG_SIZE",
            |seg| seg.parse::<u32>().ok(),
            65536,
        );

    let elf_path = env::var("ELF_PATH").expect("ELF PATH is missing");

    let proof_results_path = run_info
        .env_or(
            "PROOF_RESULTS_PATH",
            "/tmp/contracts",
        );
    let vk_path = run_info
        .env_or(
            "VERIFYING_KEY_PATH",
            "/tmp/input",
        );

    let mut client_config: ClientCfg =
        ClientCfg::new("local".to_string(), vk_path.to_owned());

    let mut prover_client = ProverClient::new(&client_config).await;

    // Set input
    let mut public_inputstream = Vec::new();
    bincode::serialize_into(&mut public_inputstream, &run_info.public_input)
        .expect("public_input serialization failed");

    let mut private_inputstream = Vec::new();
    bincode::serialize_into(&mut private_inputstream, &run_info.private_input)
        .expect("private_input serialization failed");

    let mut prover_input = ProverInput {
        elf: read(elf_path).unwrap(),
        seg_size,
        public_inputstream,
        private_inputstream,
        ..Default::default()
    };

    let start = Instant::now();
    match run_info.run_type {
        // only excute the guest program without generating the proof.
        Execute => benchmarkable!{
            execute(&mut prover_client, &mut prover_input).await;
        },
        // excute the guest program and generate the proof
        Prove => benchmarkable!{
            prove(&mut prover_client, &mut prover_input, &vk_path, &proof_results_path).await;
        },
        Verify => unreachable!(),
    }
    let end = Instant::now();
    let elapsed = end.duration_since(start);
    println!("Elapsed time: {:?} secs", elapsed.as_secs());
    Ok(())
}