1use ark_crypto_primitives::sponge::{
  2    constraints::CryptographicSpongeVar,
  3    poseidon::{constraints::PoseidonSpongeVar, PoseidonConfig, PoseidonSponge},
  4    Absorb, CryptographicSponge,
  5};
  6use ark_ec::{AffineRepr, CurveGroup, Group};
  7use ark_ff::{BigInteger, Field, PrimeField};
  8use ark_r1cs_std::fields::fp::FpVar;
  9use ark_relations::r1cs::{ConstraintSystemRef, SynthesisError};
 10
 11use crate::transcript::Transcript;
 12
 13/// PoseidonTranscript implements the Transcript trait using the Poseidon hash
 14pub struct PoseidonTranscript<C: CurveGroup>
 15where
 16    <C as Group>::ScalarField: Absorb,
 17{
 18    sponge: PoseidonSponge<C::ScalarField>,
 19}
 20
 21impl<C: CurveGroup> Transcript<C> for PoseidonTranscript<C>
 22where
 23    <C as Group>::ScalarField: Absorb,
 24{
 25    type TranscriptConfig = PoseidonConfig<C::ScalarField>;
 26
 27    fn new(poseidon_config: &Self::TranscriptConfig) -> Self {
 28        let sponge = PoseidonSponge::<C::ScalarField>::new(poseidon_config);
 29        Self { sponge }
 30    }
 31    fn absorb(&mut self, v: &C::ScalarField) {
 32        self.sponge.absorb(&v);
 33    }
 34    fn absorb_vec(&mut self, v: &[C::ScalarField]) {
 35        self.sponge.absorb(&v);
 36    }
 37    fn absorb_point(&mut self, p: &C) {
 38        self.sponge.absorb(&prepare_point(p));
 39    }
 40    fn get_challenge(&mut self) -> C::ScalarField {
 41        let c = self.sponge.squeeze_field_elements(1);
 42        self.sponge.absorb(&c[0]);
 43        c[0]
 44    }
 45    fn get_challenges(&mut self, n: usize) -> Vec<C::ScalarField> {
 46        let c = self.sponge.squeeze_field_elements(n);
 47        self.sponge.absorb(&c);
 48        c
 49    }
 50}
 51
 52// Returns the point coordinates in Fr, so it can be absrobed by the transcript. It does not work
 53// over bytes in order to have a logic that can be reproduced in-circuit.
 54fn prepare_point<C: CurveGroup>(p: &C) -> Vec<C::ScalarField> {
 55    let binding = p.into_affine();
 56    let p_coords = &binding.xy().unwrap();
 57    let x_bi = p_coords
 58        .0
 59        .to_base_prime_field_elements()
 60        .next()
 61        .expect("a")
 62        .into_bigint();
 63    let y_bi = p_coords
 64        .1
 65        .to_base_prime_field_elements()
 66        .next()
 67        .expect("a")
 68        .into_bigint();
 69    vec![
 70        C::ScalarField::from_le_bytes_mod_order(x_bi.to_bytes_le().as_ref()),
 71        C::ScalarField::from_le_bytes_mod_order(y_bi.to_bytes_le().as_ref()),
 72    ]
 73}
 74
 75/// PoseidonTranscriptVar implements the gadget compatible with PoseidonTranscript
 76pub struct PoseidonTranscriptVar<F: PrimeField> {
 77    sponge: PoseidonSpongeVar<F>,
 78}
 79impl<F: PrimeField> PoseidonTranscriptVar<F> {
 80    pub fn new(cs: ConstraintSystemRef<F>, poseidon_config: &PoseidonConfig<F>) -> Self {
 81        let sponge = PoseidonSpongeVar::<F>::new(cs, poseidon_config);
 82        Self { sponge }
 83    }
 84    pub fn absorb(&mut self, v: FpVar<F>) -> Result<(), SynthesisError> {
 85        self.sponge.absorb(&v)
 86    }
 87    pub fn absorb_vec(&mut self, v: &[FpVar<F>]) -> Result<(), SynthesisError> {
 88        self.sponge.absorb(&v)
 89    }
 90    pub fn get_challenge(&mut self) -> Result<FpVar<F>, SynthesisError> {
 91        let c = self.sponge.squeeze_field_elements(1)?;
 92        self.sponge.absorb(&c[0])?;
 93        Ok(c[0].clone())
 94    }
 95    pub fn get_challenges(&mut self, n: usize) -> Result<Vec<FpVar<F>>, SynthesisError> {
 96        let c = self.sponge.squeeze_field_elements(n)?;
 97        self.sponge.absorb(&c)?;
 98        Ok(c)
 99    }
100}
101
102#[cfg(test)]
103pub mod tests {
104    use super::*;
105    use ark_bls12_377::{Fr, G1Projective};
106    use ark_crypto_primitives::sponge::poseidon::find_poseidon_ark_and_mds;
107    use ark_r1cs_std::{alloc::AllocVar, fields::fp::FpVar, R1CSVar};
108    use ark_relations::r1cs::ConstraintSystem;
109
110    /// WARNING the method poseidon_test_config is for tests only
111    #[cfg(test)]
112    pub fn poseidon_test_config<F: PrimeField>() -> PoseidonConfig<F> {
113        let full_rounds = 8;
114        let partial_rounds = 31;
115        let alpha = 5;
116        let rate = 2;
117
118        let (ark, mds) = find_poseidon_ark_and_mds::<F>(
119            F::MODULUS_BIT_SIZE as u64,
120            rate,
121            full_rounds,
122            partial_rounds,
123            0,
124        );
125
126        PoseidonConfig::new(
127            full_rounds as usize,
128            partial_rounds as usize,
129            alpha,
130            mds,
131            ark,
132            rate,
133            1,
134        )
135    }
136
137    #[test]
138    fn test_transcript_and_transcriptvar() {
139        // use 'native' transcript
140        let config = poseidon_test_config::<Fr>();
141        let mut tr = PoseidonTranscript::<G1Projective>::new(&config);
142        tr.absorb(&Fr::from(42_u32));
143        let c = tr.get_challenge();
144
145        // use 'gadget' transcript
146        let cs = ConstraintSystem::<Fr>::new_ref();
147        let mut tr_var = PoseidonTranscriptVar::<Fr>::new(cs.clone(), &config);
148        let v = FpVar::<Fr>::new_witness(cs.clone(), || Ok(Fr::from(42_u32))).unwrap();
149        tr_var.absorb(v).unwrap();
150        let c_var = tr_var.get_challenge().unwrap();
151
152        // assert that native & gadget transcripts return the same challenge
153        assert_eq!(c, c_var.value().unwrap());
154    }
155}