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

benches/bench.rs

@@ -1,14 +1,18 @@
-use ark_ff::PrimeField;
+use ark_ec::short_weierstrass::Projective;
+use ark_ff::{PrimeField, UniformRand};
 use ark_r1cs_std::{
     alloc::AllocVar,
     eq::EqGadget,
-    fields::{nonnative::NonNativeFieldVar, FieldVar},
+    fields::{fp::FpVar, nonnative::NonNativeFieldVar, FieldVar},
+    groups::{curves::short_weierstrass::ProjectiveVar, CurveVar},
+    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 +166,64 @@ fn inverse<TargetField: PrimeField, BaseField: PrimeField, R: RngCore>(
     );
 }
 
+macro_rules! cond_select_bench_individual {
+    ($bench_name:ident, $bench_base_field:ty, $var:ty, $native_type:ty, $constant:ident) => {
+        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) = {
+                // value array
+                let values: Vec<$native_type> =
+                    (0..128).map(|_| <$native_type>::rand(rng)).collect();
+                let values_const: Vec<$var> =
+                    values.iter().map(|x| <$var>::$constant(*x)).collect();
+
+                // index array
+                let position: Vec<bool> = (0..7).map(|_| rng.gen()).collect();
+                let position_var: Vec<Boolean<$bench_base_field>> = 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 _ =
+                    <$var>::conditionally_select_power_of_two_vector(&position_var, &values_const);
+
+                let constraints_after = cs.num_constraints();
+                let nonzeros_after = get_density(&cs);
+
+                (
+                    constraints_after - constraints_before,
+                    nonzeros_after - nonzeros_before,
+                )
+            };
+
+            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!(
+            "{} takes: {} constraints, {} non-zeros",
+            stringify!($bench_name),
+            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 +297,26 @@ 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!(
+        FpVar_select,
+        ark_mnt6_753::Fr,
+        FpVar<ark_mnt6_753::Fr>,
+        ark_mnt6_753::Fr,
+        Constant
+    );
+    cond_select_bench_individual!(
+        Boolean_select,
+        ark_mnt6_753::Fr,
+        Boolean<ark_mnt6_753::Fr>,
+        bool,
+        Constant
+    );
+    pub type FBaseVar = FpVar<ark_mnt6_753::Fq>;
+    cond_select_bench_individual!(
+        SWProjective_select,
+        ark_mnt6_753::Fq,
+        ProjectiveVar<ark_mnt6_753::g1::Config, FBaseVar>,
+        Projective<ark_mnt6_753::g1::Config>,
+        constant
+    );
 }

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/bits/boolean.rs

@@ -216,6 +216,56 @@ impl<F: Field> CondSelectGadget<F> for AllocatedBool<F> {
             _ => unreachable!("Impossible"),
         }
     }
+
+    /// Using the hybrid method 5.3 from <https://github.com/mir-protocol/r1cs-workshop/blob/master/workshop.pdf>.
+    fn conditionally_select_power_of_two_vector(
+        position: &[Boolean<F>],
+        values: &[Self],
+    ) -> Result<Self, SynthesisError> {
+        let cs = position[0].cs();
+        let n = position.len();
+
+        // split n into l and m, where l + m = n
+        // total cost is 2^m + 2^l - l - 2, so we'd rather maximize l than m
+        let m = n / 2;
+        let l = n - m;
+
+        let two_to_l = 1 << l;
+        let two_to_m = 1 << m;
+
+        // we only need the lower L bits
+        let lower_bits = &mut position[m..].to_vec();
+        let sub_tree = sum_of_conditions(lower_bits)?;
+
+        // index for the chunk
+        let mut index = 0;
+        for x in lower_bits {
+            index *= 2;
+            index += if x.value()? { 1 } else { 0 };
+        }
+        let chunk_size = 1 << l;
+        let root_vals: Vec<Self> = values
+            .chunks(chunk_size)
+            .map(|chunk| chunk[index].clone())
+            .collect();
+
+        let upper_elems = {
+            for i in 0..two_to_m {
+                let mut x = LinearCombination::zero();
+                for j in 0..two_to_l {
+                    let v = values[i * two_to_l + j].value()?;
+                    x = &x + sub_tree[j].clone() * v.into();
+                }
+                cs.enforce_constraint(x, lc!() + Variable::One, lc!() + root_vals[i].variable)?;
+            }
+
+            Ok(root_vals)
+        }?;
+
+        // apply the repeated selection method, to select one of 2^m subtree results
+        let upper_bits = &mut position[..m].to_vec();
+        repeated_selection(upper_bits, upper_elems)
+    }
 }
 
 /// Represents a boolean value in the constraint system which is guaranteed
