1use ark_ff::PrimeField;
  2use ark_r1cs_std::fields::{fp::FpVar, FieldVar};
  3
  4/// Returns (b, b^2, b^4, ..., b^{2^{t-1}})
  5pub fn exponential_powers<F: PrimeField>(b: F, t: usize) -> Vec<F> {
  6    let mut r = vec![F::zero(); t];
  7    r[0] = b;
  8    for i in 1..t {
  9        r[i] = r[i - 1].square();
 10    }
 11    r
 12}
 13
 14/// The in-circuit version of `exponential_powers`
 15pub fn exponential_powers_var<F: PrimeField>(b: FpVar<F>, t: usize) -> Vec<FpVar<F>> {
 16    let mut r = vec![FpVar::zero(); t];
 17    r[0] = b;
 18    for i in 1..t {
 19        r[i] = &r[i - 1] * &r[i - 1];
 20    }
 21    r
 22}
 23
 24/// Returns (a, a^2, a^3, ..., a^{n-1})
 25pub fn all_powers<F: PrimeField>(a: F, n: usize) -> Vec<F> {
 26    let mut r = vec![F::zero(); n];
 27    for (i, r_i) in r.iter_mut().enumerate() {
 28        *r_i = a.pow([i as u64]);
 29    }
 30    r
 31}
 32
 33/// The in-circuit version of `all_powers`
 34pub fn all_powers_var<F: PrimeField>(a: FpVar<F>, n: usize) -> Vec<FpVar<F>> {
 35    if n == 0 {
 36        return vec![];
 37    }
 38    let mut r = vec![FpVar::zero(); n];
 39    r[0] = FpVar::one();
 40    for i in 1..n {
 41        r[i] = &r[i - 1] * &a;
 42    }
 43    r
 44}
 45
 46/// returns a vector containing βᵢ* = βᵢ + α ⋅ δᵢ
 47pub fn betas_star<F: PrimeField>(betas: &[F], deltas: &[F], alpha: F) -> Vec<F> {
 48    betas
 49        .iter()
 50        .zip(
 51            deltas
 52                .iter()
 53                .map(|delta_i| alpha * delta_i)
 54                .collect::<Vec<F>>(),
 55        )
 56        .map(|(beta_i, delta_i_alpha)| *beta_i + delta_i_alpha)
 57        .collect()
 58}
 59
 60/// The in-circuit version of `betas_star`
 61pub fn betas_star_var<F: PrimeField>(
 62    betas: &[FpVar<F>],
 63    deltas: &[FpVar<F>],
 64    alpha: &FpVar<F>,
 65) -> Vec<FpVar<F>> {
 66    betas
 67        .iter()
 68        .zip(deltas)
 69        .map(|(beta_i, delta_i)| beta_i + alpha * delta_i)
 70        .collect::<Vec<FpVar<F>>>()
 71}
 72
 73#[cfg(test)]
 74mod tests {
 75    use std::error::Error;
 76
 77    use ark_bn254::Fr;
 78    use ark_r1cs_std::{alloc::AllocVar, R1CSVar};
 79    use ark_relations::r1cs::ConstraintSystem;
 80    use ark_std::{test_rng, UniformRand};
 81
 82    use super::*;
 83
 84    #[test]
 85    fn test_exponential_powers() -> Result<(), Box<dyn Error>> {
 86        let rng = &mut test_rng();
 87
 88        for t in 1..10 {
 89            let cs = ConstraintSystem::<Fr>::new_ref();
 90
 91            let b = Fr::rand(rng);
 92            let b_var = FpVar::new_witness(cs.clone(), || Ok(b))?;
 93
 94            let r = exponential_powers(b, t);
 95            let r_var = exponential_powers_var(b_var, t);
 96
 97            assert_eq!(r, r_var.value()?);
 98            assert!(cs.is_satisfied()?);
 99        }
100
101        Ok(())
102    }
103
104    #[test]
105    fn test_all_powers() -> Result<(), Box<dyn Error>> {
106        let rng = &mut test_rng();
107
108        for n in 1..10 {
109            let cs = ConstraintSystem::<Fr>::new_ref();
110
111            let a = Fr::rand(rng);
112            let a_var = FpVar::new_witness(cs.clone(), || Ok(a))?;
113
114            let r = all_powers(a, n);
115            let r_var = all_powers_var(a_var, n);
116
117            assert_eq!(r, r_var.value()?);
118            assert!(cs.is_satisfied()?);
119        }
120
121        Ok(())
122    }
123
124    #[test]
125    fn test_betas_star() -> Result<(), Box<dyn Error>> {
126        let rng = &mut test_rng();
127
128        for t in 1..10 {
129            let cs = ConstraintSystem::<Fr>::new_ref();
130
131            let betas = (0..t).map(|_| Fr::rand(rng)).collect::<Vec<_>>();
132            let deltas = (0..t).map(|_| Fr::rand(rng)).collect::<Vec<_>>();
133            let alpha = Fr::rand(rng);
134
135            let betas_var = Vec::new_witness(cs.clone(), || Ok(betas.clone()))?;
136            let deltas_var = Vec::new_witness(cs.clone(), || Ok(deltas.clone()))?;
137            let alpha_var = FpVar::new_witness(cs.clone(), || Ok(alpha))?;
138
139            let r = betas_star(&betas, &deltas, alpha);
140            let r_var = betas_star_var(&betas_var, &deltas_var, &alpha_var);
141            assert_eq!(r, r_var.value()?);
142            assert!(cs.is_satisfied()?);
143        }
144
145        Ok(())
146    }
147}