openvm_stark_sdk/config/
baby_bear_poseidon2.rs

1use std::{
2    io::{self, Read, Write},
3    marker::PhantomData,
4    sync::OnceLock,
5};
6
7use openvm_stark_backend::{
8    codec::{
9        decode_extension_field32, decode_prime_field32, encode_extension_field32,
10        encode_prime_field32, DecodableConfig, EncodableConfig,
11    },
12    hasher::Hasher,
13    p3_symmetric::{PaddingFreeSponge, Permutation, TruncatedPermutation},
14    prover::{Coordinator, CpuColMajorBackend, ReferenceDevice},
15    transcript::duplex_sponge,
16    FiatShamirTranscript, StarkEngine, StarkProtocolConfig, SystemParams, TranscriptLog,
17};
18use p3_baby_bear::{default_babybear_poseidon2_16, BabyBear, Poseidon2BabyBear};
19use p3_field::{extension::BinomialExtensionField, PrimeCharacteristicRing};
20
21const RATE: usize = 8;
22/// permutation width
23const WIDTH: usize = 16; // rate + capacity
24pub const CHUNK: usize = 8;
25pub const DIGEST_SIZE: usize = CHUNK;
26
27type Perm = Poseidon2BabyBear<WIDTH>;
28// Generic over P: CryptographicPermutation<[F; WIDTH]>
29type Hash<P> = PaddingFreeSponge<P, WIDTH, RATE, DIGEST_SIZE>;
30type Compress<P> = TruncatedPermutation<P, 2, CHUNK, WIDTH>;
31type PermHasher<P> = Hasher<F, Digest, Hash<P>, Compress<P>>;
32// Defined below
33type SC = BabyBearPoseidon2Config;
34
35// Convenience type aliases
36pub type F = BabyBear;
37pub type EF = BinomialExtensionField<BabyBear, 4>;
38pub const D_EF: usize = 4;
39pub type Digest = [F; DIGEST_SIZE];
40pub type DuplexSponge = duplex_sponge::DuplexSponge<F, Perm, WIDTH, RATE>;
41pub type DuplexSpongeRecorder = duplex_sponge::DuplexSpongeRecorder<F, Perm, WIDTH, RATE>;
42pub type DuplexSpongeValidator = duplex_sponge::DuplexSpongeValidator<F, Perm, WIDTH, RATE>;
43
44#[derive(Clone, Debug, derive_new::new)]
45pub struct BabyBearPoseidon2Config {
46    params: SystemParams,
47    hasher: PermHasher<Perm>,
48}
49
50impl StarkProtocolConfig for BabyBearPoseidon2Config {
51    type F = F;
52    type EF = EF;
53    type Digest = Digest;
54    type Hasher = PermHasher<Perm>;
55
56    fn params(&self) -> &SystemParams {
57        &self.params
58    }
59
60    fn hasher(&self) -> &Self::Hasher {
61        &self.hasher
62    }
63}
64
65impl BabyBearPoseidon2Config {
66    pub fn new_from_perm(params: SystemParams, perm: Perm) -> Self {
67        let hasher = Hasher::new(
68            PaddingFreeSponge::new(perm.clone()),
69            TruncatedPermutation::new(perm),
70        );
71        Self { params, hasher }
72    }
73
74    pub fn default_from_params(params: SystemParams) -> Self {
75        let perm = default_babybear_poseidon2_16();
76        Self::new_from_perm(params, perm)
77    }
78}
79
80impl EncodableConfig for BabyBearPoseidon2Config {
81    fn encode_base_field<W: Write>(val: &F, writer: &mut W) -> io::Result<()> {
82        encode_prime_field32(val, writer)
83    }
84
85    fn encode_extension_field<W: Write>(val: &EF, writer: &mut W) -> io::Result<()> {
86        encode_extension_field32::<Self::F, _, _>(val, writer)
87    }
88
89    fn encode_digest<W: Write>(digest: &Self::Digest, writer: &mut W) -> io::Result<()> {
90        for val in digest {
91            encode_prime_field32(val, writer)?;
92        }
93        Ok(())
94    }
95}
96
97impl DecodableConfig for BabyBearPoseidon2Config {
98    fn decode_base_field<R: Read>(reader: &mut R) -> io::Result<F> {
99        decode_prime_field32(reader)
100    }
101
102    fn decode_extension_field<R: Read>(reader: &mut R) -> io::Result<EF> {
103        decode_extension_field32::<F, _, _>(reader)
104    }
105
106    fn decode_digest<R: Read>(reader: &mut R) -> io::Result<Digest> {
107        let mut result = Digest::default();
108        for val in &mut result {
109            *val = decode_prime_field32(reader)?;
110        }
111        Ok(result)
112    }
113}
114
115pub struct BabyBearPoseidon2RefEngine<TS = DuplexSponge> {
116    device: ReferenceDevice<SC>,
117    _transcript: PhantomData<TS>,
118}
119
120impl<TS> StarkEngine for BabyBearPoseidon2RefEngine<TS>
121where
122    TS: FiatShamirTranscript<SC> + From<Perm>,
123{
124    type SC = SC;
125    type PB = CpuColMajorBackend<SC>;
126    type PD = ReferenceDevice<SC>;
127    type TS = TS;
128
129    fn new(params: SystemParams) -> Self {
130        let config = BabyBearPoseidon2Config::default_from_params(params);
131        Self {
132            device: ReferenceDevice::new(config),
133            _transcript: PhantomData,
134        }
135    }
136
137    fn config(&self) -> &SC {
138        self.device.config()
139    }
140
141    fn device(&self) -> &Self::PD {
142        &self.device
143    }
144
145    fn initial_transcript(&self) -> Self::TS {
146        TS::from(default_babybear_poseidon2_16())
147    }
148
149    fn prover_from_transcript(
150        &self,
151        transcript: TS,
152    ) -> Coordinator<Self::SC, Self::PB, Self::PD, Self::TS> {
153        Coordinator::new(CpuColMajorBackend::new(), self.device.clone(), transcript)
154    }
155}
156
157// ---- Optimized CPU engine (behind `cpu-backend` feature) ----
158
159#[cfg(feature = "cpu-backend")]
160mod cpu_engine {
161    use openvm_cpu_backend::{CpuBackend, CpuDevice};
162
163    use super::*;
164
165    /// SIMD-optimized transcript backed by Plonky3's `DuplexChallenger`.
166    ///
167    /// Produces identical Fiat-Shamir challenges as the standard `DuplexSponge` transcript,
168    /// but with ~4x faster proof-of-work grinding on aarch64 NEON (8x on x86 AVX2).
169    #[derive(Clone, Debug)]
170    pub struct CpuTranscript {
171        inner: openvm_stark_backend::p3_challenger::DuplexChallenger<
172            BabyBear,
173            Poseidon2BabyBear<WIDTH>,
174            WIDTH,
175            RATE,
176        >,
177    }
178
179    impl From<Perm> for CpuTranscript {
180        fn from(perm: Perm) -> Self {
181            Self {
182                inner: openvm_stark_backend::p3_challenger::DuplexChallenger::new(perm),
183            }
184        }
185    }
186
187    impl FiatShamirTranscript<BabyBearPoseidon2Config> for CpuTranscript {
188        #[inline]
189        fn observe(&mut self, value: BabyBear) {
190            openvm_stark_backend::p3_challenger::CanObserve::observe(&mut self.inner, value);
191        }
192
193        #[inline]
194        fn sample(&mut self) -> BabyBear {
195            openvm_stark_backend::p3_challenger::CanSample::sample(&mut self.inner)
196        }
197
198        fn observe_commit(&mut self, digest: [BabyBear; RATE]) {
199            for x in digest {
200                openvm_stark_backend::p3_challenger::CanObserve::observe(&mut self.inner, x);
201            }
202        }
203
204        fn grind(&mut self, bits: usize) -> BabyBear {
205            openvm_stark_backend::p3_challenger::GrindingChallenger::grind(&mut self.inner, bits)
206        }
207    }
208
209    /// Row-major CPU engine for BabyBear + Poseidon2.
210    ///
211    /// Default transcript is [`CpuTranscript`], which uses Plonky3's `DuplexChallenger`
212    /// for SIMD-optimized proof-of-work grinding (~4x faster on NEON, ~8x on AVX2).
213    pub struct BabyBearPoseidon2CpuEngine<TS = CpuTranscript> {
214        device: CpuDevice<SC>,
215        _transcript: PhantomData<TS>,
216    }
217
218    impl<TS> StarkEngine for BabyBearPoseidon2CpuEngine<TS>
219    where
220        TS: FiatShamirTranscript<SC> + From<Poseidon2BabyBear<WIDTH>>,
221    {
222        type SC = SC;
223        type PB = CpuBackend<SC>;
224        type PD = CpuDevice<SC>;
225        type TS = TS;
226
227        fn new(params: SystemParams) -> Self {
228            let config = BabyBearPoseidon2Config::default_from_params(params);
229            Self {
230                device: CpuDevice::new(config),
231                _transcript: PhantomData,
232            }
233        }
234
235        fn config(&self) -> &SC {
236            self.device.config()
237        }
238
239        fn device(&self) -> &Self::PD {
240            &self.device
241        }
242
243        fn initial_transcript(&self) -> Self::TS {
244            TS::from(default_babybear_poseidon2_16())
245        }
246
247        fn prover_from_transcript(
248            &self,
249            transcript: TS,
250        ) -> Coordinator<Self::SC, Self::PB, Self::PD, Self::TS> {
251            Coordinator::new(CpuBackend::new(), self.device.clone(), transcript)
252        }
253    }
254}
255
256#[cfg(feature = "cpu-backend")]
257pub use cpu_engine::{BabyBearPoseidon2CpuEngine, CpuTranscript};
258
259// Fixed Poseidon2 configuration
260pub fn poseidon2_perm() -> &'static Poseidon2BabyBear<WIDTH> {
261    static PERM: OnceLock<Poseidon2BabyBear<WIDTH>> = OnceLock::new();
262    PERM.get_or_init(default_babybear_poseidon2_16)
263}
264
265pub fn poseidon2_compress_with_capacity(
266    left: [F; CHUNK],
267    right: [F; CHUNK],
268) -> ([F; CHUNK], [F; CHUNK]) {
269    let mut state = [F::ZERO; WIDTH];
270    state[..CHUNK].copy_from_slice(&left);
271    state[CHUNK..].copy_from_slice(&right);
272    poseidon2_perm().permute_mut(&mut state);
273    (
274        state[..CHUNK].try_into().unwrap(),
275        state[CHUNK..].try_into().unwrap(),
276    )
277}
278
279pub fn default_duplex_sponge() -> DuplexSponge {
280    DuplexSponge::from(poseidon2_perm().clone())
281}
282
283pub fn default_duplex_sponge_recorder() -> DuplexSpongeRecorder {
284    DuplexSpongeRecorder::from(poseidon2_perm().clone())
285}
286
287pub fn default_duplex_sponge_validator(
288    logs: TranscriptLog<F, [F; WIDTH]>,
289) -> DuplexSpongeValidator {
290    DuplexSpongeValidator::new(poseidon2_perm().clone(), logs)
291}
292
293#[cfg(test)]
294mod poseidon2_constant_tests {
295    use p3_baby_bear::{
296        BABYBEAR_RC16_EXTERNAL_FINAL, BABYBEAR_RC16_EXTERNAL_INITIAL, BABYBEAR_RC16_INTERNAL,
297    };
298    use zkhash::{
299        ark_ff::PrimeField as _, fields::babybear::FpBabyBear as HorizenBabyBear,
300        poseidon2::poseidon2_instance_babybear::RC16,
301    };
302
303    use super::*;
304
305    fn horizen_to_p3(horizen_babybear: HorizenBabyBear) -> BabyBear {
306        BabyBear::from_u64(horizen_babybear.into_bigint().0[0])
307    }
308
309    #[allow(clippy::type_complexity)]
310    pub fn horizen_round_consts_16() -> ((Vec<[BabyBear; 16]>, Vec<[BabyBear; 16]>), Vec<BabyBear>)
311    {
312        let p3_rc16: Vec<Vec<BabyBear>> = RC16
313            .iter()
314            .map(|round| {
315                round
316                    .iter()
317                    .map(|babybear| horizen_to_p3(*babybear))
318                    .collect()
319            })
320            .collect();
321
322        let rounds_f = 8;
323        let rounds_p = 13;
324        let rounds_f_beginning = rounds_f / 2;
325        let p_end = rounds_f_beginning + rounds_p;
326        let initial: Vec<[BabyBear; 16]> = p3_rc16[..rounds_f_beginning]
327            .iter()
328            .cloned()
329            .map(|round| round.try_into().unwrap())
330            .collect();
331        let terminal: Vec<[BabyBear; 16]> = p3_rc16[p_end..]
332            .iter()
333            .cloned()
334            .map(|round| round.try_into().unwrap())
335            .collect();
336        let internal_round_constants: Vec<BabyBear> = p3_rc16[rounds_f_beginning..p_end]
337            .iter()
338            .map(|round| round[0])
339            .collect();
340        ((initial, terminal), internal_round_constants)
341    }
342
343    /// Uses HorizenLabs Poseidon2 round constants, but plonky3 Mat4 and also
344    /// with a p3 Monty reduction factor.
345    #[test]
346    fn test_horizen_p3_rc_equality() {
347        let ((external_initial, external_terminal), internal_constants) = horizen_round_consts_16();
348        assert_eq!(external_initial, BABYBEAR_RC16_EXTERNAL_INITIAL.to_vec());
349        assert_eq!(external_terminal, BABYBEAR_RC16_EXTERNAL_FINAL.to_vec());
350        assert_eq!(internal_constants, BABYBEAR_RC16_INTERNAL.to_vec());
351    }
352}
openvm_stark_sdk/config/baby_bear_poseidon2.rs

openvm_stark_sdk/config/
baby_bear_poseidon2.rs
