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