1use ark_crypto_primitives::sponge::{
  2    constraints::CryptographicSpongeVar,
  3    poseidon::{constraints::PoseidonSpongeVar, PoseidonConfig, PoseidonSponge},
  4    Absorb, CryptographicSponge,
  5};
  6use ark_ff::PrimeField;
  7use ark_r1cs_std::fields::fp::FpVar;
  8use ark_relations::r1cs::{ConstraintSystemRef, SynthesisError};
  9
 10use crate::transcript::Transcript;
 11
 12/// PoseidonTranscript implements the Transcript trait using the Poseidon hash
 13pub struct PoseidonTranscript<F: PrimeField + Absorb> {
 14    sponge: PoseidonSponge<F>,
 15}
 16
 17impl<F: PrimeField + Absorb> Transcript<F> for PoseidonTranscript<F> {
 18    type TranscriptConfig = PoseidonConfig<F>;
 19
 20    fn new(poseidon_config: &Self::TranscriptConfig) -> Self {
 21        let sponge = PoseidonSponge::<F>::new(poseidon_config);
 22        Self { sponge }
 23    }
 24    fn absorb(&mut self, v: &F) {
 25        self.sponge.absorb(&v);
 26    }
 27    fn absorb_vec(&mut self, v: &[F]) {
 28        self.sponge.absorb(&v);
 29    }
 30    fn get_challenge(&mut self) -> F {
 31        let c = self.sponge.squeeze_field_elements(1);
 32        self.sponge.absorb(&c[0]);
 33        c[0]
 34    }
 35    fn get_challenges(&mut self, n: usize) -> Vec<F> {
 36        let c = self.sponge.squeeze_field_elements(n);
 37        self.sponge.absorb(&c);
 38        c
 39    }
 40}
 41
 42/// PoseidonTranscriptVar implements the gadget compatible with PoseidonTranscript
 43pub struct PoseidonTranscriptVar<F: PrimeField> {
 44    sponge: PoseidonSpongeVar<F>,
 45}
 46impl<F: PrimeField> PoseidonTranscriptVar<F> {
 47    pub fn new(cs: ConstraintSystemRef<F>, poseidon_config: &PoseidonConfig<F>) -> Self {
 48        let sponge = PoseidonSpongeVar::<F>::new(cs, poseidon_config);
 49        Self { sponge }
 50    }
 51    pub fn absorb(&mut self, v: FpVar<F>) -> Result<(), SynthesisError> {
 52        self.sponge.absorb(&v)
 53    }
 54    pub fn absorb_vec(&mut self, v: &[FpVar<F>]) -> Result<(), SynthesisError> {
 55        self.sponge.absorb(&v)
 56    }
 57    pub fn get_challenge(&mut self) -> Result<FpVar<F>, SynthesisError> {
 58        let c = self.sponge.squeeze_field_elements(1)?;
 59        self.sponge.absorb(&c[0])?;
 60        Ok(c[0].clone())
 61    }
 62    pub fn get_challenges(&mut self, n: usize) -> Result<Vec<FpVar<F>>, SynthesisError> {
 63        let c = self.sponge.squeeze_field_elements(n)?;
 64        self.sponge.absorb(&c)?;
 65        Ok(c)
 66    }
 67}
 68
 69#[cfg(test)]
 70mod tests {
 71    use super::*;
 72    use ark_bls12_381::Fr;
 73    use ark_crypto_primitives::sponge::poseidon::find_poseidon_ark_and_mds;
 74    use ark_r1cs_std::{alloc::AllocVar, fields::fp::FpVar, R1CSVar};
 75    use ark_relations::r1cs::ConstraintSystem;
 76
 77    /// WARNING the method poseidon_test_config is for tests only
 78    #[cfg(test)]
 79    pub fn poseidon_test_config<F: PrimeField>() -> PoseidonConfig<F> {
 80        let full_rounds = 8;
 81        let partial_rounds = 31;
 82        let alpha = 5;
 83        let rate = 2;
 84
 85        let (ark, mds) = find_poseidon_ark_and_mds::<F>(
 86            F::MODULUS_BIT_SIZE as u64,
 87            rate,
 88            full_rounds,
 89            partial_rounds,
 90            0,
 91        );
 92
 93        PoseidonConfig::new(
 94            full_rounds as usize,
 95            partial_rounds as usize,
 96            alpha,
 97            mds,
 98            ark,
 99            rate,
100            1,
101        )
102    }
103
104    #[test]
105    fn test_transcript_and_transcriptvar() {
106        // use 'native' transcript
107        let config = poseidon_test_config::<Fr>();
108        let mut tr = PoseidonTranscript::<Fr>::new(&config);
109        tr.absorb(&Fr::from(42_u32));
110        let c = tr.get_challenge();
111
112        // use 'gadget' transcript
113        let cs = ConstraintSystem::<Fr>::new_ref();
114        let mut tr_var = PoseidonTranscriptVar::<Fr>::new(cs.clone(), &config);
115        let v = FpVar::<Fr>::new_witness(cs.clone(), || Ok(Fr::from(42_u32))).unwrap();
116        tr_var.absorb(v).unwrap();
117        let c_var = tr_var.get_challenge().unwrap();
118
119        // assert that native & gadget transcripts return the same challenge
120        assert_eq!(c, c_var.value().unwrap());
121    }
122}