@@ -950,6 +1000,56 @@ impl<F: Field> CondSelectGadget<F> for Boolean<F> {
             },
         }
     }
+
+    /// Using the hybrid method 5.3 from <https://github.com/mir-protocol/r1cs-workshop/blob/master/workshop.pdf>.
+    fn conditionally_select_power_of_two_vector(
+        position: &[Boolean<F>],
+        values: &[Self],
+    ) -> Result<Self, SynthesisError> {
+        let cs = position[0].cs();
+        let n = position.len();
+
+        // split n into l and m, where l + m = n
+        // total cost is 2^m + 2^l - l - 2, so we'd rather maximize l than m
+        let m = n / 2;
+        let l = n - m;
+
+        let two_to_l = 1 << l;
+        let two_to_m = 1 << m;
+
+        // we only need the lower L bits
+        let lower_bits = &mut position[m..].to_vec();
+        let sub_tree = sum_of_conditions(lower_bits)?;
+
+        // index for the chunk
+        let mut index = 0;
+        for x in lower_bits {
+            index *= 2;
+            index += if x.value()? { 1 } else { 0 };
+        }
+        let chunk_size = 1 << l;
+        let root_vals: Vec<Self> = values
+            .chunks(chunk_size)
+            .map(|chunk| chunk[index].clone())
+            .collect();
+
+        let upper_elems = {
+            for i in 0..two_to_m {
+                let mut x = LinearCombination::zero();
+                for j in 0..two_to_l {
+                    let v = values[i * two_to_l + j].value()?;
+                    x = &x + sub_tree[j].clone() * v.into();
+                }
+                cs.enforce_constraint(x, lc!() + Variable::One, root_vals[i].lc())?;
+            }
+
+            Ok(root_vals)
+        }?;
+
+        // apply the repeated selection method, to select one of 2^m subtree results
+        let upper_bits = &mut position[..m].to_vec();
+        repeated_selection(upper_bits, upper_elems)
+    }
 }
 
 #[cfg(test)]
@@ -958,6 +1058,7 @@ mod test {
     use crate::prelude::*;
     use ark_ff::{BitIteratorBE, BitIteratorLE, Field, One, PrimeField, UniformRand, Zero};
     use ark_relations::r1cs::{ConstraintSystem, Namespace, SynthesisError};
+    use ark_std::rand::Rng;
     use ark_test_curves::bls12_381::Fr;
 
     #[test]
@@ -1818,4 +1919,43 @@ mod test {
 
         Ok(())
     }
+
+    #[test]
+    fn test_bool_random_access() {
+        let mut rng = ark_std::test_rng();
+
+        for _ in 0..100 {
+            let cs = ConstraintSystem::<Fr>::new_ref();
+
+            // value array
+            let values: Vec<bool> = (0..128).map(|_| rng.gen()).collect();
+            let values_const: Vec<Boolean<Fr>> =
+                values.iter().map(|x| Boolean::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!(
+                Boolean::conditionally_select_power_of_two_vector(&position_var, &values_const)
+                    .unwrap()
+                    .value()
+                    .unwrap(),
+                values[index]
+            )
+        }
+    }
 }