openvm_poseidon2_air/
babybear.rs

1use std::array::from_fn;
2
3use lazy_static::lazy_static;
4use openvm_stark_backend::p3_field::{
5    integers::QuotientMap, PrimeCharacteristicRing, PrimeField32,
6};
7use openvm_stark_sdk::p3_baby_bear::BabyBear;
8use zkhash::{
9    ark_ff::PrimeField as _, fields::babybear::FpBabyBear as HorizenBabyBear,
10    poseidon2::poseidon2_instance_babybear::RC16,
11};
12
13use super::{
14    Poseidon2Constants, BABY_BEAR_POSEIDON2_HALF_FULL_ROUNDS, BABY_BEAR_POSEIDON2_PARTIAL_ROUNDS,
15    POSEIDON2_WIDTH,
16};
17
18pub(crate) fn horizen_to_p3_babybear(horizen_babybear: HorizenBabyBear) -> BabyBear {
19    BabyBear::from_u64(horizen_babybear.into_bigint().0[0])
20}
21
22pub(crate) fn horizen_round_consts() -> Poseidon2Constants<BabyBear> {
23    let p3_rc16: Vec<Vec<BabyBear>> = RC16
24        .iter()
25        .map(|round| {
26            round
27                .iter()
28                .map(|babybear| horizen_to_p3_babybear(*babybear))
29                .collect()
30        })
31        .collect();
32    let p_end = BABY_BEAR_POSEIDON2_HALF_FULL_ROUNDS + BABY_BEAR_POSEIDON2_PARTIAL_ROUNDS;
33
34    let beginning_full_round_constants: [[BabyBear; POSEIDON2_WIDTH];
35        BABY_BEAR_POSEIDON2_HALF_FULL_ROUNDS] = from_fn(|i| p3_rc16[i].clone().try_into().unwrap());
36    let partial_round_constants: [BabyBear; BABY_BEAR_POSEIDON2_PARTIAL_ROUNDS] =
37        from_fn(|i| p3_rc16[i + BABY_BEAR_POSEIDON2_HALF_FULL_ROUNDS][0]);
38    let ending_full_round_constants: [[BabyBear; POSEIDON2_WIDTH];
39        BABY_BEAR_POSEIDON2_HALF_FULL_ROUNDS] =
40        from_fn(|i| p3_rc16[i + p_end].clone().try_into().unwrap());
41
42    Poseidon2Constants {
43        beginning_full_round_constants,
44        partial_round_constants,
45        ending_full_round_constants,
46    }
47}
48
49lazy_static! {
50    pub static ref BABYBEAR_BEGIN_EXT_CONSTS: [[BabyBear; POSEIDON2_WIDTH]; BABY_BEAR_POSEIDON2_HALF_FULL_ROUNDS] =
51        horizen_round_consts().beginning_full_round_constants;
52    pub static ref BABYBEAR_PARTIAL_CONSTS: [BabyBear; BABY_BEAR_POSEIDON2_PARTIAL_ROUNDS] =
53        horizen_round_consts().partial_round_constants;
54    pub static ref BABYBEAR_END_EXT_CONSTS: [[BabyBear; POSEIDON2_WIDTH]; BABY_BEAR_POSEIDON2_HALF_FULL_ROUNDS] =
55        horizen_round_consts().ending_full_round_constants;
56}
57
58/// The vector `[-2, 1, 2, 1/2, 3, 4, -1/2, -3, -4, 1/2^8, 1/4, 1/8, 1/2^27, -1/2^8, -1/16,
59/// -1/2^27]` saved as an array of BabyBear elements. Copied from plonky3's Poseidon2
60/// implementation to preserve the exact constraint structure.
61pub const INTERNAL_DIAG_MONTY_16: [BabyBear; 16] = BabyBear::new_array([
62    BabyBear::ORDER_U32 - 2,
63    1,
64    2,
65    (BabyBear::ORDER_U32 + 1) >> 1,
66    3,
67    4,
68    (BabyBear::ORDER_U32 - 1) >> 1,
69    BabyBear::ORDER_U32 - 3,
70    BabyBear::ORDER_U32 - 4,
71    BabyBear::ORDER_U32 - ((BabyBear::ORDER_U32 - 1) >> 8),
72    BabyBear::ORDER_U32 - ((BabyBear::ORDER_U32 - 1) >> 2),
73    BabyBear::ORDER_U32 - ((BabyBear::ORDER_U32 - 1) >> 3),
74    BabyBear::ORDER_U32 - 15,
75    (BabyBear::ORDER_U32 - 1) >> 8,
76    (BabyBear::ORDER_U32 - 1) >> 4,
77    15,
78]);
79
80pub(crate) fn babybear_internal_linear_layer<R: PrimeCharacteristicRing>(
81    state: &mut [R; POSEIDON2_WIDTH],
82    diag_m1_matrix: &[BabyBear; POSEIDON2_WIDTH],
83) {
84    let sum: R = state.iter().cloned().sum();
85    for (val, &diag_elem) in state.iter_mut().zip(diag_m1_matrix.iter()) {
86        let diag_sub = R::PrimeSubfield::from_int(diag_elem.as_canonical_u32());
87        let diag_r = R::from_prime_subfield(diag_sub);
88        *val = sum.clone() + diag_r * val.clone();
89    }
90}
openvm_poseidon2_air/babybear.rs

openvm_poseidon2_air/
babybear.rs
