1use crate::{
  2    boolean::Boolean,
  3    fields::{fp::FpVar, FieldVar},
  4    prelude::*,
  5    ToBitsGadget,
  6};
  7use algebra::PrimeField;
  8use core::cmp::Ordering;
  9use r1cs_core::{lc, SynthesisError, Variable};
 10
 11impl<F: PrimeField> FpVar<F> {
 12    /// This function enforces the ordering between `self` and `other`. The
 13    /// constraint system will not be satisfied otherwise. If `self` should
 14    /// also be checked for equality, e.g. `self <= other` instead of `self < other`, set
 15    /// `should_also_check_quality` to `true`. This variant verifies `self` and `other`
 16    /// are `<= (p-1)/2`.
 17    pub fn enforce_cmp(
 18        &self,
 19        other: &FpVar<F>,
 20        ordering: Ordering,
 21        should_also_check_equality: bool,
 22    ) -> Result<(), SynthesisError> {
 23        let (left, right) = self.process_cmp_inputs(other, ordering, should_also_check_equality)?;
 24        left.enforce_smaller_than(&right)
 25    }
 26
 27    /// This function enforces the ordering between `self` and `other`. The
 28    /// constraint system will not be satisfied otherwise. If `self` should
 29    /// also be checked for equality, e.g. `self <= other` instead of `self < other`, set
 30    /// `should_also_check_quality` to `true`. This variant assumes `self` and `other`
 31    /// are `<= (p-1)/2` and does not generate constraints to verify that.
 32    pub fn enforce_cmp_unchecked(
 33        &self,
 34        other: &FpVar<F>,
 35        ordering: Ordering,
 36        should_also_check_equality: bool,
 37    ) -> Result<(), SynthesisError> {
 38        let (left, right) = self.process_cmp_inputs(other, ordering, should_also_check_equality)?;
 39        left.enforce_smaller_than_unchecked(&right)
 40    }
 41
 42    /// This function checks the ordering between `self` and `other`. It outputs self
 43    /// `Boolean` that contains the result - `1` if true, `0` otherwise. The
 44    /// constraint system will be satisfied in any case. If `self` should
 45    /// also be checked for equality, e.g. `self <= other` instead of `self < other`, set
 46    /// `should_also_check_quality` to `true`. This variant verifies `self` and `other`
 47    /// are `<= (p-1)/2`.
 48    pub fn is_cmp(
 49        &self,
 50        other: &FpVar<F>,
 51        ordering: Ordering,
 52        should_also_check_equality: bool,
 53    ) -> Result<Boolean<F>, SynthesisError> {
 54        let (left, right) = self.process_cmp_inputs(other, ordering, should_also_check_equality)?;
 55        left.is_smaller_than(&right)
 56    }
 57
 58    /// This function checks the ordering between `self` and `other`. It outputs a
 59    /// `Boolean` that contains the result - `1` if true, `0` otherwise. The
 60    /// constraint system will be satisfied in any case. If `self` should
 61    /// also be checked for equality, e.g. `self <= other` instead of `self < other`, set
 62    /// `should_also_check_quality` to `true`. This variant assumes `self` and `other`
 63    /// are `<= (p-1)/2` and does not generate constraints to verify that.
 64    pub fn is_cmp_unchecked(
 65        &self,
 66        other: &FpVar<F>,
 67        ordering: Ordering,
 68        should_also_check_equality: bool,
 69    ) -> Result<Boolean<F>, SynthesisError> {
 70        let (left, right) = self.process_cmp_inputs(other, ordering, should_also_check_equality)?;
 71        left.is_smaller_than_unchecked(&right)
 72    }
 73
 74    fn process_cmp_inputs(
 75        &self,
 76        other: &Self,
 77        ordering: Ordering,
 78        should_also_check_equality: bool,
 79    ) -> Result<(Self, Self), SynthesisError> {
 80        let (left, right) = match ordering {
 81            Ordering::Less => (self, other),
 82            Ordering::Greater => (other, self),
 83            Ordering::Equal => Err(SynthesisError::Unsatisfiable)?,
 84        };
 85        let right_for_check = if should_also_check_equality {
 86            right + F::one()
 87        } else {
 88            right.clone()
 89        };
 90
 91        Ok((left.clone(), right_for_check))
 92    }
 93
 94    // Helper function to enforce `self <= (p-1)/2`.
 95    pub fn enforce_smaller_or_equal_than_mod_minus_one_div_two(
 96        &self,
 97    ) -> Result<(), SynthesisError> {
 98        // It's okay to use `to_non_unique_bits` bits here because we're enforcing
 99        // self <= (p-1)/2, which implies self < p.
100        let _ = Boolean::enforce_smaller_or_equal_than_le(
101            &self.to_non_unique_bits_le()?,
102            F::modulus_minus_one_div_two(),
103        )?;
104        Ok(())
105    }
106
107    /// Helper function to check `self < other` and output a result bit. This function
108    /// verifies `self` and `other` are `<= (p-1)/2`.
109    fn is_smaller_than(&self, other: &FpVar<F>) -> Result<Boolean<F>, SynthesisError> {
110        self.enforce_smaller_or_equal_than_mod_minus_one_div_two()?;
111        other.enforce_smaller_or_equal_than_mod_minus_one_div_two()?;
112        self.is_smaller_than_unchecked(other)
113    }
114
115    /// Helper function to check `self < other` and output a result bit. This function
116    /// assumes `self` and `other` are `<= (p-1)/2` and does not generate constraints
117    /// to verify that.
118    fn is_smaller_than_unchecked(&self, other: &FpVar<F>) -> Result<Boolean<F>, SynthesisError> {
119        Ok((self - other)
120            .double()?
121            .to_bits_le()?
122            .first()
123            .unwrap()
124            .clone())
125    }
126
127    /// Helper function to enforce `self < other`. This function verifies `self` and `other`
128    /// are `<= (p-1)/2`.
129    fn enforce_smaller_than(&self, other: &FpVar<F>) -> Result<(), SynthesisError> {
130        self.enforce_smaller_or_equal_than_mod_minus_one_div_two()?;
131        other.enforce_smaller_or_equal_than_mod_minus_one_div_two()?;
132        self.enforce_smaller_than_unchecked(other)
133    }
134
135    /// Helper function to enforce `self < other`. This function assumes `self` and `other`
136    /// are `<= (p-1)/2` and does not generate constraints to verify that.
137    fn enforce_smaller_than_unchecked(&self, other: &FpVar<F>) -> Result<(), SynthesisError> {
138        let cs = [self, other].cs().unwrap();
139        let is_smaller_than = self.is_smaller_than_unchecked(other)?;
140        let lc_one = lc!() + Variable::One;
141        cs.enforce_constraint(is_smaller_than.lc(), lc_one.clone(), lc_one)
142    }
143}
144
145#[cfg(test)]
146mod test {
147    use rand::{Rng, SeedableRng};
148    use rand_xorshift::XorShiftRng;
149    use std::cmp::Ordering;
150
151    use crate::{alloc::AllocVar, fields::fp::FpVar};
152    use algebra::{bls12_381::Fr, PrimeField, UniformRand};
153    use r1cs_core::ConstraintSystem;
154
155    #[test]
156    fn test_cmp() {
157        let mut rng = &mut XorShiftRng::from_seed([
158            0x5d, 0xbe, 0x62, 0x59, 0x8d, 0x31, 0x3d, 0x76, 0x32, 0x37, 0xdb, 0x17, 0xe5, 0xbc,
159            0x06, 0x54,
160        ]);
161        fn rand_in_range<R: Rng>(rng: &mut R) -> Fr {
162            let pminusonedivtwo: Fr = Fr::modulus_minus_one_div_two().into();
163            let mut r;
164            loop {
165                r = Fr::rand(rng);
166                if r <= pminusonedivtwo {
167                    break;
168                }
169            }
170            r
171        }
172        for i in 0..10 {
173            let cs = ConstraintSystem::<Fr>::new_ref();
174            let a = rand_in_range(&mut rng);
175            let a_var = FpVar::<Fr>::new_witness(cs.ns("a"), || Ok(a)).unwrap();
176            let b = rand_in_range(&mut rng);
177            let b_var = FpVar::<Fr>::new_witness(cs.ns("b"), || Ok(b)).unwrap();
178
179            match a.cmp(&b) {
180                Ordering::Less => {
181                    a_var.enforce_cmp(&b_var, Ordering::Less, false).unwrap();
182                    a_var.enforce_cmp(&b_var, Ordering::Less, true).unwrap();
183                }
184                Ordering::Greater => {
185                    a_var.enforce_cmp(&b_var, Ordering::Greater, false).unwrap();
186                    a_var.enforce_cmp(&b_var, Ordering::Greater, true).unwrap();
187                }
188                _ => {}
189            }
190
191            if i == 0 {
192                println!("number of constraints: {}", cs.num_constraints());
193            }
194            assert!(cs.is_satisfied().unwrap());
195        }
196        println!("Finished with satisfaction tests");
197
198        for _i in 0..10 {
199            let cs = ConstraintSystem::<Fr>::new_ref();
200            let a = rand_in_range(&mut rng);
201            let a_var = FpVar::<Fr>::new_witness(cs.ns("a"), || Ok(a)).unwrap();
202            let b = rand_in_range(&mut rng);
203            let b_var = FpVar::<Fr>::new_witness(cs.ns("b"), || Ok(b)).unwrap();
204
205            match b.cmp(&a) {
206                Ordering::Less => {
207                    a_var.enforce_cmp(&b_var, Ordering::Less, false).unwrap();
208                    a_var.enforce_cmp(&b_var, Ordering::Less, true).unwrap();
209                }
210                Ordering::Greater => {
211                    a_var.enforce_cmp(&b_var, Ordering::Greater, false).unwrap();
212                    a_var.enforce_cmp(&b_var, Ordering::Greater, true).unwrap();
213                }
214                _ => {}
215            }
216
217            assert!(!cs.is_satisfied().unwrap());
218        }
219
220        for _i in 0..10 {
221            let cs = ConstraintSystem::<Fr>::new_ref();
222            let a = rand_in_range(&mut rng);
223            let a_var = FpVar::<Fr>::new_witness(cs.ns("a"), || Ok(a)).unwrap();
224            a_var.enforce_cmp(&a_var, Ordering::Less, false).unwrap();
225
226            assert!(!cs.is_satisfied().unwrap());
227        }
228
229        for _i in 0..10 {
230            let cs = ConstraintSystem::<Fr>::new_ref();
231            let a = rand_in_range(&mut rng);
232            let a_var = FpVar::<Fr>::new_witness(cs.ns("a"), || Ok(a)).unwrap();
233            a_var.enforce_cmp(&a_var, Ordering::Less, true).unwrap();
234            assert!(cs.is_satisfied().unwrap());
235        }
236    }
237}