risc0_r0vm/
lib.rs

1// Copyright 2025 RISC Zero, Inc.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15mod actors;
16mod api;
17
18use std::{io, net::SocketAddr, path::PathBuf, rc::Rc};
19
20use clap::{Args, Parser, ValueEnum};
21use risc0_circuit_rv32im::execute::Segment;
22use risc0_zkvm::{
23    compute_image_id, get_prover_server, ApiServer, ExecutorEnv, ExecutorImpl, ProverOpts,
24    ProverServer, VerifierContext,
25};
26
27use self::actors::protocol::TaskKind;
28
29/// Runs a RISC-V ELF binary within the RISC Zero ZKVM.
30#[derive(Parser)]
31#[command(about, version, author)]
32struct Cli {
33    #[command(flatten)]
34    mode: Mode,
35
36    /// Receipt output file.
37    #[arg(long)]
38    receipt: Option<PathBuf>,
39
40    /// The hash function to use to produce a proof.
41    #[arg(long, value_enum, default_value_t = HashFn::Poseidon2)]
42    hashfn: HashFn,
43
44    /// Whether to prove executions ending in error status.
45    //
46    // When false, only prove execution sessions that end in a successful
47    // [ExitCode] (i.e. `Halted(0)` or `Paused(0)`. When set to true, any
48    // completed execution session will be proven, including indicated
49    // errors (e.g. `Halted(1)`) and sessions ending in `Fault`.
50    #[arg(long)]
51    prove_guest_errors: bool,
52
53    /// File to read initial input from.
54    ///
55    /// Reads input from stdin if an initial input file is not provided.
56    #[arg(long)]
57    initial_input: Option<PathBuf>,
58
59    /// Display verbose output.
60    #[arg(short, long, action = clap::ArgAction::Count)]
61    verbose: u8,
62
63    /// Add environment variables in the form of NAME=value.
64    #[arg(long, action = clap::ArgAction::Append)]
65    env: Vec<String>,
66
67    /// Write "pprof" protobuf output of the guest's run to this file.
68    /// You can use google's pprof (<https://github.com/google/pprof>)
69    /// to read it.
70    #[arg(long, env = "RISC0_PPROF_OUT")]
71    pprof_out: Option<PathBuf>,
72
73    /// The receipt kind produced by the r0vm prover
74    #[arg(long, value_enum, default_value_t = ReceiptKind::Composite)]
75    receipt_kind: ReceiptKind,
76
77    /// Compute the image_id for the specified ELF
78    #[arg(long)]
79    id: bool,
80
81    #[arg(long)]
82    with_debugger: bool,
83
84    /// The address to connect to or listen on.
85    #[arg(long)]
86    addr: Option<SocketAddr>,
87
88    #[arg(long)]
89    api: Option<SocketAddr>,
90
91    #[arg(long)]
92    storage: Option<PathBuf>,
93
94    #[arg(long)]
95    po2: Option<u32>,
96
97    /// Number of GPUs to use.
98    #[arg(long)]
99    num_gpus: Option<usize>,
100}
101
102#[derive(Args)]
103#[group(required = true)]
104struct Mode {
105    #[arg(long)]
106    rpc: bool,
107
108    #[arg(long)]
109    port: Option<u16>,
110
111    /// The ELF to execute
112    #[arg(long)]
113    elf: Option<PathBuf>,
114
115    /// The image to execute
116    #[arg(long)]
117    image: Option<PathBuf>,
118
119    /// Prove a pre-recorded segment.
120    #[arg(long)]
121    segment: Option<PathBuf>,
122
123    /// Start the manager.
124    #[arg(long)]
125    manager: bool,
126
127    /// Start a worker.
128    #[arg(long, value_enum, value_delimiter(','))]
129    worker: Vec<TaskKind>,
130
131    #[arg(long)]
132    config: Option<PathBuf>,
133}
134
135#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, ValueEnum)]
136enum HashFn {
137    #[value(name = "sha-256")]
138    Sha256,
139    #[value(name = "poseidon2")]
140    Poseidon2,
141}
142
143#[derive(Clone, PartialEq, ValueEnum)]
144enum ReceiptKind {
145    #[value(name = "composite")]
146    Composite,
147    #[value(name = "succinct")]
148    Succinct,
149    #[value(name = "groth16")]
150    Groth16,
151}
152
153pub fn main() {
154    let args = Cli::parse();
155
156    if args.mode.manager || !args.mode.worker.is_empty() || args.mode.config.is_some() {
157        self::actors::async_main(&args).unwrap();
158        return;
159    }
160
161    tracing_subscriber::fmt()
162        .with_env_filter(tracing_subscriber::filter::EnvFilter::from_default_env())
163        .init();
164
165    if args.mode.rpc {
166        self::actors::rpc_main(args.num_gpus).unwrap();
167        return;
168    }
169
170    if args.num_gpus.is_some_and(|v| v != 1) {
171        eprintln!("num_gpus > 1 or 0 for current mode unsupported.");
172        return;
173    }
174
175    if args.id {
176        let blob = std::fs::read(args.mode.elf.unwrap()).unwrap();
177        let image_id = compute_image_id(&blob).unwrap();
178        println!("{image_id}");
179        return;
180    }
181
182    if let Some(port) = args.mode.port {
183        run_server(port);
184        return;
185    }
186
187    if let Some(path) = args.mode.segment {
188        let bytes = std::fs::read(path).unwrap();
189        let segment = Segment::decode(&bytes).unwrap();
190        segment.execute().unwrap();
191        return;
192    }
193
194    let env = {
195        let mut builder = ExecutorEnv::builder();
196
197        for var in args.env.iter() {
198            let (name, value) = var
199                .split_once('=')
200                .expect("Environment variables should be of the form NAME=value");
201            builder.env_var(name, value);
202        }
203
204        if let Some(input) = args.initial_input.as_ref() {
205            builder.stdin(std::fs::File::open(input).unwrap());
206        } else {
207            builder.stdin(io::stdin());
208        }
209
210        if let Some(pprof_out) = args.pprof_out.as_ref() {
211            builder.enable_profiler(pprof_out);
212        }
213
214        builder.build().unwrap()
215    };
216
217    // TODO(povw): Add PoVW here.
218    let session = {
219        let mut exec = if let Some(ref elf_path) = args.mode.elf {
220            let elf_contents = std::fs::read(elf_path).unwrap();
221            ExecutorImpl::from_elf(env, &elf_contents).unwrap()
222        } else if let Some(ref image_path) = args.mode.image {
223            let image_contents = std::fs::read(image_path).unwrap();
224            let image = bincode::deserialize(&image_contents).unwrap();
225            ExecutorImpl::new(env, image).unwrap()
226        } else {
227            unreachable!()
228        };
229        if args.with_debugger {
230            exec.run_with_debugger().unwrap();
231            return;
232        } else {
233            exec.run().unwrap()
234        }
235    };
236
237    let prover = args.get_prover();
238    let ctx = VerifierContext::default();
239    let receipt = prover.prove_session(&ctx, &session).unwrap().receipt;
240
241    let receipt_data = bincode::serialize(&receipt).unwrap();
242    let receipt_bytes = bytemuck::cast_slice(&receipt_data);
243    if let Some(receipt_file) = args.receipt.as_ref() {
244        std::fs::write(receipt_file, receipt_bytes).expect("Unable to write receipt file");
245        if args.verbose > 0 {
246            eprintln!(
247                "Wrote {} bytes of receipt to {}",
248                receipt_data.len(),
249                receipt_file.display()
250            );
251        }
252    }
253}
254
255impl Cli {
256    fn get_prover(&self) -> Rc<dyn ProverServer> {
257        let hashfn = match self.hashfn {
258            HashFn::Sha256 => "sha-256",
259            HashFn::Poseidon2 => "poseidon2",
260        };
261        let opts = ProverOpts::default()
262            .with_hashfn(hashfn.to_string())
263            .with_prove_guest_errors(self.prove_guest_errors)
264            .with_receipt_kind(match self.receipt_kind {
265                ReceiptKind::Composite => risc0_zkvm::ReceiptKind::Composite,
266                ReceiptKind::Succinct => risc0_zkvm::ReceiptKind::Succinct,
267                ReceiptKind::Groth16 => risc0_zkvm::ReceiptKind::Groth16,
268            });
269        get_prover_server(&opts).unwrap()
270    }
271}
272
273fn run_server(port: u16) {
274    let addr = format!("127.0.0.1:{port}");
275    let server = ApiServer::new_tcp(addr);
276    server.run().unwrap()
277}
risc0_r0vm/lib.rs

risc0_r0vm/
lib.rs
