w3f-ring-proof 0.0.10

use ark_ec::{AdditiveGroup, AffineRepr, CurveGroup};
use ark_ff::{BigInteger, One, PrimeField, Zero};
use ark_std::rand::Rng;
use ark_std::{vec, vec::Vec};
use w3f_plonk_common::domain::Domain;
use w3f_plonk_common::gadgets::ec::AffineColumn;

use crate::piop::FixedColumns;

pub const ZK_ROWS: usize = 3;

/// Plonk Interactive Oracle Proofs (PIOP) parameters.
#[derive(Clone)]
pub struct PiopParams<G: AffineRepr<BaseField: PrimeField>> {
    /// Domain over which the piop is represented.
    pub domain: Domain<G::BaseField>,
    /// Number of bits used to represent a jubjub scalar.
    pub scalar_bitlen: usize,
    /// Length of the part of the column representing the public keys (including the padding).
    pub keyset_part_size: usize,
    /// Blinding base point.
    pub h: G,
    /// Summation base point.
    pub seed: G,
    /// The point used to pad the list of public keys.
    pub padding: G,
}

impl<G: AffineRepr<BaseField: PrimeField>> PiopParams<G> {
    /// Initialize PIOP parameters.
    ///
    /// - `domain`: polynomials evaluation domain.
    /// - `h`: Blinding base point.
    /// - `seed`: Accumulation base point
    /// - `padding`: The point used to pad the list of public keys.
    ///
    /// All points should be of an unknown discrete log.
    pub fn setup(domain: Domain<G::BaseField>, h: G, seed: G, padding: G) -> Self {
        let scalar_bitlen = G::ScalarField::MODULUS_BIT_SIZE as usize;
        // 1 accounts for the last cells of the points and bits columns that remain unconstrained
        let keyset_part_size = domain.capacity - scalar_bitlen - 1;
        Self {
            domain,
            scalar_bitlen,
            keyset_part_size,
            h,
            seed,
            padding,
        }
    }

    pub fn rand<R: Rng>(domain_size: usize, rng: &mut R) -> Self {
        let domain = Domain::with_zk_rows(domain_size, ZK_ROWS);
        let h = G::rand(rng);
        let seed = G::rand(rng);
        let padding = G::rand(rng);
        Self::setup(domain, h, seed, padding)
    }

    pub fn fixed_columns(&self, keys: &[G]) -> FixedColumns<G::BaseField, G> {
        let ring_selector = self.keyset_part_selector();
        let ring_selector = self.domain.public_column(ring_selector);
        let points = self.points_column(&keys);
        FixedColumns {
            points,
            ring_selector,
        }
    }

    pub fn points_column(&self, keys: &[G]) -> AffineColumn<G::BaseField, G> {
        assert!(keys.len() <= self.keyset_part_size);
        let padding_len = self.keyset_part_size - keys.len();
        let padding = vec![self.padding; padding_len];
        let points = [keys, &padding, &self.power_of_2_multiples_of_h()].concat();
        assert_eq!(points.len(), self.domain.capacity - 1);
        AffineColumn::public_column(points, &self.domain)
    }

    pub fn power_of_2_multiples_of_h(&self) -> Vec<G> {
        let mut h = self.h.into_group();
        let mut multiples = Vec::with_capacity(self.scalar_bitlen);
        multiples.push(h);
        for _ in 1..self.scalar_bitlen {
            h.double_in_place();
            multiples.push(h);
        }
        CurveGroup::normalize_batch(&multiples)
    }

    pub fn scalar_part(&self, e: G::ScalarField) -> Vec<bool> {
        let bits_with_trailing_zeroes = e.into_bigint().to_bits_le();
        let significant_bits = &bits_with_trailing_zeroes[..self.scalar_bitlen];
        significant_bits.to_vec()
    }

    pub fn keyset_part_selector(&self) -> Vec<G::BaseField> {
        [
            vec![G::BaseField::one(); self.keyset_part_size],
            vec![G::BaseField::zero(); self.scalar_bitlen],
        ]
        .concat()
    }

    pub fn blind_pk(&self, pk_k: G, blinding: G::ScalarField) -> G {
        let blinded_pk = pk_k + self.h * blinding;
        blinded_pk.into_affine()
    }
}

#[cfg(test)]
mod tests {
    use ark_ed_on_bls12_381_bandersnatch::{EdwardsAffine, Fr};
    use ark_std::ops::Mul;
    use ark_std::{test_rng, UniformRand};

    use w3f_plonk_common::domain::Domain;
    use w3f_plonk_common::test_helpers::cond_sum;

    use crate::piop::params::PiopParams;

    #[test]
    fn test_powers_of_h() {
        let rng = &mut test_rng();
        let h = EdwardsAffine::rand(rng);
        let seed = EdwardsAffine::rand(rng);
        let padding = EdwardsAffine::rand(rng);
        let domain = Domain::no_zk(1024);

        let params = PiopParams::<EdwardsAffine>::setup(domain, h, seed, padding);
        let t = Fr::rand(rng);
        let t_bits = params.scalar_part(t);
        let th = cond_sum(&t_bits, &params.power_of_2_multiples_of_h());
        assert_eq!(th, params.h.mul(t));
    }
}