From 05fee346868f910ff28db35b55b1c63ca0d77541 Mon Sep 17 00:00:00 2001
From: Tom Shen <tomshen@berkeley.edu>
Date: Thu, 4 Feb 2021 13:52:29 -0800
Subject: [PATCH] Add Mux gadget (#48)

---
 CHANGELOG.md      |  3 ++-
 src/bits/uint8.rs | 38 ++++++++++++++++++++++++++++++++++++++
 src/select.rs     | 47 ++++++++++++++++++++++++++++++++++++++++++++++-
 3 files changed, 86 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7301fbe..809e1d0 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,7 +2,7 @@
 
 ### Breaking changes
 - #12 Make the output of the `ToBitsGadget` impl for `FpVar` fixed-size
-
+- #48 Add `Clone` trait bound to `CondSelectGadget`.
 ### Features
 
 
@@ -18,6 +18,7 @@
 - #35 Construct a `FpVar` from bits
 - #36 Implement `ToConstraintFieldGadget` for `Vec<Uint8>`
 - #40, #43 Faster scalar multiplication for Short Weierstrass curves by relying on affine formulae
+- #46 Add mux gadget as an auto-impl in `CondSelectGadget` to support random access of an array
 
 ### Bug fixes
 - #8 Fix bug in `three_bit_cond_neg_lookup` when using a constant lookup bit
diff --git a/src/bits/uint8.rs b/src/bits/uint8.rs
index d5ee2b6..69e2eeb 100644
--- a/src/bits/uint8.rs
+++ b/src/bits/uint8.rs
@@ -503,4 +503,42 @@ mod test {
 
         Ok(())
     }
+
+    #[test]
+    fn test_uint8_random_access() {
+        let mut rng = XorShiftRng::seed_from_u64(1231275789u64);
+
+        for _ in 0..100 {
+            let cs = ConstraintSystem::<Fr>::new_ref();
+
+            // value array
+            let values: Vec<u8> = (0..128).map(|_| rng.gen()).collect();
+            let values_const: Vec<UInt8<Fr>> = values.iter().map(|x| UInt8::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!(
+                UInt8::conditionally_select_power_of_two_vector(&position_var, &values_const)
+                    .unwrap()
+                    .value()
+                    .unwrap(),
+                values[index]
+            )
+        }
+    }
 }
diff --git a/src/select.rs b/src/select.rs
index f79bb1e..bbc2c3c 100644
--- a/src/select.rs
+++ b/src/select.rs
@@ -1,11 +1,12 @@
 use crate::prelude::*;
 use ark_ff::Field;
 use ark_relations::r1cs::SynthesisError;
-
+use ark_std::vec::Vec;
 /// Generates constraints for selecting between one of two values.
 pub trait CondSelectGadget<ConstraintF: Field>
 where
     Self: Sized,
+    Self: Clone,
 {
     /// If `cond == &Boolean::TRUE`, then this returns `true_value`; else,
     /// returns `false_value`.
@@ -18,6 +19,50 @@ where
         true_value: &Self,
         false_value: &Self,
     ) -> Result<Self, SynthesisError>;
+
+    /// Returns an element of `values` whose index in represented by `position`.
+    /// `position` is an array of boolean that represents an unsigned integer in big endian order.
+    ///
+    /// # Example
+    /// To get the 6th element of `values`, convert unsigned integer 6 (`0b110`) to `position = [True, True, False]`,
+    /// and call `conditionally_select_power_of_two_vector(position, values)`.
+    fn conditionally_select_power_of_two_vector(
+        position: &[Boolean<ConstraintF>],
+        values: &[Self],
+    ) -> Result<Self, SynthesisError> {
+        let m = values.len();
+        let n = position.len();
+
+        // Assert m is a power of 2, and n = log(m)
+        assert!(m.is_power_of_two());
+        assert_eq!(1 << n, m);
+
+        let mut cur_mux_values = values.to_vec();
+
+        // Traverse the evaluation tree from bottom to top in level order traversal.
+        // This is method 5.1 from https://github.com/mir-protocol/r1cs-workshop/blob/master/workshop.pdf
+        // TODO: Add method 5.2/5.3
+        for i in 0..n {
+            // Size of current layer.
+            let cur_size = 1 << (n - i);
+            assert_eq!(cur_mux_values.len(), cur_size);
+
+            let mut next_mux_values = Vec::new();
+            for j in (0..cur_size).step_by(2) {
+                let cur = Self::conditionally_select(
+                    &position[n - 1 - i],
+                    // true case
+                    &cur_mux_values[j + 1],
+                    // false case
+                    &cur_mux_values[j],
+                )?;
+                next_mux_values.push(cur);
+            }
+            cur_mux_values = next_mux_values;
+        }
+
+        Ok(cur_mux_values[0].clone())
+    }
 }
 
 /// Performs a lookup in a 4-element table using two bits.