M r1cs-std/src/bits/uint64.rs → r1cs-std/src/bits/uint64.rs

@@ -7,13 +7,14 @@ boolean::{AllocatedBit, Boolean},
     prelude::*,
     Assignment, Vec,
 };
+use core::borrow::Borrow;
 
 /// Represents an interpretation of 64 `Boolean` objects as an
 /// unsigned integer.
 #[derive(Clone, Debug)]
 pub struct UInt64 {
     // Least significant bit_gadget first
-    bits:  Vec<Boolean>,
+    bits: Vec<Boolean>,
     value: Option<u64>,
 }
 
@@ -40,7 +41,7 @@ }
     }
 
     /// Allocate a `UInt64` in the constraint system
-    pub fn alloc<ConstraintF, CS>(mut cs: CS, value: Option<u64>) -> Result<Self, SynthesisError>
+    pub fn _alloc<ConstraintF, CS>(mut cs: CS, value: Option<u64>) -> Result<Self, SynthesisError>
     where
         ConstraintF: Field,
         CS: ConstraintSystem<ConstraintF>,
@@ -55,7 +56,7 @@ val >>= 1;
                 }
 
                 v
-            },
+            }
             None => vec![None; 64],
         };
 
@@ -94,19 +95,19 @@ &Boolean::Constant(b) => {
                     if b {
                         value.as_mut().map(|v| *v |= 1);
                     }
-                },
+                }
                 &Boolean::Is(ref b) => match b.get_value() {
                     Some(true) => {
                         value.as_mut().map(|v| *v |= 1);
-                    },
-                    Some(false) => {},
+                    }
+                    Some(false) => {}
                     None => value = None,
                 },
                 &Boolean::Not(ref b) => match b.get_value() {
                     Some(false) => {
                         value.as_mut().map(|v| *v |= 1);
-                    },
-                    Some(true) => {},
+                    }
+                    Some(true) => {}
                     None => value = None,
                 },
             }
@@ -128,7 +129,7 @@ .cloned()
             .collect();
 
         UInt64 {
-            bits:  new_bits,
+            bits: new_bits,
             value: self.value.map(|v| v.rotate_right(by as u32)),
         }
     }
@@ -193,12 +194,12 @@ // Accumulate the value
             match op.value {
                 Some(val) => {
                     result_value.as_mut().map(|v| *v += u128::from(val));
-                },
+                }
                 None => {
                     // If any of our operands have unknown value, we won't
                     // know the value of the result
                     result_value = None;
-                },
+                }
             }
 
             // Iterate over each bit_gadget of the operand and add the operand to
@@ -211,18 +212,18 @@ all_constants = false;
 
                         // Add coeff * bit_gadget
                         lc += (coeff, bit.get_variable());
-                    },
+                    }
                     Boolean::Not(ref bit) => {
                         all_constants = false;
 
                         // Add coeff * (1 - bit_gadget) = coeff * ONE - coeff * bit_gadget
                         lc = lc + (coeff, CS::one()) - (coeff, bit.get_variable());
-                    },
+                    }
                     Boolean::Constant(bit) => {
                         if bit {
                             lc += (coeff, CS::one());
                         }
-                    },
+                    }
                 }
 
                 coeff.double_in_place();
@@ -269,12 +270,39 @@ // Discard carry bits that we don't care about
         result_bits.truncate(64);
 
         Ok(UInt64 {
-            bits:  result_bits,
+            bits: result_bits,
             value: modular_value,
         })
     }
 }
 
+impl<ConstraintF: Field> AllocGadget<u64, ConstraintF> for UInt64 {
+    fn alloc<F, T, CS: ConstraintSystem<ConstraintF>>(
+        mut cs: CS,
+        value_gen: F,
+    ) -> Result<Self, SynthesisError>
+    where
+        F: FnOnce() -> Result<T, SynthesisError>,
+        T: Borrow<u64>,
+    {
+        let val = value_gen()?.borrow().clone();
+
+        Self::_alloc(&mut cs.ns(|| "alloc u64"), Some(val))
+    }
+
+    fn alloc_input<F, T, CS: ConstraintSystem<ConstraintF>>(
+        mut cs: CS,
+        value_gen: F,
+    ) -> Result<Self, SynthesisError>
+    where
+        F: FnOnce() -> Result<T, SynthesisError>,
+        T: Borrow<u64>,
+    {
+        let val = value_gen()?.borrow().clone();
+        Self::_alloc(&mut cs.ns(|| "alloc u64"), Some(val))
+    }
+}
+
 impl<ConstraintF: Field> ToBytesGadget<ConstraintF> for UInt64 {
     #[inline]
     fn to_bytes<CS: ConstraintSystem<ConstraintF>>(
@@ -302,7 +330,7 @@ };
         let mut bytes = Vec::new();
         for (i, chunk8) in self.to_bits_le().chunks(8).enumerate() {
             let byte = UInt8 {
-                bits:  chunk8.to_vec(),
+                bits: chunk8.to_vec(),
                 value: value_chunks[i],
             };
             bytes.push(byte);
@@ -310,6 +338,13 @@ }
 
         Ok(bytes)
     }
