1use ark_crypto_primitives::sponge::{
 2    poseidon::{PoseidonConfig, PoseidonSponge},
 3    Absorb, CryptographicSponge,
 4};
 5use ark_ec::{AffineRepr, CurveGroup, Group};
 6use ark_ff::{BigInteger, Field, PrimeField};
 7
 8pub struct PoseidonTranscript<C: CurveGroup>
 9where
10    <C as Group>::ScalarField: Absorb,
11{
12    sponge: PoseidonSponge<C::ScalarField>,
13}
14
15impl<C: CurveGroup> PoseidonTranscript<C>
16where
17    <C as Group>::ScalarField: Absorb,
18{
19    pub fn new(poseidon_config: &PoseidonConfig<C::ScalarField>) -> Self {
20        let sponge = PoseidonSponge::<C::ScalarField>::new(poseidon_config);
21        Self { sponge }
22    }
23    pub fn absorb(&mut self, v: &C::ScalarField) {
24        self.sponge.absorb(&v);
25    }
26    pub fn absorb_point(&mut self, p: &C) {
27        self.sponge.absorb(&prepare_point(p));
28    }
29    pub fn get_challenge(&mut self) -> C::ScalarField {
30        let c = self.sponge.squeeze_field_elements(1);
31        self.sponge.absorb(&c[0]);
32        c[0]
33    }
34}
35
36// Returns the point coordinates in Fr, so it can be absrobed by the transcript. It does not work
37// over bytes in order to have a logic that can be reproduced in-circuit.
38fn prepare_point<C: CurveGroup>(p: &C) -> Vec<C::ScalarField> {
39    let binding = p.into_affine();
40    let p_coords = &binding.xy().unwrap();
41    let x_bi = p_coords
42        .0
43        .to_base_prime_field_elements()
44        .next()
45        .expect("a")
46        .into_bigint();
47    let y_bi = p_coords
48        .1
49        .to_base_prime_field_elements()
50        .next()
51        .expect("a")
52        .into_bigint();
53    vec![
54        C::ScalarField::from_le_bytes_mod_order(x_bi.to_bytes_le().as_ref()),
55        C::ScalarField::from_le_bytes_mod_order(y_bi.to_bytes_le().as_ref()),
56    ]
57}
58
59#[cfg(test)]
60pub mod tests {
61    use super::*;
62    use ark_crypto_primitives::sponge::poseidon::find_poseidon_ark_and_mds;
63
64    /// WARNING the method poseidon_test_config is for tests only
65    #[cfg(test)]
66    pub fn poseidon_test_config<F: PrimeField>() -> PoseidonConfig<F> {
67        let full_rounds = 8;
68        let partial_rounds = 31;
69        let alpha = 5;
70        let rate = 2;
71
72        let (ark, mds) = find_poseidon_ark_and_mds::<F>(
73            F::MODULUS_BIT_SIZE as u64,
74            rate,
75            full_rounds,
76            partial_rounds,
77            0,
78        );
79
80        PoseidonConfig::new(
81            full_rounds as usize,
82            partial_rounds as usize,
83            alpha,
84            mds,
85            ark,
86            rate,
87            1,
88        )
89    }
90}