openvm_circuit/utils/
stark_utils.rs

Help
1use itertools::multiunzip;
2use openvm_instructions::{exe::VmExe, program::Program};
3use openvm_stark_backend::{
4    config::{StarkGenericConfig, Val},
5    p3_field::PrimeField32,
6    verifier::VerificationError,
7    Chip,
8};
9use openvm_stark_sdk::{
10    config::{
11        baby_bear_poseidon2::{BabyBearPoseidon2Config, BabyBearPoseidon2Engine},
12        setup_tracing, FriParameters,
13    },
14    engine::{StarkEngine, StarkFriEngine, VerificationDataWithFriParams},
15    p3_baby_bear::BabyBear,
16    utils::ProofInputForTest,
17};
18
19use crate::arch::{
20    vm::{VirtualMachine, VmExecutor},
21    Streams, VmConfig, VmMemoryState,
22};
23
24pub fn air_test<VC>(config: VC, exe: impl Into<VmExe<BabyBear>>)
25where
26    VC: VmConfig<BabyBear>,
27    VC::Executor: Chip<BabyBearPoseidon2Config>,
28    VC::Periphery: Chip<BabyBearPoseidon2Config>,
29{
30    air_test_with_min_segments(config, exe, Streams::default(), 1);
31}
32
33/// Executes and proves the VM and returns the final memory state.
34pub fn air_test_with_min_segments<VC>(
35    config: VC,
36    exe: impl Into<VmExe<BabyBear>>,
37    input: impl Into<Streams<BabyBear>>,
38    min_segments: usize,
39) -> Option<VmMemoryState<BabyBear>>
40where
41    VC: VmConfig<BabyBear>,
42    VC::Executor: Chip<BabyBearPoseidon2Config>,
43    VC::Periphery: Chip<BabyBearPoseidon2Config>,
44{
45    air_test_impl(config, exe, input, min_segments, true)
46}
47
48/// Executes and proves the VM and returns the final memory state.
49/// If `debug` is true, runs the debug prover.
50pub fn air_test_impl<VC>(
51    config: VC,
52    exe: impl Into<VmExe<BabyBear>>,
53    input: impl Into<Streams<BabyBear>>,
54    min_segments: usize,
55    debug: bool,
56) -> Option<VmMemoryState<BabyBear>>
57where
58    VC: VmConfig<BabyBear>,
59    VC::Executor: Chip<BabyBearPoseidon2Config>,
60    VC::Periphery: Chip<BabyBearPoseidon2Config>,
61{
62    setup_tracing();
63    let mut log_blowup = 1;
64    while config.system().max_constraint_degree > (1 << log_blowup) + 1 {
65        log_blowup += 1;
66    }
67    let engine = BabyBearPoseidon2Engine::new(FriParameters::new_for_testing(log_blowup));
68    let vm = VirtualMachine::new(engine, config);
69    let pk = vm.keygen();
70    let mut result = vm.execute_and_generate(exe, input).unwrap();
71    let final_memory = Option::take(&mut result.final_memory);
72    let global_airs = vm.config().create_chip_complex().unwrap().airs();
73    if debug {
74        for proof_input in &result.per_segment {
75            let (airs, pks, air_proof_inputs): (Vec<_>, Vec<_>, Vec<_>) =
76                multiunzip(proof_input.per_air.iter().map(|(air_id, air_proof_input)| {
77                    (
78                        global_airs[*air_id].clone(),
79                        pk.per_air[*air_id].clone(),
80                        air_proof_input.clone(),
81                    )
82                }));
83            vm.engine.debug(&airs, &pks, &air_proof_inputs);
84        }
85    }
86    let proofs = vm.prove(&pk, result);
87
88    assert!(proofs.len() >= min_segments);
89    vm.verify(&pk.get_vk(), proofs)
90        .expect("segment proofs should verify");
91    final_memory
92}
93
94/// Generates the VM STARK circuit, in the form of AIRs and traces, but does not
95/// do any proving. Output is the payload of everything the prover needs.
96///
97/// The output AIRs and traces are sorted by height in descending order.
98pub fn gen_vm_program_test_proof_input<SC: StarkGenericConfig, VC>(
99    program: Program<Val<SC>>,
100    input_stream: impl Into<Streams<Val<SC>>> + Clone,
101    #[allow(unused_mut)] mut config: VC,
102) -> ProofInputForTest<SC>
103where
104    Val<SC>: PrimeField32,
105    VC: VmConfig<Val<SC>> + Clone,
106    VC::Executor: Chip<SC>,
107    VC::Periphery: Chip<SC>,
108{
109    cfg_if::cfg_if! {
110        if #[cfg(feature = "bench-metrics")] {
111            // Run once with metrics collection enabled, which can improve runtime performance
112            config.system_mut().profiling = true;
113            {
114                let executor = VmExecutor::<Val<SC>, VC>::new(config.clone());
115                executor.execute(program.clone(), input_stream.clone()).unwrap();
116            }
117            // Run again with metrics collection disabled and measure trace generation time
118            config.system_mut().profiling = false;
119            let start = std::time::Instant::now();
120        }
121    }
122
123    let airs = config.create_chip_complex().unwrap().airs();
124    let executor = VmExecutor::<Val<SC>, VC>::new(config);
125
126    let mut result = executor
127        .execute_and_generate(program, input_stream)
128        .unwrap();
129    assert_eq!(
130        result.per_segment.len(),
131        1,
132        "only proving one segment for now"
133    );
134
135    let result = result.per_segment.pop().unwrap();
136    #[cfg(feature = "bench-metrics")]
137    metrics::gauge!("execute_and_trace_gen_time_ms").set(start.elapsed().as_millis() as f64);
138    // Filter out unused AIRS (where trace is empty)
139    let (used_airs, per_air) = result
140        .per_air
141        .into_iter()
142        .map(|(air_id, x)| (airs[air_id].clone(), x))
143        .unzip();
144    ProofInputForTest {
145        airs: used_airs,
146        per_air,
147    }
148}
149
150type ExecuteAndProveResult<SC> = Result<VerificationDataWithFriParams<SC>, VerificationError>;
151
152/// Executes program and runs simple STARK prover test (keygen, prove, verify).
153pub fn execute_and_prove_program<SC: StarkGenericConfig, E: StarkFriEngine<SC>, VC>(
154    program: Program<Val<SC>>,
155    input_stream: impl Into<Streams<Val<SC>>> + Clone,
156    config: VC,
157    engine: &E,
158) -> ExecuteAndProveResult<SC>
159where
160    Val<SC>: PrimeField32,
161    VC: VmConfig<Val<SC>> + Clone,
162    VC::Executor: Chip<SC>,
163    VC::Periphery: Chip<SC>,
164{
165    let span = tracing::info_span!("execute_and_prove_program").entered();
166    let test_proof_input = gen_vm_program_test_proof_input(program, input_stream, config);
167    let vparams = test_proof_input.run_test(engine)?;
168    span.exit();
169    Ok(vparams)
170}
openvm_circuit/utils/stark_utils.rs

openvm_circuit/utils/
stark_utils.rs
