openvm_stark_sdk/config/
baby_bear_poseidon2_root.rs

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

openvm_stark_sdk/config/
baby_bear_poseidon2_root.rs
