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