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