openvm_stark_sdk/
utils.rs

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