openvm_stark_sdk/config/
baby_bear_poseidon2_root.rs

1use ff::PrimeField;
2use openvm_stark_backend::{
3    config::StarkConfig, interaction::fri_log_up::FriLogUpPhase,
4    p3_challenger::MultiField32Challenger, p3_commit::ExtensionMmcs,
5    p3_field::extension::BinomialExtensionField,
6};
7use p3_baby_bear::BabyBear;
8use p3_bn254_fr::{Bn254Fr, FFBn254Fr, Poseidon2Bn254};
9use p3_dft::Radix2DitParallel;
10use p3_fri::{FriConfig, TwoAdicFriPcs};
11use p3_merkle_tree::MerkleTreeMmcs;
12use p3_poseidon2::ExternalLayerConstants;
13use p3_symmetric::{CryptographicPermutation, MultiField32PaddingFreeSponge, TruncatedPermutation};
14use zkhash::{
15    ark_ff::{BigInteger, PrimeField as _},
16    fields::bn256::FpBN256 as ark_FpBN256,
17    poseidon2::poseidon2_instance_bn256::RC3,
18};
19
20use super::FriParameters;
21use crate::{
22    assert_sc_compatible_with_serde,
23    config::{
24        fri_params::SecurityParameters, log_up_params::log_up_security_params_baby_bear_100_bits,
25    },
26    engine::{StarkEngine, StarkFriEngine},
27};
28
29const WIDTH: usize = 3;
30/// Poseidon rate in F. <Poseidon RATE>(2) * <# of F in a N>(8) = 16
31const RATE: usize = 16;
32const DIGEST_WIDTH: usize = 1;
33
34/// A configuration for  recursion.
35type Val = BabyBear;
36type Challenge = BinomialExtensionField<Val, 4>;
37type Perm = Poseidon2Bn254<WIDTH>;
38type Hash<P> = MultiField32PaddingFreeSponge<Val, Bn254Fr, P, WIDTH, RATE, DIGEST_WIDTH>;
39type Compress<P> = TruncatedPermutation<P, 2, 1, WIDTH>;
40type ValMmcs<P> = MerkleTreeMmcs<BabyBear, Bn254Fr, Hash<P>, Compress<P>, 1>;
41type ChallengeMmcs<P> = ExtensionMmcs<Val, Challenge, ValMmcs<P>>;
42type Dft = Radix2DitParallel<Val>;
43type Challenger<P> = MultiField32Challenger<Val, Bn254Fr, P, WIDTH, 2>;
44type Pcs<P> = TwoAdicFriPcs<Val, Dft, ValMmcs<P>, ChallengeMmcs<P>>;
45type RapPhase<P> = FriLogUpPhase<Val, Challenge, Challenger<P>>;
46
47pub type BabyBearPermutationRootConfig<P> =
48    StarkConfig<Pcs<P>, RapPhase<P>, Challenge, Challenger<P>>;
49pub type BabyBearPoseidon2RootConfig = BabyBearPermutationRootConfig<Perm>;
50pub type BabyBearPoseidon2RootEngine = BabyBearPermutationRootEngine<Perm>;
51
52assert_sc_compatible_with_serde!(BabyBearPoseidon2RootConfig);
53
54pub struct BabyBearPermutationRootEngine<P>
55where
56    P: CryptographicPermutation<[Bn254Fr; WIDTH]> + Clone,
57{
58    pub fri_params: FriParameters,
59    pub config: BabyBearPermutationRootConfig<P>,
60    pub perm: P,
61    pub max_constraint_degree: usize,
62}
63
64impl<P> StarkEngine<BabyBearPermutationRootConfig<P>> for BabyBearPermutationRootEngine<P>
65where
66    P: CryptographicPermutation<[Bn254Fr; WIDTH]> + Clone,
67{
68    fn config(&self) -> &BabyBearPermutationRootConfig<P> {
69        &self.config
70    }
71
72    fn max_constraint_degree(&self) -> Option<usize> {
73        Some(self.max_constraint_degree)
74    }
75
76    fn new_challenger(&self) -> Challenger<P> {
77        Challenger::new(self.perm.clone()).unwrap()
78    }
79}
80
81/// `pcs_log_degree` is the upper bound on the log_2(PCS polynomial degree).
82pub fn default_engine() -> BabyBearPoseidon2RootEngine {
83    default_engine_impl(SecurityParameters::standard_fast())
84}
85
86/// `pcs_log_degree` is the upper bound on the log_2(PCS polynomial degree).
87fn default_engine_impl(security_params: SecurityParameters) -> BabyBearPoseidon2RootEngine {
88    let perm = root_perm();
89    engine_from_perm(perm, security_params)
90}
91
92/// `pcs_log_degree` is the upper bound on the log_2(PCS polynomial degree).
93pub fn default_config(perm: &Perm) -> BabyBearPoseidon2RootConfig {
94    config_from_perm(perm, SecurityParameters::standard_fast())
95}
96
97pub fn engine_from_perm<P>(
98    perm: P,
99    security_params: SecurityParameters,
100) -> BabyBearPermutationRootEngine<P>
101where
102    P: CryptographicPermutation<[Bn254Fr; WIDTH]> + Clone,
103{
104    let fri_params = security_params.fri_params;
105    let max_constraint_degree = fri_params.max_constraint_degree();
106    let config = config_from_perm(&perm, security_params);
107    BabyBearPermutationRootEngine {
108        config,
109        perm,
110        fri_params,
111        max_constraint_degree,
112    }
113}
114
115pub fn config_from_perm<P>(
116    perm: &P,
117    security_params: SecurityParameters,
118) -> BabyBearPermutationRootConfig<P>
119where
120    P: CryptographicPermutation<[Bn254Fr; WIDTH]> + Clone,
121{
122    let hash = Hash::new(perm.clone()).unwrap();
123    let compress = Compress::new(perm.clone());
124    let val_mmcs = ValMmcs::new(hash, compress);
125    let challenge_mmcs = ChallengeMmcs::new(val_mmcs.clone());
126    let dft = Dft::default();
127    let SecurityParameters {
128        fri_params,
129        log_up_params,
130    } = security_params;
131    let fri_config = FriConfig {
132        log_blowup: fri_params.log_blowup,
133        log_final_poly_len: fri_params.log_final_poly_len,
134        num_queries: fri_params.num_queries,
135        proof_of_work_bits: fri_params.proof_of_work_bits,
136        mmcs: challenge_mmcs,
137    };
138    let pcs = Pcs::new(dft, val_mmcs, fri_config);
139    let rap_phase = FriLogUpPhase::new(log_up_params);
140    BabyBearPermutationRootConfig::new(pcs, rap_phase)
141}
142
143/// The permutation for outer recursion.
144pub fn root_perm() -> Perm {
145    const ROUNDS_F: usize = 8;
146    const ROUNDS_P: usize = 56;
147    let mut round_constants = bn254_poseidon2_rc3();
148    let internal_end = (ROUNDS_F / 2) + ROUNDS_P;
149    let terminal = round_constants.split_off(internal_end);
150    let internal_round_constants = round_constants.split_off(ROUNDS_F / 2);
151    let internal_round_constants = internal_round_constants
152        .into_iter()
153        .map(|vec| vec[0])
154        .collect::<Vec<_>>();
155    let initial = round_constants;
156
157    let external_round_constants = ExternalLayerConstants::new(initial, terminal);
158    Perm::new(external_round_constants, internal_round_constants)
159}
160
161fn bn254_from_ark_ff(input: ark_FpBN256) -> Bn254Fr {
162    let bytes = input.into_bigint().to_bytes_le();
163
164    let mut res = <FFBn254Fr as ff::PrimeField>::Repr::default();
165
166    for (i, digit) in res.as_mut().iter_mut().enumerate() {
167        *digit = bytes[i];
168    }
169
170    let value = FFBn254Fr::from_repr(res);
171
172    if value.is_some().into() {
173        Bn254Fr {
174            value: value.unwrap(),
175        }
176    } else {
177        panic!("Invalid field element")
178    }
179}
180
181fn bn254_poseidon2_rc3() -> Vec<[Bn254Fr; 3]> {
182    RC3.iter()
183        .map(|vec| {
184            vec.iter()
185                .cloned()
186                .map(bn254_from_ark_ff)
187                .collect::<Vec<_>>()
188                .try_into()
189                .unwrap()
190        })
191        .collect()
192}
193
194impl StarkFriEngine<BabyBearPoseidon2RootConfig> for BabyBearPoseidon2RootEngine {
195    fn new(fri_params: FriParameters) -> Self {
196        let security_params = SecurityParameters {
197            fri_params,
198            log_up_params: log_up_security_params_baby_bear_100_bits(),
199        };
200        default_engine_impl(security_params)
201    }
202    fn fri_params(&self) -> FriParameters {
203        self.fri_params
204    }
205}