openvm_stark_sdk/config/
baby_bear_poseidon2_root.rs

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

openvm_stark_sdk/config/
baby_bear_poseidon2_root.rs
