1use ark_ec::{AffineRepr, CurveGroup};
  2use ark_ff::PrimeField;
  3use std::marker::PhantomData;
  4
  5use ark_r1cs_std::fields::fp::FpVar;
  6
  7use ark_crypto_primitives::sponge::poseidon::{
  8    constraints::PoseidonSpongeVar, PoseidonConfig, PoseidonSponge,
  9};
 10use ark_crypto_primitives::sponge::{
 11    constraints::CryptographicSpongeVar, Absorb, CryptographicSponge,
 12};
 13use ark_relations::r1cs::{ConstraintSystemRef, SynthesisError};
 14
 15pub struct Transcript<F: PrimeField + Absorb, C: CurveGroup>
 16where
 17    <C as CurveGroup>::BaseField: Absorb,
 18{
 19    // where F is the Constraint Field (eq. C::ScalarField)
 20    sponge: PoseidonSponge<F>,
 21    _c: PhantomData<C>,
 22}
 23impl<F: PrimeField + Absorb, C: CurveGroup> Transcript<F, C>
 24where
 25    <C as CurveGroup>::BaseField: Absorb,
 26{
 27    pub fn new(poseidon_config: &PoseidonConfig<F>) -> Self {
 28        let sponge = PoseidonSponge::<F>::new(poseidon_config);
 29        Transcript {
 30            sponge,
 31            _c: PhantomData,
 32        }
 33    }
 34    pub fn add(&mut self, v: &F) {
 35        self.sponge.absorb(&v);
 36    }
 37    pub fn add_vec(&mut self, v: &[F]) {
 38        self.sponge.absorb(&v);
 39    }
 40    pub fn add_point(&mut self, v: &C) {
 41        let v_affine = v.into_affine();
 42        let xy = v_affine.xy().unwrap(); // WIP
 43        self.sponge.absorb(&vec![xy.0, xy.1]);
 44    }
 45    pub fn get_challenge(&mut self) -> F {
 46        let c = self.sponge.squeeze_field_elements(1);
 47        self.add(&c[0]);
 48        c[0]
 49    }
 50    pub fn get_challenge_vec(&mut self, n: usize) -> Vec<F> {
 51        let c = self.sponge.squeeze_field_elements(n);
 52        self.sponge.absorb(&c);
 53        c
 54    }
 55}
 56
 57pub struct TranscriptVar<F: PrimeField> {
 58    // where F is the Constraint Field
 59    sponge: PoseidonSpongeVar<F>,
 60}
 61impl<F: PrimeField> TranscriptVar<F> {
 62    pub fn new(cs: ConstraintSystemRef<F>, poseidon_config: &PoseidonConfig<F>) -> Self {
 63        let sponge = PoseidonSpongeVar::<F>::new(cs, poseidon_config);
 64        Self { sponge }
 65    }
 66    pub fn add(&mut self, v: FpVar<F>) -> Result<(), SynthesisError> {
 67        self.sponge.absorb(&v)
 68    }
 69    pub fn get_challenge(&mut self) -> Result<FpVar<F>, SynthesisError> {
 70        let c = self.sponge.squeeze_field_elements(1)?;
 71        self.sponge.absorb(&c[0])?;
 72        Ok(c[0].clone())
 73    }
 74    pub fn get_challenge_vec(&mut self, n: usize) -> Result<Vec<FpVar<F>>, SynthesisError> {
 75        let c = self.sponge.squeeze_field_elements(n)?;
 76        self.sponge.absorb(&c)?;
 77        Ok(c)
 78    }
 79}
 80
 81use ark_crypto_primitives::sponge::poseidon::find_poseidon_ark_and_mds;
 82
 83// WARNING this is for test only
 84pub fn poseidon_test_config<F: PrimeField>() -> PoseidonConfig<F> {
 85    let full_rounds = 8;
 86    let partial_rounds = 31;
 87    let alpha = 5;
 88    let rate = 2;
 89
 90    let (ark, mds) = find_poseidon_ark_and_mds::<F>(
 91        F::MODULUS_BIT_SIZE as u64,
 92        rate,
 93        full_rounds,
 94        partial_rounds,
 95        0,
 96    );
 97
 98    PoseidonConfig::new(
 99        full_rounds as usize,
100        partial_rounds as usize,
101        alpha,
102        mds,
103        ark,
104        rate,
105        1,
106    )
107}
108
109#[cfg(test)]
110mod tests {
111    use super::*;
112    use ark_mnt4_298::{Fr, G1Projective}; // scalar field
113    use ark_r1cs_std::{alloc::AllocVar, fields::fp::FpVar, R1CSVar};
114    use ark_relations::r1cs::ConstraintSystem;
115
116    #[test]
117    fn test_transcript_and_transcriptvar() {
118        let config = poseidon_test_config::<Fr>();
119        let mut tr = Transcript::<Fr, G1Projective>::new(&config);
120        tr.add(&Fr::from(42_u32));
121        let c = tr.get_challenge();
122
123        let cs = ConstraintSystem::<Fr>::new_ref();
124        let mut tr_var = TranscriptVar::<Fr>::new(cs.clone(), &config);
125        let v = FpVar::<Fr>::new_witness(cs.clone(), || Ok(Fr::from(42_u32))).unwrap();
126        tr_var.add(v).unwrap();
127        let c_var = tr_var.get_challenge().unwrap();
128
129        assert_eq!(c, c_var.value().unwrap());
130    }
131}