openvm_stark_sdk/
utils.rs

1use std::{cmp::Reverse, iter::zip};
2
3use itertools::Itertools;
4use openvm_stark_backend::{
5    config::StarkGenericConfig,
6    p3_field::FieldAlgebra,
7    p3_matrix::Matrix,
8    prover::{
9        cpu::{CpuBackend, CpuDevice},
10        types::AirProvingContext,
11    },
12    verifier::VerificationError,
13    AirRef,
14};
15use rand::{rngs::StdRng, Rng, SeedableRng};
16
17use crate::engine::{StarkFriEngine, VerificationDataWithFriParams};
18
19/// `stark-backend::prover::types::ProofInput` without specifying AIR IDs.
20pub struct ProofInputForTest<SC: StarkGenericConfig> {
21    pub airs: Vec<AirRef<SC>>,
22    pub per_air: Vec<AirProvingContext<CpuBackend<SC>>>,
23}
24
25impl<SC: StarkGenericConfig> ProofInputForTest<SC> {
26    pub fn run_test(
27        self,
28        engine: &impl StarkFriEngine<SC = SC, PB = CpuBackend<SC>, PD = CpuDevice<SC>>,
29    ) -> Result<VerificationDataWithFriParams<SC>, VerificationError> {
30        assert_eq!(self.airs.len(), self.per_air.len());
31        engine.run_test(self.airs, self.per_air)
32    }
33    /// Sort AIRs by their trace height in descending order. This should not be used outside
34    /// static-verifier because a dynamic verifier should support any AIR order.
35    /// This is related to an implementation detail of FieldMerkleTreeMMCS which is used in most
36    /// configs. Reference: <https://github.com/Plonky3/Plonky3/blob/27b3127dab047e07145c38143379edec2960b3e1/merkle-tree/src/merkle_tree.rs#L53>
37    pub fn sort_chips(&mut self) {
38        let airs = std::mem::take(&mut self.airs);
39        let air_proving_ctxs = std::mem::take(&mut self.per_air);
40        let (airs, air_proof_inputs): (Vec<_>, Vec<_>) = zip(airs, air_proving_ctxs)
41            .sorted_by_key(|(_, air_proving_ctx)| {
42                Reverse(
43                    air_proving_ctx
44                        .common_main
45                        .as_ref()
46                        .map(|trace| trace.height())
47                        .unwrap_or(0),
48                )
49            })
50            .unzip();
51        self.airs = airs;
52        self.per_air = air_proof_inputs;
53    }
54}
55
56/// Deterministic seeded RNG, for testing use
57pub fn create_seeded_rng() -> StdRng {
58    let seed = [42; 32];
59    StdRng::from_seed(seed)
60}
61
62pub fn create_seeded_rng_with_seed(seed: u64) -> StdRng {
63    let seed_be = seed.to_be_bytes();
64    let mut seed = [0u8; 32];
65    seed[24..32].copy_from_slice(&seed_be);
66    StdRng::from_seed(seed)
67}
68
69// Returns row major matrix
70pub fn generate_random_matrix<F: FieldAlgebra>(
71    mut rng: impl Rng,
72    height: usize,
73    width: usize,
74) -> Vec<Vec<F>> {
75    (0..height)
76        .map(|_| {
77            (0..width)
78                .map(|_| F::from_wrapped_u32(rng.gen()))
79                .collect_vec()
80        })
81        .collect_vec()
82}
83
84pub fn to_field_vec<F: FieldAlgebra>(v: Vec<u32>) -> Vec<F> {
85    v.into_iter().map(F::from_canonical_u32).collect()
86}
87
88/// A macro to create a `Vec<Arc<dyn AnyRap<_>>>` from a list of AIRs because Rust cannot infer the
89/// type correctly when using `vec!`.
90#[macro_export]
91macro_rules! any_rap_arc_vec {
92    [$($e:expr),*] => {
93        {
94            let chips: Vec<std::sync::Arc<dyn openvm_stark_backend::rap::AnyRap<_>>> = vec![$(std::sync::Arc::new($e)),*];
95            chips
96        }
97    };
98}
99
100#[macro_export]
101macro_rules! assert_sc_compatible_with_serde {
102    ($sc:ty) => {
103        static_assertions::assert_impl_all!(openvm_stark_backend::keygen::types::MultiStarkProvingKey<$sc>: serde::Serialize, serde::de::DeserializeOwned);
104        static_assertions::assert_impl_all!(openvm_stark_backend::keygen::types::MultiStarkVerifyingKey<$sc>: serde::Serialize, serde::de::DeserializeOwned);
105        static_assertions::assert_impl_all!(openvm_stark_backend::proof::Proof<$sc>: serde::Serialize, serde::de::DeserializeOwned);
106    };
107}