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 that `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 is_smaller_than = self.is_smaller_than_unchecked(other)?;
144        let lc_one = lc!() + Variable::One;
145        [self, other]
146            .cs()
147            .enforce_constraint(is_smaller_than.lc(), lc_one.clone(), lc_one)
148    }
149}
150
151#[cfg(test)]
152mod test {
153    use rand::{Rng, SeedableRng};
154    use rand_xorshift::XorShiftRng;
155    use std::cmp::Ordering;
156
157    use crate::{alloc::AllocVar, fields::fp::FpVar};
158    use algebra::{bls12_381::Fr, PrimeField, UniformRand};
159    use r1cs_core::ConstraintSystem;
160
161    #[test]
162    fn test_cmp() {
163        let mut rng = &mut XorShiftRng::from_seed([
164            0x5d, 0xbe, 0x62, 0x59, 0x8d, 0x31, 0x3d, 0x76, 0x32, 0x37, 0xdb, 0x17, 0xe5, 0xbc,
165            0x06, 0x54,
166        ]);
167        fn rand_in_range<R: Rng>(rng: &mut R) -> Fr {
168            let pminusonedivtwo: Fr = Fr::modulus_minus_one_div_two().into();
169            let mut r;
170            loop {
171                r = Fr::rand(rng);
172                if r <= pminusonedivtwo {
173                    break;
174                }
175            }
176            r
177        }
178        for i in 0..10 {
179            let cs = ConstraintSystem::<Fr>::new_ref();
180            let a = rand_in_range(&mut rng);
181            let a_var = FpVar::<Fr>::new_witness(cs.clone(), || Ok(a)).unwrap();
182            let b = rand_in_range(&mut rng);
183            let b_var = FpVar::<Fr>::new_witness(cs.clone(), || Ok(b)).unwrap();
184
185            match a.cmp(&b) {
186                Ordering::Less => {
187                    a_var.enforce_cmp(&b_var, Ordering::Less, false).unwrap();
188                    a_var.enforce_cmp(&b_var, Ordering::Less, true).unwrap();
189                }
190                Ordering::Greater => {
191                    a_var.enforce_cmp(&b_var, Ordering::Greater, false).unwrap();
192                    a_var.enforce_cmp(&b_var, Ordering::Greater, true).unwrap();
193                }
194                _ => {}
195            }
196
197            if i == 0 {
198                println!("number of constraints: {}", cs.num_constraints());
199            }
200            assert!(cs.is_satisfied().unwrap());
201        }
202        println!("Finished with satisfaction tests");
203
204        for _i in 0..10 {
205            let cs = ConstraintSystem::<Fr>::new_ref();
206            let a = rand_in_range(&mut rng);
207            let a_var = FpVar::<Fr>::new_witness(cs.clone(), || Ok(a)).unwrap();
208            let b = rand_in_range(&mut rng);
209            let b_var = FpVar::<Fr>::new_witness(cs.clone(), || Ok(b)).unwrap();
210
211            match b.cmp(&a) {
212                Ordering::Less => {
213                    a_var.enforce_cmp(&b_var, Ordering::Less, false).unwrap();
214                    a_var.enforce_cmp(&b_var, Ordering::Less, true).unwrap();
215                }
216                Ordering::Greater => {
217                    a_var.enforce_cmp(&b_var, Ordering::Greater, false).unwrap();
218                    a_var.enforce_cmp(&b_var, Ordering::Greater, true).unwrap();
219                }
220                _ => {}
221            }
222
223            assert!(!cs.is_satisfied().unwrap());
224        }
225
226        for _i in 0..10 {
227            let cs = ConstraintSystem::<Fr>::new_ref();
228            let a = rand_in_range(&mut rng);
229            let a_var = FpVar::<Fr>::new_witness(cs.clone(), || Ok(a)).unwrap();
230            a_var.enforce_cmp(&a_var, Ordering::Less, false).unwrap();
231
232            assert!(!cs.is_satisfied().unwrap());
233        }
234
235        for _i in 0..10 {
236            let cs = ConstraintSystem::<Fr>::new_ref();
237            let a = rand_in_range(&mut rng);
238            let a_var = FpVar::<Fr>::new_witness(cs.clone(), || Ok(a)).unwrap();
239            a_var.enforce_cmp(&a_var, Ordering::Less, true).unwrap();
240            assert!(cs.is_satisfied().unwrap());
241        }
242    }
243}