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