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::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 absorbed 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 zero_point = (&C::BaseField::zero(), &C::BaseField::zero());
 65    let xy = affine.xy().unwrap_or(zero_point);
 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> TranscriptVar<F> for PoseidonTranscriptVar<F> {
 88    type TranscriptVarConfig = PoseidonConfig<F>;
 89
 90    fn new(cs: ConstraintSystemRef<F>, poseidon_config: &Self::TranscriptVarConfig) -> Self {
 91        let sponge = PoseidonSpongeVar::<F>::new(cs, poseidon_config);
 92        Self { sponge }
 93    }
 94    fn absorb(&mut self, v: FpVar<F>) -> Result<(), SynthesisError> {
 95        self.sponge.absorb(&v)
 96    }
 97    fn absorb_vec(&mut self, v: &[FpVar<F>]) -> Result<(), SynthesisError> {
 98        self.sponge.absorb(&v)
 99    }
100    fn get_challenge(&mut self) -> Result<FpVar<F>, SynthesisError> {
101        let c = self.sponge.squeeze_field_elements(1)?;
102        self.sponge.absorb(&c[0])?;
103        Ok(c[0].clone())
104    }
105
106    /// returns the bit representation of the challenge, we use its output in-circuit for the
107    /// `GC.scalar_mul_le` method.
108    fn get_challenge_nbits(&mut self, nbits: usize) -> Result<Vec<Boolean<F>>, SynthesisError> {
109        self.sponge.squeeze_bits(nbits)
110    }
111    fn get_challenges(&mut self, n: usize) -> Result<Vec<FpVar<F>>, SynthesisError> {
112        let c = self.sponge.squeeze_field_elements(n)?;
113        self.sponge.absorb(&c)?;
114        Ok(c)
115    }
116}
117
118/// This Poseidon configuration generator agrees with Circom's Poseidon(4) in the case of BN254's scalar field
119pub fn poseidon_canonical_config<F: PrimeField>() -> PoseidonConfig<F> {
120    // 120 bit security target as in
121    // https://eprint.iacr.org/2019/458.pdf
122    // t = rate + 1
123
124    let full_rounds = 8;
125    let partial_rounds = 60;
126    let alpha = 5;
127    let rate = 4;
128
129    let (ark, mds) = ark_crypto_primitives::sponge::poseidon::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#[cfg(test)]
149pub mod tests {
150    use super::*;
151    use ark_bn254::{constraints::GVar, Fq, Fr, G1Projective as G1};
152    use ark_grumpkin::Projective;
153    use ark_r1cs_std::{alloc::AllocVar, groups::CurveVar, R1CSVar};
154    use ark_relations::r1cs::ConstraintSystem;
155    use std::ops::Mul;
156
157    // Test with value taken from https://github.com/iden3/circomlibjs/blob/43cc582b100fc3459cf78d903a6f538e5d7f38ee/test/poseidon.js#L32
158    #[test]
159    fn check_against_circom_poseidon() {
160        use ark_bn254::Fr;
161        use ark_crypto_primitives::sponge::{poseidon::PoseidonSponge, CryptographicSponge};
162        use std::str::FromStr;
163
164        let config = poseidon_canonical_config::<Fr>();
165        let mut poseidon_sponge: PoseidonSponge<_> = CryptographicSponge::new(&config);
166        let v: Vec<Fr> = vec!["1", "2", "3", "4"]
167            .into_iter()
168            .map(|x| Fr::from_str(x).unwrap())
169            .collect();
170        poseidon_sponge.absorb(&v);
171        poseidon_sponge.squeeze_field_elements::<Fr>(1);
172        assert!(
173            poseidon_sponge.state[0]
174                == Fr::from_str(
175                    "18821383157269793795438455681495246036402687001665670618754263018637548127333"
176                )
177                .unwrap()
178        );
179    }
180
181    #[test]
182    fn test_transcript_and_transcriptvar_get_challenge() {
183        // use 'native' transcript
184        let config = poseidon_canonical_config::<Fr>();
185        let mut tr = PoseidonTranscript::<G1>::new(&config);
186        tr.absorb(&Fr::from(42_u32));
187        let c = tr.get_challenge();
188
189        // use 'gadget' transcript
190        let cs = ConstraintSystem::<Fr>::new_ref();
191        let mut tr_var = PoseidonTranscriptVar::<Fr>::new(cs.clone(), &config);
192        let v = FpVar::<Fr>::new_witness(cs.clone(), || Ok(Fr::from(42_u32))).unwrap();
193        tr_var.absorb(v).unwrap();
194        let c_var = tr_var.get_challenge().unwrap();
195
196        // assert that native & gadget transcripts return the same challenge
197        assert_eq!(c, c_var.value().unwrap());
198    }
199
200    #[test]
201    fn test_transcript_and_transcriptvar_nbits() {
202        let nbits = crate::constants::N_BITS_RO;
203
204        // use 'native' transcript
205        let config = poseidon_canonical_config::<Fq>();
206        let mut tr = PoseidonTranscript::<Projective>::new(&config);
207        tr.absorb(&Fq::from(42_u32));
208
209        // get challenge from native transcript
210        let c_bits = tr.get_challenge_nbits(nbits);
211
212        // use 'gadget' transcript
213        let cs = ConstraintSystem::<Fq>::new_ref();
214        let mut tr_var = PoseidonTranscriptVar::<Fq>::new(cs.clone(), &config);
215        let v = FpVar::<Fq>::new_witness(cs.clone(), || Ok(Fq::from(42_u32))).unwrap();
216        tr_var.absorb(v).unwrap();
217
218        // get challenge from circuit transcript
219        let c_var = tr_var.get_challenge_nbits(nbits).unwrap();
220
221        let P = G1::generator();
222        let PVar = GVar::new_witness(cs.clone(), || Ok(P)).unwrap();
223
224        // multiply point P by the challenge in different formats, to ensure that we get the same
225        // result natively and in-circuit
226
227        // native c*P
228        let c_Fr = Fr::from_bigint(BigInteger::from_bits_le(&c_bits)).unwrap();
229        let cP_native = P.mul(c_Fr);
230
231        // native c*P using mul_bits_be (notice the .rev to convert the LE to BE)
232        let cP_native_bits = P.mul_bits_be(c_bits.into_iter().rev());
233
234        // in-circuit c*P using scalar_mul_le
235        let cPVar = PVar.scalar_mul_le(c_var.iter()).unwrap();
236
237        // check that they are equal
238        assert_eq!(
239            cP_native.into_affine(),
240            cPVar.value().unwrap().into_affine()
241        );
242        assert_eq!(
243            cP_native_bits.into_affine(),
244            cPVar.value().unwrap().into_affine()
245        );
246    }
247}