1use ark_ff::PrimeField;
  2use ark_r1cs_std::fields::{fp::FpVar, FieldVar};
  3use num_integer::Integer;
  4
  5/// Returns (b, b^2, b^4, ..., b^{2^{t-1}})
  6pub fn exponential_powers<F: PrimeField>(b: F, t: usize) -> Vec<F> {
  7    let mut r = vec![F::zero(); t];
  8    r[0] = b;
  9    for i in 1..t {
 10        r[i] = r[i - 1].square();
 11    }
 12    r
 13}
 14
 15/// The in-circuit version of `exponential_powers`
 16pub fn exponential_powers_var<F: PrimeField>(b: FpVar<F>, t: usize) -> Vec<FpVar<F>> {
 17    let mut r = vec![FpVar::zero(); t];
 18    r[0] = b;
 19    for i in 1..t {
 20        r[i] = &r[i - 1] * &r[i - 1];
 21    }
 22    r
 23}
 24
 25/// Returns (a, a^2, a^3, ..., a^{n-1})
 26pub fn all_powers<F: PrimeField>(a: F, n: usize) -> Vec<F> {
 27    let mut r = vec![F::zero(); n];
 28    for (i, r_i) in r.iter_mut().enumerate() {
 29        *r_i = a.pow([i as u64]);
 30    }
 31    r
 32}
 33
 34/// The in-circuit version of `all_powers`
 35pub fn all_powers_var<F: PrimeField>(a: FpVar<F>, n: usize) -> Vec<FpVar<F>> {
 36    if n == 0 {
 37        return vec![];
 38    }
 39    let mut r = vec![FpVar::zero(); n];
 40    r[0] = FpVar::one();
 41    for i in 1..n {
 42        r[i] = &r[i - 1] * &a;
 43    }
 44    r
 45}
 46
 47/// returns a vector containing βᵢ* = βᵢ + α ⋅ δᵢ
 48pub fn betas_star<F: PrimeField>(betas: &[F], deltas: &[F], alpha: F) -> Vec<F> {
 49    betas
 50        .iter()
 51        .zip(
 52            deltas
 53                .iter()
 54                .map(|delta_i| alpha * delta_i)
 55                .collect::<Vec<F>>(),
 56        )
 57        .map(|(beta_i, delta_i_alpha)| *beta_i + delta_i_alpha)
 58        .collect()
 59}
 60
 61/// The in-circuit version of `betas_star`
 62pub fn betas_star_var<F: PrimeField>(
 63    betas: &[FpVar<F>],
 64    deltas: &[FpVar<F>],
 65    alpha: &FpVar<F>,
 66) -> Vec<FpVar<F>> {
 67    betas
 68        .iter()
 69        .zip(deltas)
 70        .map(|(beta_i, delta_i)| beta_i + alpha * delta_i)
 71        .collect::<Vec<FpVar<F>>>()
 72}
 73
 74/// Returns the product of selected elements in `betas`.
 75/// For every index `j`, whether `betas[j]` is selected is determined by the
 76/// `j`-th bit in the binary (little endian) representation of `i`.
 77///
 78/// If `betas = (β, β^2, β^4, ..., β^{2^{t-1}})`, then the result is equal to
 79/// `β^i`.
 80pub fn pow_i<F: PrimeField>(mut i: usize, betas: &[F]) -> F {
 81    let mut j = 0;
 82    let mut r = F::one();
 83    while i > 0 {
 84        if i.is_odd() {
 85            r *= betas[j];
 86        }
 87        i >>= 1;
 88        j += 1;
 89    }
 90    r
 91}
 92
 93/// The in-circuit version of `pow_i`
 94#[allow(dead_code)] // Will remove this once we have the decider circuit for Protogalaxy
 95pub fn pow_i_var<F: PrimeField>(mut i: usize, betas: &[FpVar<F>]) -> FpVar<F> {
 96    let mut j = 0;
 97    let mut r = FieldVar::one();
 98    while i > 0 {
 99        if i.is_odd() {
100            r *= &betas[j];
101        }
102        i >>= 1;
103        j += 1;
104    }
105    r
106}
107
108#[cfg(test)]
109mod tests {
110    use std::error::Error;
111
112    use ark_bn254::Fr;
113    use ark_r1cs_std::{alloc::AllocVar, R1CSVar};
114    use ark_relations::r1cs::ConstraintSystem;
115    use ark_std::{test_rng, UniformRand};
116    use rand::Rng;
117
118    use super::*;
119
120    #[test]
121    fn test_exponential_powers() -> Result<(), Box<dyn Error>> {
122        let rng = &mut test_rng();
123
124        for t in 1..10 {
125            let cs = ConstraintSystem::<Fr>::new_ref();
126
127            let b = Fr::rand(rng);
128            let b_var = FpVar::new_witness(cs.clone(), || Ok(b))?;
129
130            let r = exponential_powers(b, t);
131            let r_var = exponential_powers_var(b_var, t);
132
133            assert_eq!(r, r_var.value()?);
134            assert!(cs.is_satisfied()?);
135        }
136
137        Ok(())
138    }
139
140    #[test]
141    fn test_all_powers() -> Result<(), Box<dyn Error>> {
142        let rng = &mut test_rng();
143
144        for n in 1..10 {
145            let cs = ConstraintSystem::<Fr>::new_ref();
146
147            let a = Fr::rand(rng);
148            let a_var = FpVar::new_witness(cs.clone(), || Ok(a))?;
149
150            let r = all_powers(a, n);
151            let r_var = all_powers_var(a_var, n);
152
153            assert_eq!(r, r_var.value()?);
154            assert!(cs.is_satisfied()?);
155        }
156
157        Ok(())
158    }
159
160    #[test]
161    fn test_betas_star() -> Result<(), Box<dyn Error>> {
162        let rng = &mut test_rng();
163
164        for t in 1..10 {
165            let cs = ConstraintSystem::<Fr>::new_ref();
166
167            let betas = (0..t).map(|_| Fr::rand(rng)).collect::<Vec<_>>();
168            let deltas = (0..t).map(|_| Fr::rand(rng)).collect::<Vec<_>>();
169            let alpha = Fr::rand(rng);
170
171            let betas_var = Vec::new_witness(cs.clone(), || Ok(betas.clone()))?;
172            let deltas_var = Vec::new_witness(cs.clone(), || Ok(deltas.clone()))?;
173            let alpha_var = FpVar::new_witness(cs.clone(), || Ok(alpha))?;
174
175            let r = betas_star(&betas, &deltas, alpha);
176            let r_var = betas_star_var(&betas_var, &deltas_var, &alpha_var);
177            assert_eq!(r, r_var.value()?);
178            assert!(cs.is_satisfied()?);
179        }
180
181        Ok(())
182    }
183
184    #[test]
185    fn test_pow_i() -> Result<(), Box<dyn Error>> {
186        let rng = &mut test_rng();
187
188        for t in 1..10 {
189            let cs = ConstraintSystem::<Fr>::new_ref();
190
191            let betas = (0..t).map(|_| Fr::rand(rng)).collect::<Vec<_>>();
192            let i = rng.gen_range(0..(1 << t));
193
194            let betas_var = Vec::new_witness(cs.clone(), || Ok(betas.clone()))?;
195
196            let r = pow_i(i, &betas);
197            let r_var = pow_i_var(i, &betas_var);
198            assert_eq!(r, r_var.value()?);
199            assert!(cs.is_satisfied()?);
200        }
201
202        Ok(())
203    }
204}