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::{boolean::Boolean, 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_challenge_nbits(&mut self, nbits: usize) -> Vec<bool> {
 46        self.sponge.squeeze_bits(nbits)
 47    }
 48    fn get_challenges(&mut self, n: usize) -> Vec<C::ScalarField> {
 49        let c = self.sponge.squeeze_field_elements(n);
 50        self.sponge.absorb(&c);
 51        c
 52    }
 53}
 54
 55// Returns the point coordinates in Fr, so it can be absrobed by the transcript. It does not work
 56// over bytes in order to have a logic that can be reproduced in-circuit.
 57fn prepare_point<C: CurveGroup>(p: &C) -> Vec<C::ScalarField> {
 58    let binding = p.into_affine();
 59    let p_coords = &binding.xy().unwrap();
 60    let x_bi = p_coords
 61        .0
 62        .to_base_prime_field_elements()
 63        .next()
 64        .expect("a")
 65        .into_bigint();
 66    let y_bi = p_coords
 67        .1
 68        .to_base_prime_field_elements()
 69        .next()
 70        .expect("a")
 71        .into_bigint();
 72    vec![
 73        C::ScalarField::from_le_bytes_mod_order(x_bi.to_bytes_le().as_ref()),
 74        C::ScalarField::from_le_bytes_mod_order(y_bi.to_bytes_le().as_ref()),
 75    ]
 76}
 77
 78/// PoseidonTranscriptVar implements the gadget compatible with PoseidonTranscript
 79pub struct PoseidonTranscriptVar<F: PrimeField> {
 80    sponge: PoseidonSpongeVar<F>,
 81}
 82impl<F: PrimeField> PoseidonTranscriptVar<F> {
 83    pub fn new(cs: ConstraintSystemRef<F>, poseidon_config: &PoseidonConfig<F>) -> Self {
 84        let sponge = PoseidonSpongeVar::<F>::new(cs, poseidon_config);
 85        Self { sponge }
 86    }
 87    pub fn absorb(&mut self, v: FpVar<F>) -> Result<(), SynthesisError> {
 88        self.sponge.absorb(&v)
 89    }
 90    pub fn absorb_vec(&mut self, v: &[FpVar<F>]) -> Result<(), SynthesisError> {
 91        self.sponge.absorb(&v)
 92    }
 93    pub fn get_challenge(&mut self) -> Result<FpVar<F>, SynthesisError> {
 94        let c = self.sponge.squeeze_field_elements(1)?;
 95        self.sponge.absorb(&c[0])?;
 96        Ok(c[0].clone())
 97    }
 98
 99    /// returns the bit representation of the challenge, we use its output in-circuit for the
100    /// `GC.scalar_mul_le` method.
101    pub fn get_challenge_nbits(&mut self, nbits: usize) -> Result<Vec<Boolean<F>>, SynthesisError> {
102        self.sponge.squeeze_bits(nbits)
103    }
104    pub fn get_challenges(&mut self, n: usize) -> Result<Vec<FpVar<F>>, SynthesisError> {
105        let c = self.sponge.squeeze_field_elements(n)?;
106        self.sponge.absorb(&c)?;
107        Ok(c)
108    }
109}
110
111#[cfg(test)]
112pub mod tests {
113    use super::*;
114    use ark_crypto_primitives::sponge::poseidon::find_poseidon_ark_and_mds;
115    use ark_pallas::{constraints::GVar, Fq, Fr, Projective};
116    use ark_r1cs_std::{alloc::AllocVar, fields::fp::FpVar, groups::CurveVar, R1CSVar};
117    use ark_relations::r1cs::ConstraintSystem;
118    use ark_vesta::Projective as E2Projective;
119    use std::ops::Mul;
120
121    /// WARNING the method poseidon_test_config is for tests only
122    #[cfg(test)]
123    pub fn poseidon_test_config<F: PrimeField>() -> PoseidonConfig<F> {
124        let full_rounds = 8;
125        let partial_rounds = 31;
126        let alpha = 5;
127        let rate = 2;
128
129        let (ark, mds) = find_poseidon_ark_and_mds::<F>(
130            F::MODULUS_BIT_SIZE as u64,
131            rate,
132            full_rounds,
133            partial_rounds,
134            0,
135        );
136
137        PoseidonConfig::new(
138            full_rounds as usize,
139            partial_rounds as usize,
140            alpha,
141            mds,
142            ark,
143            rate,
144            1,
145        )
146    }
147
148    #[test]
149    fn test_transcript_and_transcriptvar_get_challenge() {
150        // use 'native' transcript
151        let config = poseidon_test_config::<Fr>();
152        let mut tr = PoseidonTranscript::<Projective>::new(&config);
153        tr.absorb(&Fr::from(42_u32));
154        let c = tr.get_challenge();
155
156        // use 'gadget' transcript
157        let cs = ConstraintSystem::<Fr>::new_ref();
158        let mut tr_var = PoseidonTranscriptVar::<Fr>::new(cs.clone(), &config);
159        let v = FpVar::<Fr>::new_witness(cs.clone(), || Ok(Fr::from(42_u32))).unwrap();
160        tr_var.absorb(v).unwrap();
161        let c_var = tr_var.get_challenge().unwrap();
162
163        // assert that native & gadget transcripts return the same challenge
164        assert_eq!(c, c_var.value().unwrap());
165    }
166
167    #[test]
168    fn test_transcript_and_transcriptvar_nbits() {
169        let nbits = crate::constants::N_BITS_CHALLENGE;
170
171        // use 'native' transcript
172        let config = poseidon_test_config::<Fq>();
173        let mut tr = PoseidonTranscript::<E2Projective>::new(&config);
174        tr.absorb(&Fq::from(42_u32));
175
176        // get challenge from native transcript
177        let c_bits = tr.get_challenge_nbits(nbits);
178
179        // use 'gadget' transcript
180        let cs = ConstraintSystem::<Fq>::new_ref();
181        let mut tr_var = PoseidonTranscriptVar::<Fq>::new(cs.clone(), &config);
182        let v = FpVar::<Fq>::new_witness(cs.clone(), || Ok(Fq::from(42_u32))).unwrap();
183        tr_var.absorb(v).unwrap();
184
185        // get challenge from circuit transcript
186        let c_var = tr_var.get_challenge_nbits(nbits).unwrap();
187
188        let P = Projective::generator();
189        let PVar = GVar::new_witness(cs.clone(), || Ok(P)).unwrap();
190
191        // multiply point P by the challenge in different formats, to ensure that we get the same
192        // result natively and in-circuit
193
194        // native c*P
195        let c_Fr = Fr::from_bigint(BigInteger::from_bits_le(&c_bits)).unwrap();
196        let cP_native = P.mul(c_Fr);
197
198        // native c*P using mul_bits_be (notice the .rev to convert the LE to BE)
199        let cP_native_bits = P.mul_bits_be(c_bits.into_iter().rev());
200
201        // in-circuit c*P using scalar_mul_le
202        let cPVar = PVar.scalar_mul_le(c_var.iter()).unwrap();
203
204        // check that they are equal
205        assert_eq!(
206            cP_native.into_affine(),
207            cPVar.value().unwrap().into_affine()
208        );
209        assert_eq!(
210            cP_native_bits.into_affine(),
211            cPVar.value().unwrap().into_affine()
212        );
213    }
214}