Import uint library

codes.py

@@ -22,6 +22,10 @@
     pub fn get_value(&self) -> Option<S> {
         self.val.clone()
     }
+    pub fn add_constant<CS: ConstraintSystem<S>>(&mut self, num: S) {
+        self.lc = self.lc.clone() + (num, CS::one());
+        self.val = self.val.map(|v| v + num);
+    }
     pub fn alloc<CS: ConstraintSystem<S>>(
         cs: &mut CS,
         val: Option<S>,
@@ -45,6 +49,17 @@
     }
 }
 
+impl<S: PrimeField> Sub<&Number<S>> for &Number<S> {
+    type Output = Number<S>;
+
+    fn sub(self, other: &Number<S>) -> Self::Output {
+        Self::Output {
+            lc: self.lc.clone() - &other.lc,
+            val: self.val.zip(other.val).map(|(slf, othr)| slf - othr),
+        }
+    }
+}
+
 impl<S: PrimeField> Add<(S, &Number<S>)> for &Number<S> {
     type Output = Number<S>;
 
@@ -59,3 +74,122 @@
     }
 }
 """
+
+UINT_RS = """
+use super::Number;
+use bellman::gadgets::boolean::{AllocatedBit, Boolean};
+use bellman::{ConstraintSystem, LinearCombination, SynthesisError};
+use ff::{PrimeField, PrimeFieldBits};
+
+#[derive(Clone)]
+pub struct UnsignedInteger<S: PrimeField + PrimeFieldBits> {
+    bits: Vec<AllocatedBit>,
+    num: Number<S>,
+}
+
+impl<S: PrimeField + PrimeFieldBits> UnsignedInteger<S> {
+    pub fn get_lc(&self) -> &LinearCombination<S> {
+        self.num.get_lc()
+    }
+    pub fn get_number(&self) -> &Number<S> {
+        &self.num
+    }
+    pub fn get_value(&self) -> Option<S> {
+        self.num.get_value()
+    }
+    pub fn bits(&self) -> &Vec<AllocatedBit> {
+        &self.bits
+    }
+    pub fn num_bits(&self) -> usize {
+        self.bits.len()
+    }
+    pub fn alloc<CS: ConstraintSystem<S>>(
+        cs: &mut CS,
+        val: S,
+        bits: usize,
+    ) -> Result<Self, SynthesisError> {
+        let alloc = Number::<S>::alloc(&mut *cs, Some(val))?;
+        Self::constrain(cs, alloc.into(), bits)
+    }
+    pub fn alloc_32<CS: ConstraintSystem<S>>(
+        cs: &mut CS,
+        val: u32,
+    ) -> Result<Self, SynthesisError> {
+        Self::alloc(cs, S::from(val as u64), 32)
+    }
+    pub fn alloc_64<CS: ConstraintSystem<S>>(
+        cs: &mut CS,
+        val: u64,
+    ) -> Result<Self, SynthesisError> {
+        Self::alloc(cs, S::from(val), 64)
+    }
+    pub fn constrain<CS: ConstraintSystem<S>>(
+        cs: &mut CS,
+        num: Number<S>,
+        num_bits: usize,
+    ) -> Result<Self, SynthesisError> {
+        let mut bits = Vec::new();
+        let mut coeff = S::one();
+        let mut all = LinearCombination::<S>::zero();
+        let bit_vals: Option<Vec<bool>> = num
+            .get_value()
+            .map(|v| v.to_le_bits().iter().map(|b| *b).collect());
+        for i in 0..num_bits {
+            let bit = AllocatedBit::alloc(&mut *cs, bit_vals.as_ref().map(|b| b[i]))?;
+            all = all + (coeff, bit.get_variable());
+            bits.push(bit);
+            coeff = coeff.double();
+        }
+        cs.enforce(
+            || "check",
+            |lc| lc + &all,
+            |lc| lc + CS::one(),
+            |lc| lc + num.get_lc(),
+        );
+
+        Ok(Self { num, bits })
+    }
+
+    // ~198 constraints
+    pub fn lt<CS: ConstraintSystem<S>>(
+        &self,
+        cs: &mut CS,
+        other: &UnsignedInteger<S>,
+    ) -> Result<Boolean, SynthesisError> {
+        assert_eq!(self.num_bits(), other.num_bits());
+        let num_bits = self.num_bits();
+
+        // Imagine a and b are two sigend (num_bits + 1) bits numbers
+        let two_bits = S::from(2).pow_vartime(&[num_bits as u64 + 1, 0, 0, 0]);
+        let mut sub = &self.num - &other.num;
+        sub.add_constant::<CS>(two_bits);
+
+        let sub_bits = UnsignedInteger::<S>::constrain(&mut *cs, sub, num_bits + 2)?;
+        Ok(Boolean::Is(sub_bits.bits()[num_bits].clone()))
+    }
+
+    pub fn gt<CS: ConstraintSystem<S>>(
+        &self,
+        cs: &mut CS,
+        other: &UnsignedInteger<S>,
+    ) -> Result<Boolean, SynthesisError> {
+        other.lt(cs, self)
+    }
+
+    pub fn lte<CS: ConstraintSystem<S>>(
+        &self,
+        cs: &mut CS,
+        other: &UnsignedInteger<S>,
+    ) -> Result<Boolean, SynthesisError> {
+        Ok(self.gt(cs, other)?.not())
+    }
+
+    pub fn gte<CS: ConstraintSystem<S>>(
+        &self,
+        cs: &mut CS,
+        other: &UnsignedInteger<S>,
+    ) -> Result<Boolean, SynthesisError> {
+        Ok(self.lt(cs, other)?.not())
+    }
+}
+"""

main.py

@@ -17,8 +17,14 @@
 rand = "0.8.5"
 """
 
-MAIN_RS_TEMPLATE = """mod number;
+MAIN_RS_TEMPLATE = """
+mod number;
+mod uint;
+#[allow(unused_imports)]
 use number::Number;
+#[allow(unused_imports)]
+use uint::UnsignedInteger;
+
 use bellman::groth16::{
     create_random_proof, generate_random_parameters, prepare_verifying_key, verify_proof,
 };
@@ -183,7 +189,7 @@ def mul(self, a, b):
         return ret
 
     def compile(self, path, project_name):
-        libs = {"number": codes.NUMBER_RS}
+        libs = {"number": codes.NUMBER_RS, "uint": codes.UINT_RS}
 
         proj = os.path.join(path, project_name)
         cargo_toml = os.path.join(proj, "Cargo.toml")
@@ -194,7 +200,7 @@ def compile(self, path, project_name):
 
         for (file_name, content) in libs.items():
             with io.open(os.path.join(src, file_name + ".rs"), "w") as f:
-                f.write(content)
+                f.write("#![allow(dead_code)]\n" + content)
 
         main_rs = os.path.join(src, "main.rs")