generate_fixtures/
generate-fixtures.rs

1use std::{fs, path::Path};
2
3use eyre::Result;
4use openvm_benchmarks_utils::{get_elf_path, get_fixtures_dir, get_programs_dir, read_elf_file};
5use openvm_circuit::arch::{
6    PreflightExecutor, SingleSegmentVmProver, VmBuilder, VmExecutionConfig,
7};
8use openvm_continuations::verifier::internal::types::InternalVmVerifierInput;
9use openvm_native_circuit::{NativeConfig, NATIVE_MAX_TRACE_HEIGHTS};
10use openvm_native_recursion::hints::Hintable;
11use openvm_sdk::{
12    config::{AggregationTreeConfig, AppConfig, AppFriParams, SdkVmConfig},
13    prover::AggStarkProver,
14    CpuSdk, StdIn, F, SC,
15};
16use openvm_stark_sdk::{
17    config::baby_bear_poseidon2::BabyBearPoseidon2Engine, engine::StarkFriEngine,
18    openvm_stark_backend::proof::Proof,
19};
20use serde::Serialize;
21use tracing::info_span;
22use tracing_subscriber::{fmt, EnvFilter};
23
24const PROGRAMS_AND_INPUTS: &[(&str, Option<u64>)] =
25    &[("fibonacci", Some(1 << 22)), ("kitchen-sink", None)];
26
27/// Helper function to serialize and write data to a file
28fn write_fixture<T: Serialize>(path: impl AsRef<Path>, data: &T, description: &str) -> Result<()> {
29    tracing::debug!("Writing {}", description);
30    let bytes = bitcode::serialize(data)?;
31    fs::write(path, bytes)?;
32    Ok(())
33}
34
35/// Aggregates leaf proofs into internal proofs and saves them to disk
36/// Returns the final internal proof and the count of internal proofs saved
37fn aggregate_leaf_proofs<E, NativeBuilder>(
38    agg_prover: &mut AggStarkProver<E, NativeBuilder>,
39    leaf_proofs: Vec<Proof<SC>>,
40    fixtures_dir: &Path,
41    program: &str,
42) -> Result<(Proof<SC>, usize)>
43where
44    E: StarkFriEngine<SC = SC>,
45    NativeBuilder: VmBuilder<E, VmConfig = NativeConfig>,
46    <NativeConfig as VmExecutionConfig<F>>::Executor:
47        PreflightExecutor<F, <NativeBuilder as VmBuilder<E>>::RecordArena>,
48{
49    let mut internal_node_idx = -1;
50    let mut internal_node_height = 0;
51    let mut internal_proof_count = 0;
52    let mut proofs = leaf_proofs;
53
54    // We will always generate at least one internal proof, even if there is only one leaf
55    // proof, in order to shrink the proof size
56    while proofs.len() > 1 || internal_node_height == 0 {
57        let internal_inputs = InternalVmVerifierInput::chunk_leaf_or_internal_proofs(
58            (*agg_prover.internal_prover.program_commitment()).into(),
59            &proofs,
60            agg_prover.num_children_internal,
61        );
62
63        proofs = info_span!(
64            "agg_layer",
65            group = format!("internal.{internal_node_height}")
66        )
67        .in_scope(|| -> Result<Vec<Proof<SC>>> {
68            internal_inputs
69                .into_iter()
70                .enumerate()
71                .map(|(i, input)| -> Result<Proof<SC>> {
72                    internal_node_idx += 1;
73                    let proof = info_span!("single_internal_agg", idx = internal_node_idx)
74                        .in_scope(|| {
75                            SingleSegmentVmProver::prove(
76                                &mut agg_prover.internal_prover,
77                                input.write(),
78                                NATIVE_MAX_TRACE_HEIGHTS,
79                            )
80                        })?;
81
82                    // Save proof
83                    write_fixture(
84                        fixtures_dir.join(format!(
85                            "{}.internal.{}.proof",
86                            program,
87                            internal_proof_count + i
88                        )),
89                        &proof,
90                        &format!("internal proof {}", internal_proof_count + i),
91                    )?;
92
93                    Ok(proof)
94                })
95                .collect::<Result<Vec<_>, _>>()
96        })?;
97
98        internal_proof_count += proofs.len();
99        internal_node_height += 1;
100    }
101
102    let final_internal_proof = proofs.pop().unwrap();
103    Ok((final_internal_proof, internal_proof_count))
104}
105
106fn main() -> Result<()> {
107    // Set up logging
108    fmt::fmt().with_env_filter(EnvFilter::new("info")).init();
109
110    // Create fixtures directory if it doesn't exist
111    let fixtures_dir = get_fixtures_dir();
112    fs::create_dir_all(&fixtures_dir)?;
113
114    tracing::info!("Processing {} programs", PROGRAMS_AND_INPUTS.len());
115
116    for (idx, &(program, input)) in PROGRAMS_AND_INPUTS.iter().enumerate() {
117        tracing::info!(
118            "Processing program {}/{}: {} {}",
119            idx + 1,
120            PROGRAMS_AND_INPUTS.len(),
121            program,
122            input
123                .map(|i| format!("(input: {})", i))
124                .unwrap_or_else(|| "(no input)".to_string())
125        );
126
127        let program_dir = get_programs_dir().join(program);
128
129        tracing::info!(program = %program, "Loading VM config");
130        let config_path = program_dir.join("openvm.toml");
131        let config_content = fs::read_to_string(&config_path)?;
132        let vm_config = SdkVmConfig::from_toml(&config_content)?.app_vm_config;
133
134        tracing::info!(program = %program, "Preparing ELF");
135        let elf_path = get_elf_path(&program_dir);
136        let elf = read_elf_file(&elf_path)?;
137
138        // Create app config with default parameters
139        let app_config = AppConfig::new(AppFriParams::default().fri_params, vm_config);
140
141        let sdk = CpuSdk::new(app_config.clone())?;
142        let exe = sdk.convert_to_exe(elf)?;
143
144        // Prepare stdin
145        let mut stdin = StdIn::default();
146        if let Some(input_value) = input {
147            tracing::info!(program = %program, input = %input_value, "Preparing stdin with input");
148            stdin.write(&input_value);
149        } else {
150            tracing::info!(program = %program, "No input provided for program");
151        }
152
153        tracing::info!(program = %program, "Generating app proof");
154        let app_proof = sdk.app_prover(exe)?.prove(stdin)?;
155
156        // Save app proof
157        write_fixture(
158            fixtures_dir.join(format!("{}.app.proof", program)),
159            &app_proof,
160            "app proof",
161        )?;
162
163        tracing::info!(program = %program, "Getting keys");
164        let app_pk = sdk.app_pk();
165        let agg_pk = sdk.agg_pk();
166
167        // Save keys
168        write_fixture(
169            fixtures_dir.join(format!("{}.leaf.exe", program)),
170            &app_pk.leaf_committed_exe.exe,
171            "leaf exe",
172        )?;
173
174        write_fixture(
175            fixtures_dir.join(format!("{}.leaf.pk", program)),
176            &agg_pk.leaf_vm_pk.vm_pk,
177            "leaf proving key",
178        )?;
179
180        write_fixture(
181            fixtures_dir.join(format!("{}.internal.exe", program)),
182            &agg_pk.internal_committed_exe.exe,
183            "internal exe",
184        )?;
185
186        write_fixture(
187            fixtures_dir.join(format!("{}.internal.pk", program)),
188            &agg_pk.internal_vm_pk.vm_pk,
189            "internal proving key",
190        )?;
191
192        tracing::info!(program = %program, "Creating aggregation provers");
193        let native_builder = sdk.native_builder().clone();
194        let leaf_verifier_exe = app_pk.leaf_committed_exe.exe.clone();
195
196        let tree_config = AggregationTreeConfig::default();
197        let mut agg_prover = AggStarkProver::<BabyBearPoseidon2Engine, _>::new(
198            native_builder.clone(),
199            agg_pk,
200            leaf_verifier_exe,
201            tree_config,
202        )?;
203
204        tracing::info!(program = %program, "Generating leaf proofs");
205        let leaf_proofs = agg_prover.generate_leaf_proofs(&app_proof)?;
206        tracing::info!(program = %program, leaf_proof_count = leaf_proofs.len(), "Generated leaf proofs");
207
208        // Save leaf proofs
209        for (i, leaf_proof) in leaf_proofs.iter().enumerate() {
210            write_fixture(
211                fixtures_dir.join(format!("{}.leaf.{}.proof", program, i)),
212                leaf_proof,
213                &format!("leaf proof {}", i),
214            )?;
215        }
216
217        tracing::info!(program = %program, "Generating internal proofs");
218
219        let (_final_internal_proof, internal_proof_count) =
220            aggregate_leaf_proofs(&mut agg_prover, leaf_proofs.clone(), &fixtures_dir, program)?;
221
222        tracing::info!(
223            program = %program,
224            leaf_proofs = leaf_proofs.len(),
225            total_internals = internal_proof_count,
226            "Generated and saved {} fixtures: leaf.exe, leaf.pk, internal.exe, internal.pk, app.proof, {} leaf proofs, and {} internal proofs",
227            program,
228            leaf_proofs.len(),
229            internal_proof_count
230        );
231    }
232
233    tracing::info!("Successfully processed all programs");
234    Ok(())
235}
generate_fixtures/generate-fixtures.rs

generate_fixtures/
generate-fixtures.rs
