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