+
+    fn to_bytes_strict<CS: ConstraintSystem<ConstraintF>>(
+        &self,
+        cs: CS,
+    ) -> Result<Vec<UInt8>, SynthesisError> {
+        self.to_bytes(cs)
+    }
 }
 
 impl PartialEq for UInt64 {
@@ -345,7 +380,10 @@
 #[cfg(test)]
 mod test {
     use super::UInt64;
-    use crate::{bits::boolean::Boolean, test_constraint_system::TestConstraintSystem, Vec};
+    use crate::{
+        alloc::AllocGadget, bits::boolean::Boolean, test_constraint_system::TestConstraintSystem,
+        Vec,
+    };
     use algebra::{bls12_381::Fr, One, Zero};
     use r1cs_core::ConstraintSystem;
     use rand::{Rng, SeedableRng};
@@ -366,7 +404,7 @@ for (i, bit_gadget) in b.bits.iter().enumerate() {
                 match bit_gadget {
                     &Boolean::Constant(bit_gadget) => {
                         assert!(bit_gadget == ((b.value.unwrap() >> i) & 1 == 1));
-                    },
+                    }
                     _ => unreachable!(),
                 }
             }
@@ -375,8 +413,8 @@ let expected_to_be_same = b.to_bits_le();
 
             for x in v.iter().zip(expected_to_be_same.iter()) {
                 match x {
-                    (&Boolean::Constant(true), &Boolean::Constant(true)) => {},
-                    (&Boolean::Constant(false), &Boolean::Constant(false)) => {},
+                    (&Boolean::Constant(true), &Boolean::Constant(true)) => {}
+                    (&Boolean::Constant(false), &Boolean::Constant(false)) => {}
                     _ => unreachable!(),
                 }
             }
@@ -396,9 +434,9 @@ let c: u64 = rng.gen();
 
             let mut expected = a ^ b ^ c;
 
-            let a_bit = UInt64::alloc(cs.ns(|| "a_bit"), Some(a)).unwrap();
+            let a_bit = UInt64::alloc(cs.ns(|| "a_bit"), || Ok(a)).unwrap();
             let b_bit = UInt64::constant(b);
-            let c_bit = UInt64::alloc(cs.ns(|| "c_bit"), Some(c)).unwrap();
+            let c_bit = UInt64::alloc(cs.ns(|| "c_bit"), || Ok(c)).unwrap();
 
             let r = a_bit.xor(cs.ns(|| "first xor"), &b_bit).unwrap();
             let r = r.xor(cs.ns(|| "second xor"), &c_bit).unwrap();
@@ -411,13 +449,13 @@ for b in r.bits.iter() {
                 match b {
                     &Boolean::Is(ref b) => {
                         assert!(b.get_value().unwrap() == (expected & 1 == 1));
-                    },
+                    }
                     &Boolean::Not(ref b) => {
                         assert!(!b.get_value().unwrap() == (expected & 1 == 1));
-                    },
+                    }
                     &Boolean::Constant(b) => {
                         assert!(b == (expected & 1 == 1));
-                    },
+                    }
                 }
 
                 expected >>= 1;
@@ -452,7 +490,7 @@ &Boolean::Is(_) => panic!(),
                     &Boolean::Not(_) => panic!(),
                     &Boolean::Constant(b) => {
                         assert!(b == (expected & 1 == 1));
-                    },
+                    }
                 }
 
                 expected >>= 1;
@@ -474,10 +512,10 @@ let d: u64 = rng.gen();
 
             let mut expected = (a ^ b).wrapping_add(c).wrapping_add(d);
 
-            let a_bit = UInt64::alloc(cs.ns(|| "a_bit"), Some(a)).unwrap();
+            let a_bit = UInt64::alloc(cs.ns(|| "a_bit"), || Ok(a)).unwrap();
             let b_bit = UInt64::constant(b);
             let c_bit = UInt64::constant(c);
-            let d_bit = UInt64::alloc(cs.ns(|| "d_bit"), Some(d)).unwrap();
+            let d_bit = UInt64::alloc(cs.ns(|| "d_bit"), || Ok(d)).unwrap();
 
             let r = a_bit.xor(cs.ns(|| "xor"), &b_bit).unwrap();
             let r = UInt64::addmany(cs.ns(|| "addition"), &[r, c_bit, d_bit]).unwrap();
@@ -490,10 +528,10 @@ for b in r.bits.iter() {
                 match b {
                     &Boolean::Is(ref b) => {
                         assert!(b.get_value().unwrap() == (expected & 1 == 1));
-                    },
+                    }
                     &Boolean::Not(ref b) => {
                         assert!(!b.get_value().unwrap() == (expected & 1 == 1));
-                    },
+                    }
                     &Boolean::Constant(_) => unreachable!(),
                 }
 
@@ -529,7 +567,7 @@ for b in &b.bits {
                 match b {
                     &Boolean::Constant(b) => {
                         assert_eq!(b, tmp & 1 == 1);
-                    },
+                    }
                     _ => unreachable!(),
                 }