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