1use std::{path::PathBuf, sync::Arc};
2
3use clap::Parser;
4use eyre::{eyre, Result};
5use openvm_circuit::arch::{
6 execution_mode::metered::segment_ctx::DEFAULT_MAX_MEMORY, instructions::exe::VmExe,
7};
8use openvm_continuations::CommitBytes;
9#[cfg(feature = "evm-prove")]
10use openvm_sdk::keygen::RootProvingKey;
11use openvm_sdk::{
12 config::{AggregationSystemParams, AggregationTreeConfig},
13 fs::{read_object_from_file, write_object_to_file, write_to_file_json},
14 keygen::{AggPrefixProvingKey, AggProvingKey, AppProvingKey},
15 types::{AppExecutionCommit, VerificationBaselineJson, VersionedVmStarkProof},
16 Sdk, F, SC,
17};
18use openvm_sdk_config::SdkVmConfig;
19use openvm_stark_backend::keygen::types::MultiStarkProvingKey;
20
21use super::{RunArgs, RunCargoArgs};
22use crate::{
23 args::ProvingKeyArgs,
24 commands::build,
25 default::{default_internal_recursive_pk_path, APP_PROOF_EXT, STARK_PROOF_EXT, VMEXE_EXT},
26 input::read_to_stdin,
27 util::{
28 get_agg_prefix_pk_path, get_app_baseline_path, get_app_pk_path, get_manifest_path_and_dir,
29 get_single_target_name, get_target_dir, get_target_output_dir,
30 },
31};
32
33#[derive(Parser)]
34#[command(name = "prove", about = "Generate a program proof")]
35pub struct ProveCmd {
36 #[command(subcommand)]
37 command: ProveSubCommand,
38}
39
40#[derive(Parser)]
41enum ProveSubCommand {
42 App {
43 #[arg(
44 long,
45 action,
46 help = "Path to app proof output, by default will be ./${bin_name}.app.proof",
47 help_heading = "Output"
48 )]
49 proof: Option<PathBuf>,
50
51 #[arg(
52 long,
53 action,
54 help = "Path to app proving key, by default will be ${openvm_dir}/app.pk",
55 help_heading = "OpenVM Options"
56 )]
57 app_pk: Option<PathBuf>,
58
59 #[command(flatten)]
60 run_args: RunArgs,
61
62 #[command(flatten)]
63 cargo_args: RunCargoArgs,
64
65 #[command(flatten)]
66 segmentation_args: SegmentationArgs,
67 },
68 Stark {
69 #[arg(
70 long,
71 action,
72 help = "Path to STARK proof output, by default will be ./${bin_name}.stark.proof",
73 help_heading = "Output"
74 )]
75 proof: Option<PathBuf>,
76
77 #[command(flatten)]
78 keys: ProvingKeyArgs,
79
80 #[command(flatten)]
81 run_args: RunArgs,
82
83 #[command(flatten)]
84 cargo_args: RunCargoArgs,
85
86 #[command(flatten)]
87 segmentation_args: SegmentationArgs,
88
89 #[command(flatten)]
90 agg_tree_config: AggregationTreeConfig,
91 },
92 #[cfg(feature = "evm-prove")]
93 Evm {
94 #[arg(
95 long,
96 action,
97 help = "Path to EVM proof output, by default will be ./${bin_name}.evm.proof",
98 help_heading = "Output"
99 )]
100 proof: Option<PathBuf>,
101
102 #[command(flatten)]
103 keys: ProvingKeyArgs,
104
105 #[arg(
106 long,
107 action,
108 help = "Path to root proving key, by default will be ${HOME}/.openvm/root.pk",
109 help_heading = "OpenVM Options"
110 )]
111 root_pk: Option<PathBuf>,
112
113 #[arg(
114 long,
115 action,
116 help = "Path to Halo2 proving key, by default will be ${HOME}/.openvm/halo2.pk",
117 help_heading = "OpenVM Options"
118 )]
119 halo2_pk: Option<PathBuf>,
120
121 #[command(flatten)]
122 run_args: RunArgs,
123
124 #[command(flatten)]
125 cargo_args: RunCargoArgs,
126
127 #[command(flatten)]
128 segmentation_args: SegmentationArgs,
129
130 #[command(flatten)]
131 agg_tree_config: AggregationTreeConfig,
132 },
133}
134
135#[derive(Clone, Copy, Parser)]
136pub struct SegmentationArgs {
137 #[arg(
140 long,
141 default_value_t = DEFAULT_MAX_MEMORY,
142 help_heading = "OpenVM Options"
143 )]
144 pub segment_max_memory: usize,
145}
146
147impl ProveCmd {
148 pub fn run(&self) -> Result<()> {
149 match &self.command {
150 ProveSubCommand::App {
151 app_pk,
152 proof,
153 run_args,
154 cargo_args,
155 segmentation_args,
156 } => {
157 let mut app_pk = load_app_pk(app_pk, cargo_args)?;
158 configure_app_pk(&mut app_pk, segmentation_args);
159 let sdk = Sdk::builder()
160 .app_pk(app_pk)
161 .agg_params(AggregationSystemParams::default())
162 .build()?;
163 let (exe, target_name) = load_or_build_exe(run_args, cargo_args)?;
164
165 let mut prover = sdk.app_prover(exe)?;
166 let exe_commit = CommitBytes::from(prover.app_exe_commit());
167 println!("exe commit: {exe_commit}");
168
169 let app_proof = prover.prove(read_to_stdin(&run_args.input)?)?;
170
171 let proof_path = if let Some(proof) = proof {
172 proof
173 } else {
174 &PathBuf::from(target_name).with_extension(APP_PROOF_EXT)
175 };
176 println!(
177 "App proof completed! Writing App proof to {}",
178 proof_path.display()
179 );
180 write_object_to_file(proof_path, app_proof)?;
181 }
182 ProveSubCommand::Stark {
183 keys,
184 proof,
185 run_args,
186 cargo_args,
187 segmentation_args,
188 agg_tree_config,
189 } => {
190 let mut app_pk = load_app_pk(&keys.app_pk, cargo_args)?;
191 let (exe, target_name) = load_or_build_exe(run_args, cargo_args)?;
192 configure_app_pk(&mut app_pk, segmentation_args);
193 let agg_pk = load_required_agg_pk(&keys.agg_prefix_pk, &keys.agg_pk, cargo_args)?;
194 let sdk = Sdk::builder()
195 .app_pk(app_pk)
196 .agg_pk(agg_pk)
197 .agg_tree_config(*agg_tree_config)
198 .build()?;
199 let mut prover = sdk.prover(exe)?;
200 let baseline = prover.generate_baseline();
201 let app_vm_commit = prover.app_vm_commit();
202
203 let app_commit = AppExecutionCommit {
204 app_exe_commit: CommitBytes::from(baseline.app_exe_commit),
205 app_vm_commit: CommitBytes::from(app_vm_commit),
206 };
207 println!("exe commit: {}", app_commit.app_exe_commit);
208 println!("vm commit: {}", app_commit.app_vm_commit);
209
210 let (stark_proof, _metadata) =
211 prover.prove(read_to_stdin(&run_args.input)?, &[])?;
212 let stark_proof_bytes = VersionedVmStarkProof::new(stark_proof)?;
213
214 let target_dir = target_dir_from_cargo_args(cargo_args)?;
215 let target_output_dir = get_target_output_dir(&target_dir, &cargo_args.profile);
216 let target_name_path =
217 get_single_target_name(cargo_args).unwrap_or(PathBuf::from(&target_name));
218 let baseline_path = get_app_baseline_path(&target_output_dir, target_name_path);
219 println!("Writing baseline to {}", baseline_path.display());
220 let baseline_json: VerificationBaselineJson = baseline.into();
221 write_to_file_json(&baseline_path, &baseline_json)?;
222
223 let proof_path = if let Some(proof) = proof {
224 proof
225 } else {
226 &PathBuf::from(target_name).with_extension(STARK_PROOF_EXT)
227 };
228 println!(
229 "STARK proof completed! Writing STARK proof to {}",
230 proof_path.display()
231 );
232 write_to_file_json(proof_path, stark_proof_bytes)?;
233 }
234 #[cfg(feature = "evm-prove")]
235 ProveSubCommand::Evm {
236 keys,
237 proof,
238 root_pk,
239 halo2_pk,
240 run_args,
241 cargo_args,
242 segmentation_args,
243 agg_tree_config,
244 } => {
245 let mut app_pk = load_app_pk(&keys.app_pk, cargo_args)?;
246 let (exe, target_name) = load_or_build_exe(run_args, cargo_args)?;
247
248 println!("Generating EVM proof, this may take a lot of compute and memory...");
249 configure_app_pk(&mut app_pk, segmentation_args);
250 let agg_pk = load_required_agg_pk(&keys.agg_prefix_pk, &keys.agg_pk, cargo_args)?;
251 let root_pk = load_required_root_pk(root_pk)?;
252 let halo2_pk = load_required_halo2_pk(halo2_pk)?;
253 let sdk = Sdk::builder()
254 .app_pk(app_pk)
255 .agg_pk(agg_pk)
256 .root_pk(root_pk)
257 .halo2_pk(halo2_pk)
258 .agg_tree_config(*agg_tree_config)
259 .build()?;
260 let mut prover = sdk.evm_prover(exe)?;
261 let exe_commit = CommitBytes::from(prover.stark_prover.app_prover.app_exe_commit());
262 println!("exe commit: {exe_commit}");
263 let evm_proof = prover.prove_evm(read_to_stdin(&run_args.input)?, &[])?;
264
265 let proof_path = if let Some(proof) = proof {
266 proof
267 } else {
268 &PathBuf::from(target_name).with_extension(crate::default::EVM_PROOF_EXT)
269 };
270 println!(
271 "EVM proof completed! Writing EVM proof to {}",
272 proof_path.display()
273 );
274 write_to_file_json(proof_path, evm_proof)?;
275 }
276 }
277 Ok(())
278 }
279}
280
281pub(crate) fn load_app_pk(
282 app_pk: &Option<PathBuf>,
283 cargo_args: &RunCargoArgs,
284) -> Result<AppProvingKey<SdkVmConfig>> {
285 let app_pk_path = if let Some(app_pk) = app_pk {
286 app_pk.to_path_buf()
287 } else {
288 let (manifest_path, _) = get_manifest_path_and_dir(&cargo_args.manifest.manifest_path)?;
289 let target_dir = get_target_dir(&cargo_args.manifest.target_dir, &manifest_path);
290 get_app_pk_path(&target_dir)
291 };
292
293 read_object_from_file(app_pk_path)
294}
295
296pub(crate) fn load_or_build_exe(
299 run_args: &RunArgs,
300 cargo_args: &RunCargoArgs,
301) -> Result<(VmExe<F>, String)> {
302 let exe_path = if let Some(exe) = &run_args.exe {
303 exe
304 } else {
305 let target_name = get_single_target_name(cargo_args)?;
307 let build_args = run_args.clone().into();
308 let cargo_args = cargo_args.clone().into();
309 let output_dir = build(&build_args, &cargo_args)?;
310 &output_dir.join(target_name.with_extension(VMEXE_EXT))
311 };
312
313 let app_exe = read_object_from_file(exe_path)?;
314 Ok((
315 app_exe,
316 exe_path.file_stem().unwrap().to_string_lossy().into_owned(),
317 ))
318}
319
320fn configure_app_pk(app_pk: &mut AppProvingKey<SdkVmConfig>, segmentation_args: &SegmentationArgs) {
323 Arc::get_mut(&mut app_pk.app_vm_pk)
324 .unwrap()
325 .vm_config
326 .system
327 .config
328 .set_segmentation_max_memory(segmentation_args.segment_max_memory);
329}
330
331fn target_dir_from_cargo_args(cargo_args: &RunCargoArgs) -> Result<PathBuf> {
332 let (manifest_path, _) = get_manifest_path_and_dir(&cargo_args.manifest.manifest_path)?;
333 Ok(get_target_dir(
334 &cargo_args.manifest.target_dir,
335 &manifest_path,
336 ))
337}
338
339fn resolve_agg_prefix_pk_path(
340 agg_prefix_pk: &Option<PathBuf>,
341 cargo_args: &RunCargoArgs,
342) -> Result<PathBuf> {
343 if let Some(agg_prefix_pk) = agg_prefix_pk {
344 Ok(agg_prefix_pk.to_path_buf())
345 } else {
346 let target_dir = target_dir_from_cargo_args(cargo_args)?;
347 Ok(get_agg_prefix_pk_path(&target_dir))
348 }
349}
350
351pub(crate) fn load_required_agg_pk(
352 agg_prefix_pk: &Option<PathBuf>,
353 agg_pk: &Option<PathBuf>,
354 cargo_args: &RunCargoArgs,
355) -> Result<AggProvingKey> {
356 let prefix_pk = load_required_agg_prefix_pk(agg_prefix_pk, cargo_args)?;
357 let internal_recursive_pk = load_required_internal_recursive_pk(agg_pk)?;
358 Ok(AggProvingKey {
359 prefix: prefix_pk,
360 internal_recursive: internal_recursive_pk,
361 })
362}
363
364pub(crate) fn load_required_agg_prefix_pk(
365 agg_prefix_pk: &Option<PathBuf>,
366 cargo_args: &RunCargoArgs,
367) -> Result<AggPrefixProvingKey> {
368 let agg_prefix_pk_path = resolve_agg_prefix_pk_path(agg_prefix_pk, cargo_args)?;
369 read_object_from_file(&agg_prefix_pk_path).map_err(|e| {
370 eyre!(
371 "Failed to read aggregation prefix proving key from {}: {e}\nRun 'cargo openvm keygen' first to generate it",
372 agg_prefix_pk_path.display()
373 )
374 })
375}
376
377pub(crate) fn load_required_internal_recursive_pk(
378 agg_pk: &Option<PathBuf>,
379) -> Result<Arc<MultiStarkProvingKey<SC>>> {
380 let internal_recursive_pk_path = agg_pk
381 .clone()
382 .unwrap_or_else(|| PathBuf::from(default_internal_recursive_pk_path()));
383 read_object_from_file(&internal_recursive_pk_path).map_err(|e| {
384 eyre!(
385 "Failed to read internal-recursive proving key from {}: {e}\nRun 'cargo openvm setup' first to generate it",
386 internal_recursive_pk_path.display()
387 )
388 })
389}
390
391#[cfg(feature = "evm-prove")]
392fn load_required_root_pk(root_pk: &Option<PathBuf>) -> Result<RootProvingKey> {
393 let root_pk_path = root_pk
394 .clone()
395 .unwrap_or_else(|| PathBuf::from(crate::default::default_root_pk_path()));
396 read_object_from_file(&root_pk_path).map_err(|e| {
397 eyre!(
398 "Failed to read root proving key from {}: {e}\nRun 'cargo openvm setup --evm' first to generate it",
399 root_pk_path.display()
400 )
401 })
402}
403
404#[cfg(feature = "evm-prove")]
405fn load_required_halo2_pk(
406 halo2_pk: &Option<PathBuf>,
407) -> Result<openvm_sdk::keygen::Halo2ProvingKey> {
408 let halo2_pk_path = halo2_pk
409 .clone()
410 .unwrap_or_else(|| PathBuf::from(crate::default::default_halo2_pk_path()));
411 openvm_sdk::fs::read_halo2_pk_from_file(&halo2_pk_path).map_err(|e| {
412 eyre!(
413 "Failed to read Halo2 proving key from {}: {e}\nRun 'cargo openvm setup --evm' first to generate it",
414 halo2_pk_path.display()
415 )
416 })
417}