Hybrid method for random access/conditional selection from 2^k values

benches/bench.rs

@@ -2,13 +2,15 @@ use ark_ff::PrimeField;
 use ark_r1cs_std::{
     alloc::AllocVar,
     eq::EqGadget,
-    fields::{nonnative::NonNativeFieldVar, FieldVar},
+    fields::{fp::FpVar, nonnative::NonNativeFieldVar, FieldVar},
+    prelude::Boolean,
+    select::CondSelectGadget,
 };
 use ark_relations::{
     ns,
     r1cs::{ConstraintSystem, ConstraintSystemRef, OptimizationGoal},
 };
-use ark_std::rand::RngCore;
+use ark_std::rand::{Rng, RngCore};
 
 const NUM_REPETITIONS: usize = 1;
 
@@ -162,6 +164,64 @@ fn inverse<TargetField: PrimeField, BaseField: PrimeField, R: RngCore>(
     );
 }
 
+fn cond_select<BaseField: PrimeField, R: RngCore>(
+    cs: ConstraintSystemRef<BaseField>,
+    rng: &mut R,
+) -> (usize, usize) {
+    // value array
+    let values: Vec<BaseField> = (0..128).map(|_| BaseField::rand(rng)).collect();
+    let values_const: Vec<FpVar<BaseField>> = values.iter().map(|x| FpVar::Constant(*x)).collect();
+
+    // index array
+    let position: Vec<bool> = (0..7).map(|_| rng.gen()).collect();
+    let position_var: Vec<Boolean<BaseField>> = position
+        .iter()
+        .map(|b| {
+            Boolean::new_witness(ark_relations::ns!(cs, "index_arr_element"), || Ok(*b)).unwrap()
+        })
+        .collect();
+
+    let constraints_before = cs.num_constraints();
+    let nonzeros_before = get_density(&cs);
+
+    let _ = FpVar::conditionally_select_power_of_two_vector(&position_var, &values_const);
+
+    let constraints_after = cs.num_constraints();
+    let nonzeros_after = get_density(&cs);
+
+    return (
+        constraints_after - constraints_before,
+        nonzeros_after - nonzeros_before,
+    );
+}
+
+macro_rules! cond_select_bench_individual {
+    ($bench_base_field:ty) => {
+        let rng = &mut ark_std::test_rng();
+        let mut num_constraints = 0;
+        let mut num_nonzeros = 0;
+        for _ in 0..NUM_REPETITIONS {
+            let cs_sys = ConstraintSystem::<$bench_base_field>::new();
+            let cs = ConstraintSystemRef::new(cs_sys);
+            cs.set_optimization_goal(OptimizationGoal::Constraints);
+
+            let (cur_constraints, cur_nonzeros) =
+                cond_select::<$bench_base_field, _>(cs.clone(), rng);
+
+            num_constraints += cur_constraints;
+            num_nonzeros += cur_nonzeros;
+
+            assert!(cs.is_satisfied().unwrap());
+        }
+        let average_constraints = num_constraints / NUM_REPETITIONS;
+        let average_nonzeros = num_nonzeros / NUM_REPETITIONS;
+        println!(
+            "cond_select takes: {} constraints, {} non-zeros",
+            average_constraints, average_nonzeros,
+        );
+    };
+}
+
 macro_rules! nonnative_bench_individual {
     ($bench_method:ident, $bench_name:ident, $bench_target_field:ty, $bench_base_field:ty) => {
         let rng = &mut ark_std::test_rng();
@@ -235,4 +295,5 @@ fn main() {
     nonnative_bench!(BLS12MNT4Small, ark_bls12_381::Fr, ark_mnt4_298::Fr);
     nonnative_bench!(BLS12, ark_bls12_381::Fq, ark_bls12_381::Fr);
     nonnative_bench!(MNT6BigMNT4Small, ark_mnt6_753::Fr, ark_mnt4_298::Fr);
+    cond_select_bench_individual!(ark_mnt6_753::Fr);
 }

rustfmt.toml

@@ -6,4 +6,4 @@ match_block_trailing_comma = true
 use_field_init_shorthand = true
 edition = "2018"
 condense_wildcard_suffixes = true
-merge_imports = true
+imports_granularity="Crate"

src/fields/fp/mod.rs

@@ -425,8 +425,9 @@ impl<F: PrimeField> AllocatedFp<F> {
         // The high level logic is as follows:
         // We want to check that self - other != 0. We do this by checking that
         // (self - other).inverse() exists. In more detail, we check the following:
-        // If `should_enforce == true`, then we set `multiplier = (self - other).inverse()`,
-        // and check that (self - other) * multiplier == 1. (i.e., that the inverse exists)
+        // If `should_enforce == true`, then we set `multiplier = (self -
+        // other).inverse()`, and check that (self - other) * multiplier == 1.
+        // (i.e., that the inverse exists)
         //
         // If `should_enforce == false`, then we set `multiplier == 0`, and check that
         // (self - other) * 0 == 0, which is always satisfied.
@@ -978,6 +979,23 @@ impl<F: PrimeField> CondSelectGadget<F> for FpVar<F> {
             },
         }
     }
+
+    fn add_lc(
+        val: &Self,
+        lc: LinearCombination<F>,
+    ) -> Result<LinearCombination<F>, SynthesisError> {
+        let v = val.value()?;
+        Ok(lc * v)
+    }
+
+    fn allocate_with_value(
+        var: Variable,
+        val: &Self,
+        cs: &ConstraintSystemRef<F>,
+    ) -> Result<Self, SynthesisError> {
+        let v = val.value()?;
+        Ok(AllocatedFp::new(Some(v), var, cs.clone()).into())
+    }
 }
 
 /// Uses two bits to perform a lookup into a table
