Refactor pow_i

folding-schemes/src/folding/protogalaxy/folding.rs

@@ -10,7 +10,7 @@ use ark_std::{cfg_into_iter, log2, One, Zero};
 use rayon::prelude::*;
 use std::marker::PhantomData;
 
-use super::utils::{all_powers, betas_star, exponential_powers};
+use super::utils::{all_powers, betas_star, exponential_powers, pow_i};
 use super::ProtoGalaxyError;
 use super::{CommittedInstance, Witness};
 
@@ -19,7 +19,6 @@ use crate::arith::r1cs::RelaxedR1CS;
 use crate::arith::{r1cs::R1CS, Arith};
 use crate::transcript::Transcript;
 use crate::utils::vec::*;
-use crate::utils::virtual_polynomial::bit_decompose;
 use crate::Error;
 
 #[derive(Clone, Debug)]
@@ -323,22 +322,6 @@ where
     }
 }
 
-// naive impl of pow_i for betas, assuming that betas=(b, b^2, b^4, ..., b^{2^{t-1}})
-pub fn pow_i<F: PrimeField>(i: usize, betas: &[F]) -> F {
-    // WIP check if makes more sense to do it with ifs instead of arithmetic
-
-    let n = 2_u64.pow(betas.len() as u32);
-    let b = bit_decompose(i as u64, n as usize);
-
-    let mut r: F = F::one();
-    for (j, beta_j) in betas.iter().enumerate() {
-        if b[j] {
-            r *= beta_j;
-        }
-    }
-    r
-}
-
 /// calculates F[x] using the optimized binary-tree technique
 /// described in Claim 4.4
 /// of [ProtoGalaxy](https://eprint.iacr.org/2023/1106.pdf)

folding-schemes/src/folding/protogalaxy/traits.rs

@@ -6,7 +6,7 @@ use ark_relations::r1cs::SynthesisError;
 use ark_std::{cfg_iter, log2, rand::RngCore, One, Zero};
 use rayon::prelude::*;
 
-use super::{folding::pow_i, CommittedInstance, CommittedInstanceVar, Witness};
+use super::{utils::pow_i, CommittedInstance, CommittedInstanceVar, Witness};
 use crate::{
     arith::r1cs::{RelaxedR1CS, R1CS},
     transcript::AbsorbNonNative,

folding-schemes/src/folding/protogalaxy/utils.rs

@@ -1,5 +1,6 @@
 use ark_ff::PrimeField;
 use ark_r1cs_std::fields::{fp::FpVar, FieldVar};
+use num_integer::Integer;
 
 /// Returns (b, b^2, b^4, ..., b^{2^{t-1}})
 pub fn exponential_powers<F: PrimeField>(b: F, t: usize) -> Vec<F> {
@@ -70,6 +71,40 @@ pub fn betas_star_var<F: PrimeField>(
         .collect::<Vec<FpVar<F>>>()
 }
 
+/// Returns the product of selected elements in `betas`.
+/// For every index `j`, whether `betas[j]` is selected is determined by the
+/// `j`-th bit in the binary (little endian) representation of `i`.
+///
+/// If `betas = (β, β^2, β^4, ..., β^{2^{t-1}})`, then the result is equal to
+/// `β^i`.
+pub fn pow_i<F: PrimeField>(mut i: usize, betas: &[F]) -> F {
+    let mut j = 0;
+    let mut r = F::one();
+    while i > 0 {
+        if i.is_odd() {
+            r *= betas[j];
+        }
+        i >>= 1;
+        j += 1;
+    }
+    r
+}
+
+/// The in-circuit version of `pow_i`
+#[allow(dead_code, reason = "This will be used in the decider circuit")]
+pub fn pow_i_var<F: PrimeField>(mut i: usize, betas: &[FpVar<F>]) -> FpVar<F> {
+    let mut j = 0;
+    let mut r = FpVar::one();
+    while i > 0 {
+        if i.is_odd() {
+            r *= &betas[j];
+        }
+        i >>= 1;
+        j += 1;
+    }
+    r
+}
+
 #[cfg(test)]
 mod tests {
     use std::error::Error;
@@ -78,6 +113,7 @@ mod tests {
     use ark_r1cs_std::{alloc::AllocVar, R1CSVar};
     use ark_relations::r1cs::ConstraintSystem;
     use ark_std::{test_rng, UniformRand};
+    use rand::Rng;
 
     use super::*;
 
@@ -144,4 +180,25 @@ mod tests {
 
         Ok(())
     }
+
+    #[test]
+    fn test_pow_i() -> Result<(), Box<dyn Error>> {
+        let rng = &mut test_rng();
+
+        for t in 1..10 {
+            let cs = ConstraintSystem::<Fr>::new_ref();
+
+            let betas = (0..t).map(|_| Fr::rand(rng)).collect::<Vec<_>>();
+            let i = rng.gen_range(0..(1 << t));
+
+            let betas_var = Vec::new_witness(cs.clone(), || Ok(betas.clone()))?;
+
+            let r = pow_i(i, &betas);
+            let r_var = pow_i_var(i, &betas_var);
+            assert_eq!(r, r_var.value()?);
+            assert!(cs.is_satisfied()?);
+        }
+
+        Ok(())
+    }
 }