openvm_stark_backend/test_utils/
mod.rs

1use std::sync::Arc;
2
3use itertools::Itertools;
4use p3_field::{PrimeCharacteristicRing, PrimeField64};
5use p3_matrix::dense::RowMajorMatrix;
6
7use self::dummy_airs::{
8    fib_air::air::FibonacciAir,
9    fib_selector_air::air::FibonacciSelectorAir,
10    interaction::{
11        dummy_interaction_air::DummyInteractionAir,
12        self_interaction_air::{SelfInteractionAir, SelfInteractionChip},
13    },
14    preprocessed_cached_air::air::PreprocessedCachedAir,
15};
16use crate::{
17    interaction::{BusIndex, LogUpSecurityParameters},
18    keygen::types::{MultiStarkProvingKey, MultiStarkVerifyingKey},
19    proof::Proof,
20    prover::{
21        stacked_pcs::stacked_commit, AirProvingContext, ColMajorMatrix, CommittedTraceData,
22        CpuColMajorBackend, DeviceDataTransporter, DeviceMultiStarkProvingKey, MatrixDimensions,
23        MultiRapProver, Prover, ProvingContext, TraceCommitter,
24    },
25    AirRef, StarkEngine, StarkProtocolConfig, SystemParams, WhirConfig, WhirParams,
26    WhirProximityStrategy,
27};
28
29pub mod dummy_airs;
30
31/// Macro to create a `Vec<AirRef<SC>>` from a list of AIRs.
32#[macro_export]
33macro_rules! any_air_arc_vec {
34    ($($air:expr),+ $(,)?) => {
35        vec![$(std::sync::Arc::new($air) as $crate::AirRef<_>),+]
36    };
37}
38
39#[allow(clippy::type_complexity)]
40pub fn prove_up_to_batch_constraints<E: StarkEngine>(
41    engine: &E,
42    transcript: &mut E::TS,
43    pk: &DeviceMultiStarkProvingKey<E::PB>,
44    ctx: ProvingContext<E::PB>,
45) -> (
46    <E::PD as MultiRapProver<E::PB, E::TS>>::PartialProof,
47    <E::PD as MultiRapProver<E::PB, E::TS>>::Artifacts,
48) {
49    let (_, common_main_pcs_data) = engine
50        .device()
51        .commit(
52            &ctx.common_main_traces()
53                .map(|(_, trace)| trace)
54                .collect_vec(),
55        )
56        .unwrap();
57    engine
58        .device()
59        .prove_rap_constraints(transcript, pk, &ctx, &common_main_pcs_data)
60        .unwrap()
61}
62
63fn get_fib_number<F: PrimeField64>(mut a: u64, mut b: u64, n: usize) -> u64 {
64    for _ in 0..n - 1 {
65        let c = (a + b) % F::ORDER_U64;
66        a = b;
67        b = c;
68    }
69    b
70}
71
72fn get_conditional_fib_number<F: PrimeField64>(mut a: u64, mut b: u64, sels: &[bool]) -> u64 {
73    for &s in sels[0..sels.len() - 1].iter() {
74        if s {
75            let c = (a + b) % F::ORDER_U64;
76            a = b;
77            b = c;
78        }
79    }
80    b
81}
82
83fn commit_cached_trace<SC: StarkProtocolConfig>(
84    config: &SC,
85    trace: ColMajorMatrix<SC::F>,
86) -> CommittedTraceData<CpuColMajorBackend<SC>> {
87    let params = config.params();
88    let (commitment, data) = stacked_commit(
89        config.hasher(),
90        params.l_skip,
91        params.n_stack,
92        params.log_blowup,
93        params.k_whir(),
94        &[&trace],
95    )
96    .unwrap();
97    CommittedTraceData {
98        commitment,
99        trace,
100        data: Arc::new(data),
101    }
102}
103/// Trait for object responsible for generating the collection of AIRs and trace matrices for a
104/// single test case.
105pub trait TestFixture<SC: StarkProtocolConfig> {
106    fn airs(&self) -> Vec<AirRef<SC>>;
107
108    fn generate_proving_ctx(&self) -> ProvingContext<CpuColMajorBackend<SC>>;
109
110    fn keygen<E: StarkEngine<SC = SC>>(
111        &self,
112        engine: &E,
113    ) -> (MultiStarkProvingKey<SC>, MultiStarkVerifyingKey<SC>) {
114        engine.keygen(&self.airs())
115    }
116
117    fn prove<E: StarkEngine<SC = SC>>(
118        &self,
119        engine: &E,
120        pk: &MultiStarkProvingKey<SC>,
121    ) -> Proof<SC> {
122        self.prove_from_transcript(engine, pk, &mut engine.initial_transcript())
123    }
124
125    /// Prove using CPU tracegen and transport to device.
126    fn prove_from_transcript<E: StarkEngine<SC = SC>>(
127        &self,
128        engine: &E,
129        pk: &MultiStarkProvingKey<SC>,
130        transcript: &mut E::TS,
131    ) -> Proof<SC> {
132        let ctx = self.generate_proving_ctx();
133        let device = engine.device();
134        let d_pk = device.transport_pk_to_device(pk);
135        let d_ctx = device.transport_proving_ctx_to_device(&ctx);
136        let mut prover = engine.prover_from_transcript(transcript.clone());
137        let proof = prover.prove(&d_pk, d_ctx).unwrap();
138        *transcript = prover.transcript;
139        proof
140    }
141
142    fn keygen_and_prove<E: StarkEngine<SC = SC>>(
143        &self,
144        engine: &E,
145    ) -> (MultiStarkVerifyingKey<SC>, Proof<SC>) {
146        let (pk, vk) = self.keygen(engine);
147        let proof = self.prove(engine, &pk);
148        (vk, proof)
149    }
150}
151
152pub struct FibFixture {
153    pub a: u64,
154    pub b: u64,
155    pub n: usize,
156    pub num_airs: usize,
157    pub empty_air_indices: Vec<usize>,
158}
159
160impl FibFixture {
161    pub fn new(a: u64, b: u64, n: usize) -> Self {
162        FibFixture {
163            a,
164            b,
165            n,
166            num_airs: 1,
167            empty_air_indices: vec![],
168        }
169    }
170
171    pub fn new_with_num_airs(a: u64, b: u64, n: usize, num_airs: usize) -> Self {
172        FibFixture {
173            a,
174            b,
175            n,
176            num_airs,
177            empty_air_indices: vec![],
178        }
179    }
180
181    pub fn with_empty_air_indices(mut self, empty_air_indices: impl Into<Vec<usize>>) -> Self {
182        self.empty_air_indices = empty_air_indices.into();
183        self
184    }
185}
186
187impl<SC: StarkProtocolConfig> TestFixture<SC> for FibFixture {
188    fn airs(&self) -> Vec<AirRef<SC>> {
189        let air = Arc::new(FibonacciAir);
190        vec![air; self.num_airs]
191    }
192
193    fn generate_proving_ctx(&self) -> ProvingContext<CpuColMajorBackend<SC>> {
194        use crate::test_utils::dummy_airs::fib_air::trace::generate_trace_rows;
195        let f_n = get_fib_number::<SC::F>(self.a, self.b, self.n);
196        let pis = [self.a, self.b, f_n].map(SC::F::from_u64);
197
198        ProvingContext::new(
199            (0..self.num_airs)
200                .filter(|i| !self.empty_air_indices.contains(i))
201                .map(|i| {
202                    (
203                        i,
204                        AirProvingContext::simple(
205                            ColMajorMatrix::from_row_major(&generate_trace_rows::<SC::F>(
206                                self.a, self.b, self.n,
207                            )),
208                            pis.to_vec(),
209                        ),
210                    )
211                })
212                .collect_vec(),
213        )
214    }
215}
216
217/// Interactions fixture with 1 sender and 1 receiver
218pub struct InteractionsFixture11;
219
220impl<SC: StarkProtocolConfig> TestFixture<SC> for InteractionsFixture11 {
221    fn airs(&self) -> Vec<AirRef<SC>> {
222        let sender_air = DummyInteractionAir::new(1, true, 0);
223        let receiver_air = DummyInteractionAir::new(1, false, 0);
224        any_air_arc_vec!(sender_air, receiver_air)
225    }
226
227    fn generate_proving_ctx(&self) -> ProvingContext<CpuColMajorBackend<SC>> {
228        let sender_trace = RowMajorMatrix::new(
229            [0, 1, 3, 5, 7, 4, 546, 889]
230                .into_iter()
231                .map(SC::F::from_usize)
232                .collect(),
233            2,
234        );
235
236        let receiver_trace = RowMajorMatrix::new(
237            [1, 5, 3, 4, 4, 4, 2, 5, 0, 123, 545, 889, 1, 889, 0, 456]
238                .into_iter()
239                .map(SC::F::from_usize)
240                .collect(),
241            2,
242        );
243
244        ProvingContext::new(
245            [sender_trace, receiver_trace]
246                .into_iter()
247                .enumerate()
248                .map(|(air_idx, trace)| {
249                    (
250                        air_idx,
251                        AirProvingContext::simple_no_pis(ColMajorMatrix::from_row_major(&trace)),
252                    )
253                })
254                .collect(),
255        )
256    }
257}
258
259/// Dummy interaction AIRs with cached trace: 1 sender, 1 receiver
260#[derive(derive_new::new)]
261pub struct CachedFixture11<SC> {
262    pub config: SC,
263}
264
265impl<SC: StarkProtocolConfig> TestFixture<SC> for CachedFixture11<SC> {
266    fn airs(&self) -> Vec<AirRef<SC>> {
267        let sender_air = DummyInteractionAir::new(1, true, 0).partition();
268        let receiver_air = DummyInteractionAir::new(1, false, 0).partition();
269        any_air_arc_vec!(sender_air, receiver_air)
270    }
271
272    fn generate_proving_ctx(&self) -> ProvingContext<CpuColMajorBackend<SC>> {
273        let sender_trace = ColMajorMatrix::new(
274            [0, 3, 7, 546].into_iter().map(SC::F::from_usize).collect(),
275            1,
276        );
277        let sender_cached_trace = ColMajorMatrix::new(
278            [1, 5, 4, 889].into_iter().map(SC::F::from_usize).collect(),
279            1,
280        );
281
282        let receiver_trace = ColMajorMatrix::new(
283            [1, 3, 4, 2, 0, 545, 1, 0]
284                .into_iter()
285                .map(SC::F::from_usize)
286                .collect(),
287            1,
288        );
289        let receiver_cached_trace = ColMajorMatrix::new(
290            [5, 4, 4, 5, 123, 889, 889, 456]
291                .into_iter()
292                .map(SC::F::from_usize)
293                .collect(),
294            1,
295        );
296
297        let config = &self.config;
298        let params = config.params();
299        ProvingContext::new(
300            [
301                (sender_trace, sender_cached_trace),
302                (receiver_trace, receiver_cached_trace),
303            ]
304            .map(|(common, cached)| {
305                let (commit, data) = stacked_commit(
306                    config.hasher(),
307                    params.l_skip,
308                    params.n_stack,
309                    params.log_blowup,
310                    params.k_whir(),
311                    &[&cached],
312                )
313                .unwrap();
314                assert_eq!(common.height(), cached.height());
315                let cached_data = CommittedTraceData {
316                    commitment: commit,
317                    trace: cached,
318                    data: Arc::new(data),
319                };
320                AirProvingContext {
321                    cached_mains: vec![cached_data],
322                    common_main: common,
323                    public_values: vec![],
324                }
325            })
326            .into_iter()
327            .enumerate()
328            .collect(),
329        )
330    }
331}
332
333#[derive(derive_new::new)]
334pub struct PreprocessedFibFixture {
335    pub a: u64,
336    pub b: u64,
337    pub sels: Vec<bool>,
338}
339
340impl<SC: StarkProtocolConfig> TestFixture<SC> for PreprocessedFibFixture {
341    fn airs(&self) -> Vec<AirRef<SC>> {
342        let air = Arc::new(FibonacciSelectorAir::new(self.sels.clone(), false));
343        vec![air]
344    }
345
346    fn generate_proving_ctx(&self) -> ProvingContext<CpuColMajorBackend<SC>> {
347        use crate::test_utils::dummy_airs::fib_selector_air::trace::generate_trace_rows;
348        let trace = generate_trace_rows::<SC::F>(self.a, self.b, &self.sels);
349        let f_n = get_conditional_fib_number::<SC::F>(self.a, self.b, &self.sels);
350        let pis = [self.a, self.b, f_n].map(SC::F::from_u64);
351
352        let single_ctx =
353            AirProvingContext::simple(ColMajorMatrix::from_row_major(&trace), pis.to_vec());
354        ProvingContext::new(vec![(0, single_ctx)])
355    }
356}
357
358#[derive(derive_new::new)]
359pub struct PreprocessedAndCachedFixture<SC> {
360    pub sels: Vec<bool>,
361    pub config: SC,
362    pub num_cached_parts: usize,
363}
364
365impl<SC: StarkProtocolConfig> TestFixture<SC> for PreprocessedAndCachedFixture<SC> {
366    fn airs(&self) -> Vec<AirRef<SC>> {
367        vec![Arc::new(PreprocessedCachedAir::new(
368            self.sels.clone(),
369            self.num_cached_parts,
370        ))]
371    }
372
373    fn generate_proving_ctx(&self) -> ProvingContext<CpuColMajorBackend<SC>> {
374        assert!(self.sels.len().is_power_of_two());
375
376        let common_main =
377            ColMajorMatrix::new((0..self.sels.len()).map(SC::F::from_usize).collect(), 1);
378        let cached_mains = (0..self.num_cached_parts)
379            .map(|part| {
380                let trace = ColMajorMatrix::new(
381                    self.sels
382                        .iter()
383                        .enumerate()
384                        .map(|(i, &sel)| {
385                            SC::F::from_usize(i)
386                                + SC::F::from_usize(part + 1) * SC::F::from_bool(sel)
387                        })
388                        .collect(),
389                    1,
390                );
391                commit_cached_trace(&self.config, trace)
392            })
393            .collect();
394
395        ProvingContext::new(vec![(
396            0,
397            AirProvingContext::new(cached_mains, common_main, vec![]),
398        )])
399    }
400}
401#[derive(derive_new::new)]
402pub struct SelfInteractionFixture {
403    pub widths: Vec<usize>,
404    pub log_height: usize,
405    pub bus_index: BusIndex,
406}
407
408impl<SC: StarkProtocolConfig> TestFixture<SC> for SelfInteractionFixture {
409    fn airs(&self) -> Vec<AirRef<SC>> {
410        self.widths
411            .iter()
412            .map(|&width| {
413                Arc::new(SelfInteractionAir {
414                    width,
415                    bus_index: self.bus_index,
416                }) as AirRef<SC>
417            })
418            .collect_vec()
419    }
420
421    fn generate_proving_ctx(&self) -> ProvingContext<CpuColMajorBackend<SC>> {
422        let per_trace = self
423            .widths
424            .iter()
425            .map(|&width| {
426                let chip = SelfInteractionChip {
427                    width,
428                    log_height: self.log_height,
429                };
430                chip.generate_proving_ctx::<SC>()
431            })
432            .enumerate()
433            .collect_vec();
434        ProvingContext { per_trace }
435    }
436}
437
438pub struct MixtureFixture<SC> {
439    pub fxs: Vec<MixtureFixtureEnum<SC>>,
440}
441
442pub enum MixtureFixtureEnum<SC> {
443    FibFixture(FibFixture),
444    InteractionsFixture11(InteractionsFixture11),
445    CachedFixture11(CachedFixture11<SC>),
446    PreprocessedFibFixture(PreprocessedFibFixture),
447    SelfInteractionFixture(SelfInteractionFixture),
448}
449
450impl<SC: StarkProtocolConfig> MixtureFixtureEnum<SC> {
451    fn airs(&self) -> Vec<AirRef<SC>> {
452        use crate::test_utils::MixtureFixtureEnum::*;
453        match self {
454            FibFixture(fx) => fx.airs(),
455            InteractionsFixture11(fx) => fx.airs(),
456            CachedFixture11(fx) => fx.airs(),
457            PreprocessedFibFixture(fx) => fx.airs(),
458            SelfInteractionFixture(fx) => fx.airs(),
459        }
460    }
461
462    fn generate_air_proving_ctxs(&self) -> Vec<AirProvingContext<CpuColMajorBackend<SC>>> {
463        use crate::test_utils::MixtureFixtureEnum::*;
464        let ctx = match self {
465            FibFixture(fx) => fx.generate_proving_ctx(),
466            InteractionsFixture11(fx) => fx.generate_proving_ctx(),
467            CachedFixture11(fx) => fx.generate_proving_ctx(),
468            PreprocessedFibFixture(fx) => fx.generate_proving_ctx(),
469            SelfInteractionFixture(fx) => fx.generate_proving_ctx(),
470        };
471        ctx.per_trace
472            .into_iter()
473            .map(|(_, trace_ctx)| trace_ctx)
474            .collect_vec()
475    }
476}
477
478impl<SC> MixtureFixture<SC> {
479    pub fn new(fxs: Vec<MixtureFixtureEnum<SC>>) -> Self {
480        Self { fxs }
481    }
482
483    pub fn standard(log_height: usize, config: SC) -> Self {
484        let height = 1usize << log_height;
485        let sels = (0..height).map(|i| i % 2 == 0).collect_vec();
486        let widths = vec![4, 7, 8, 8, 10, 100];
487        Self::new(vec![
488            MixtureFixtureEnum::FibFixture(FibFixture::new(8, 8, height)),
489            MixtureFixtureEnum::InteractionsFixture11(InteractionsFixture11),
490            MixtureFixtureEnum::CachedFixture11(CachedFixture11::new(config)),
491            MixtureFixtureEnum::PreprocessedFibFixture(PreprocessedFibFixture::new(7, 3, sels)),
492            MixtureFixtureEnum::SelfInteractionFixture(SelfInteractionFixture::new(
493                widths, log_height, 5,
494            )),
495        ])
496    }
497}
498
499impl<SC: StarkProtocolConfig> TestFixture<SC> for MixtureFixture<SC> {
500    fn airs(&self) -> Vec<AirRef<SC>> {
501        self.fxs.iter().flat_map(|fx| fx.airs()).collect_vec()
502    }
503
504    fn generate_proving_ctx(&self) -> ProvingContext<CpuColMajorBackend<SC>> {
505        let per_trace = self
506            .fxs
507            .iter()
508            .flat_map(|fx| fx.generate_air_proving_ctxs())
509            .enumerate()
510            .collect_vec();
511        ProvingContext { per_trace }
512    }
513}
514
515impl SystemParams {
516    /// Parameters for testing traces of height up to `2^log_trace_height` with **toy security
517    /// parameters** for faster testing.
518    ///
519    /// **These parameters should not be used in production!**
520    pub fn new_for_testing(log_trace_height: usize) -> Self {
521        let l_skip = 4;
522        let k_whir = 4;
523        let mut params = test_system_params_small(l_skip, log_trace_height - l_skip, k_whir);
524        params.max_constraint_degree = 4;
525        params
526    }
527}
528
529/// Trace heights cannot exceed `2^{l_skip + n_stack}` and stacked cells cannot exceed
530/// `w_stack * 2^{l_skip + n_stack}` when using these system params.
531pub fn test_system_params_small(l_skip: usize, n_stack: usize, k_whir: usize) -> SystemParams {
532    let log_final_poly_len = (n_stack + l_skip) % k_whir;
533    test_system_params_small_with_poly_len(l_skip, n_stack, k_whir, log_final_poly_len, 3)
534}
535
536pub fn test_system_params_small_with_poly_len(
537    l_skip: usize,
538    n_stack: usize,
539    k_whir: usize,
540    log_final_poly_len: usize,
541    max_constraint_degree: usize,
542) -> SystemParams {
543    assert!(log_final_poly_len < l_skip + n_stack);
544    let log_blowup = 1;
545    SystemParams {
546        l_skip,
547        n_stack,
548        w_stack: 1 << 12,
549        log_blowup,
550        whir: test_whir_config_small(log_blowup, l_skip + n_stack, k_whir, log_final_poly_len),
551        logup: LogUpSecurityParameters {
552            max_interaction_count: 1 << 30,
553            log_max_message_length: 7,
554            pow_bits: 2,
555        },
556        max_constraint_degree,
557    }
558}
559
560pub fn test_whir_config_small(
561    log_blowup: usize,
562    log_stacked_height: usize,
563    k_whir: usize,
564    log_final_poly_len: usize,
565) -> WhirConfig {
566    let params = WhirParams {
567        k: k_whir,
568        log_final_poly_len,
569        query_phase_pow_bits: 1,
570        folding_pow_bits: 2,
571        mu_pow_bits: 3,
572        proximity: WhirProximityStrategy::SplitUniqueList {
573            m: 3,
574            list_start_round: 1,
575        },
576    };
577    let security_bits = 5;
578    WhirConfig::new(log_blowup, log_stacked_height, params, security_bits)
579}
580
581pub fn default_test_params_small() -> SystemParams {
582    test_system_params_small(2, 8, 3)
583}
openvm_stark_backend/test_utils/mod.rs

openvm_stark_backend/test_utils/
mod.rs
