From 0abb3a7dfe656b1c7bd85cf441d0d8b8a56af015 Mon Sep 17 00:00:00 2001
From: Nirvan Tyagi <nirvan.tyagi@gmail.com>
Date: Thu, 19 Nov 2020 22:32:21 -0500
Subject: [PATCH] CondSelectGadget for UInts

---
 src/bits/uint.rs  | 25 +++++++++++++++++++++++++
 src/bits/uint8.rs | 27 +++++++++++++++++++++++++++
 2 files changed, 52 insertions(+)

diff --git a/src/bits/uint.rs b/src/bits/uint.rs
index 8c92b31..20bbcd4 100644
--- a/src/bits/uint.rs
+++ b/src/bits/uint.rs
@@ -329,6 +329,31 @@ macro_rules! make_uint {
                 }
             }
 
+            impl<ConstraintF: Field> CondSelectGadget<ConstraintF> for $name<ConstraintF> {
+                #[tracing::instrument(target = "r1cs")]
+                fn conditionally_select(
+                    cond: &Boolean<ConstraintF>,
+                    true_value: &Self,
+                    false_value: &Self,
+                ) -> Result<Self, SynthesisError> {
+                    let selected_bits = true_value.bits.iter().zip(&false_value.bits)
+                        .map(|(true_bit, false_bit)| {
+                            cond.select(true_bit, false_bit)
+                        }).collect::<Result<Vec<Boolean<ConstraintF>>, SynthesisError>>()?;
+                    let selected_value = match (cond.value(), true_value.value(), false_value.value()) {
+                        (Ok(true), Err(_), _) => None,
+                        (Ok(true), Ok(v), _) => Some(v),
+                        (Ok(false), _, Err(_)) => None,
+                        (Ok(false), _, Ok(v)) => Some(v),
+                        (Err(_), _, _) => None,
+                    };
+                    Ok(Self {
+                        bits: selected_bits,
+                        value: selected_value,
+                    })
+                }
+            }
+
             impl<ConstraintF: Field> AllocVar<$native, ConstraintF> for $name<ConstraintF> {
                 fn new_variable<T: Borrow<$native>>(
                     cs: impl Into<Namespace<ConstraintF>>,
diff --git a/src/bits/uint8.rs b/src/bits/uint8.rs
index bc50eb5..512fbab 100644
--- a/src/bits/uint8.rs
+++ b/src/bits/uint8.rs
@@ -284,6 +284,33 @@ impl<ConstraintF: Field> EqGadget<ConstraintF> for UInt8<ConstraintF> {
     }
 }
 
+impl<ConstraintF: Field> CondSelectGadget<ConstraintF> for UInt8<ConstraintF> {
+    #[tracing::instrument(target = "r1cs")]
+    fn conditionally_select(
+        cond: &Boolean<ConstraintF>,
+        true_value: &Self,
+        false_value: &Self,
+    ) -> Result<Self, SynthesisError> {
+        let selected_bits = true_value
+            .bits
+            .iter()
+            .zip(&false_value.bits)
+            .map(|(true_bit, false_bit)| cond.select(true_bit, false_bit))
+            .collect::<Result<Vec<Boolean<ConstraintF>>, SynthesisError>>()?;
+        let selected_value = match (cond.value(), true_value.value(), false_value.value()) {
+            (Ok(true), Err(_), _) => None,
+            (Ok(true), Ok(v), _) => Some(v),
+            (Ok(false), _, Err(_)) => None,
+            (Ok(false), _, Ok(v)) => Some(v),
+            (Err(_), _, _) => None,
+        };
+        Ok(Self {
+            bits: selected_bits,
+            value: selected_value,
+        })
+    }
+}
+
 impl<ConstraintF: Field> AllocVar<u8, ConstraintF> for UInt8<ConstraintF> {
     fn new_variable<T: Borrow<u8>>(
         cs: impl Into<Namespace<ConstraintF>>,