openvm_poseidon2_air/
babybear.rs

Help
1use std::array::from_fn;
2
3use lazy_static::lazy_static;
4use openvm_stark_backend::p3_field::FieldAlgebra;
5use openvm_stark_sdk::p3_baby_bear::BabyBear;
6use zkhash::{
7    ark_ff::PrimeField as _, fields::babybear::FpBabyBear as HorizenBabyBear,
8    poseidon2::poseidon2_instance_babybear::RC16,
9};
10
11use super::{
12    Poseidon2Constants, BABY_BEAR_POSEIDON2_HALF_FULL_ROUNDS, BABY_BEAR_POSEIDON2_PARTIAL_ROUNDS,
13    POSEIDON2_WIDTH,
14};
15
16pub(crate) fn horizen_to_p3_babybear(horizen_babybear: HorizenBabyBear) -> BabyBear {
17    BabyBear::from_canonical_u64(horizen_babybear.into_bigint().0[0])
18}
19
20pub(crate) fn horizen_round_consts() -> Poseidon2Constants<BabyBear> {
21    let p3_rc16: Vec<Vec<BabyBear>> = RC16
22        .iter()
23        .map(|round| {
24            round
25                .iter()
26                .map(|babybear| horizen_to_p3_babybear(*babybear))
27                .collect()
28        })
29        .collect();
30    let p_end = BABY_BEAR_POSEIDON2_HALF_FULL_ROUNDS + BABY_BEAR_POSEIDON2_PARTIAL_ROUNDS;
31
32    let beginning_full_round_constants: [[BabyBear; POSEIDON2_WIDTH];
33        BABY_BEAR_POSEIDON2_HALF_FULL_ROUNDS] = from_fn(|i| p3_rc16[i].clone().try_into().unwrap());
34    let partial_round_constants: [BabyBear; BABY_BEAR_POSEIDON2_PARTIAL_ROUNDS] =
35        from_fn(|i| p3_rc16[i + BABY_BEAR_POSEIDON2_HALF_FULL_ROUNDS][0]);
36    let ending_full_round_constants: [[BabyBear; POSEIDON2_WIDTH];
37        BABY_BEAR_POSEIDON2_HALF_FULL_ROUNDS] =
38        from_fn(|i| p3_rc16[i + p_end].clone().try_into().unwrap());
39
40    Poseidon2Constants {
41        beginning_full_round_constants,
42        partial_round_constants,
43        ending_full_round_constants,
44    }
45}
46
47lazy_static! {
48    pub static ref BABYBEAR_BEGIN_EXT_CONSTS: [[BabyBear; POSEIDON2_WIDTH]; BABY_BEAR_POSEIDON2_HALF_FULL_ROUNDS] =
49        horizen_round_consts().beginning_full_round_constants;
50    pub static ref BABYBEAR_PARTIAL_CONSTS: [BabyBear; BABY_BEAR_POSEIDON2_PARTIAL_ROUNDS] =
51        horizen_round_consts().partial_round_constants;
52    pub static ref BABYBEAR_END_EXT_CONSTS: [[BabyBear; POSEIDON2_WIDTH]; BABY_BEAR_POSEIDON2_HALF_FULL_ROUNDS] =
53        horizen_round_consts().ending_full_round_constants;
54}
55
56pub(crate) fn babybear_internal_linear_layer<FA: FieldAlgebra, const WIDTH: usize>(
57    state: &mut [FA; WIDTH],
58    int_diag_m1_matrix: &[FA::F; WIDTH],
59) {
60    let sum = state.iter().cloned().sum::<FA>();
61    for (input, diag_m1) in state.iter_mut().zip(int_diag_m1_matrix) {
62        *input = sum.clone() + FA::from_f(*diag_m1) * input.clone();
63    }
64}
openvm_poseidon2_air/babybear.rs

openvm_poseidon2_air/
babybear.rs