@@ -1067,12 +1085,14 @@ impl<'a, F: PrimeField> Sum<&'a FpVar<F>> for FpVar<F> {
 mod test {
     use crate::{
         alloc::{AllocVar, AllocationMode},
+        boolean::Boolean,
         eq::EqGadget,
         fields::fp::FpVar,
+        select::CondSelectGadget,
         R1CSVar,
     };
     use ark_relations::r1cs::ConstraintSystem;
-    use ark_std::{UniformRand, Zero};
+    use ark_std::{rand::Rng, UniformRand, Zero};
     use ark_test_curves::bls12_381::Fr;
 
     #[test]
@@ -1105,4 +1125,42 @@ mod test {
         assert!(cs.is_satisfied().unwrap());
         assert_eq!(sum.value().unwrap(), sum_expected);
     }
+
+    #[test]
+    fn test_fpvar_random_access() {
+        let mut rng = ark_std::test_rng();
+
+        for _ in 0..100 {
+            let cs = ConstraintSystem::<Fr>::new_ref();
+
+            // value array
+            let values: Vec<Fr> = (0..128).map(|_| rng.gen()).collect();
+            let values_const: Vec<FpVar<Fr>> = values.iter().map(|x| FpVar::Constant(*x)).collect();
+
+            // index array
+            let position: Vec<bool> = (0..7).map(|_| rng.gen()).collect();
+            let position_var: Vec<Boolean<Fr>> = position
+                .iter()
+                .map(|b| {
+                    Boolean::new_witness(ark_relations::ns!(cs, "index_arr_element"), || Ok(*b))
+                        .unwrap()
+                })
+                .collect();
+
+            // index
+            let mut index = 0;
+            for x in position {
+                index *= 2;
+                index += if x { 1 } else { 0 };
+            }
+
+            assert_eq!(
+                FpVar::conditionally_select_power_of_two_vector(&position_var, &values_const)
+                    .unwrap()
+                    .value()
+                    .unwrap(),
+                values[index]
+            )
+        }
+    }
 }

src/fields/mod.rs

@@ -21,7 +21,8 @@ pub mod quadratic_extension;
 pub mod fp;
 
 /// This module contains a generic implementation of "nonnative" prime field
-/// variables. It emulates `Fp` arithmetic using `Fq` operations, where `p != q`.
+/// variables. It emulates `Fp` arithmetic using `Fq` operations, where `p !=
+/// q`.
 pub mod nonnative;
 
 /// This module contains a generic implementation of the degree-12 tower

src/fields/nonnative/allocated_field_var.rs

@@ -634,10 +634,10 @@ impl<TargetField: PrimeField, BaseField: PrimeField>
     }
 
     /// Allocates a new non-native field witness with value given by the
-    /// function `f`.  Enforces that the field element has value in `[0, modulus)`,
-    /// and returns the bits of its binary representation.
-    /// The bits are in little-endian (i.e., the bit at index 0 is the LSB) and the
-    /// bit-vector is empty in non-witness allocation modes.
+    /// function `f`.  Enforces that the field element has value in `[0,
+    /// modulus)`, and returns the bits of its binary representation.
+    /// The bits are in little-endian (i.e., the bit at index 0 is the LSB) and
+    /// the bit-vector is empty in non-witness allocation modes.
     pub fn new_witness_with_le_bits<T: Borrow<TargetField>>(
         cs: impl Into<Namespace<BaseField>>,
         f: impl FnOnce() -> Result<T, SynthesisError>,

src/pairing/mod.rs

@@ -1,6 +1,5 @@
 use crate::prelude::*;
-use ark_ec::pairing::Pairing;
-use ark_ec::CurveGroup;
+use ark_ec::{pairing::Pairing, CurveGroup};
 use ark_ff::Field;
 use ark_relations::r1cs::SynthesisError;
 use core::fmt::Debug;