1use ark_ff::PrimeField;
  2use ark_relations::r1cs::SynthesisError;
  3use ark_std::{ops::BitOr, ops::BitOrAssign};
  4
  5use super::{PrimUInt, UInt};
  6
  7impl<const N: usize, T: PrimUInt, F: PrimeField> UInt<N, T, F> {
  8    fn _or(&self, other: &Self) -> Result<Self, SynthesisError> {
  9        let mut result = self.clone();
 10        for (a, b) in result.bits.iter_mut().zip(&other.bits) {
 11            *a |= b;
 12        }
 13        result.value = self.value.and_then(|a| Some(a | other.value?));
 14        Ok(result)
 15    }
 16}
 17
 18impl<'a, const N: usize, T: PrimUInt, F: PrimeField> BitOr<Self> for &'a UInt<N, T, F> {
 19    type Output = UInt<N, T, F>;
 20
 21    /// Output `self | other`.
 22    ///
 23    /// If at least one of `self` and `other` are constants, then this method
 24    /// *does not* create any constraints or variables.
 25    ///
 26    /// ```
 27    /// # fn main() -> Result<(), ark_relations::r1cs::SynthesisError> {
 28    /// // We'll use the BLS12-381 scalar field for our constraints.
 29    /// use ark_test_curves::bls12_381::Fr;
 30    /// use ark_relations::r1cs::*;
 31    /// use ark_r1cs_std::prelude::*;
 32    ///
 33    /// let cs = ConstraintSystem::<Fr>::new_ref();
 34    /// let a = UInt8::new_witness(cs.clone(), || Ok(16))?;
 35    /// let b = UInt8::new_witness(cs.clone(), || Ok(17))?;
 36    /// let c = UInt8::new_witness(cs.clone(), || Ok(16 | 17))?;
 37    ///
 38    /// (a | b).enforce_equal(&c)?;
 39    /// assert!(cs.is_satisfied().unwrap());
 40    /// # Ok(())
 41    /// # }
 42    /// ```
 43    #[tracing::instrument(target = "r1cs", skip(self, other))]
 44    fn bitor(self, other: Self) -> Self::Output {
 45        self._or(other).unwrap()
 46    }
 47}
 48
 49impl<'a, const N: usize, T: PrimUInt, F: PrimeField> BitOr<&'a Self> for UInt<N, T, F> {
 50    type Output = UInt<N, T, F>;
 51
 52    #[tracing::instrument(target = "r1cs", skip(self, other))]
 53    fn bitor(self, other: &Self) -> Self::Output {
 54        self._or(&other).unwrap()
 55    }
 56}
 57
 58impl<'a, const N: usize, T: PrimUInt, F: PrimeField> BitOr<UInt<N, T, F>> for &'a UInt<N, T, F> {
 59    type Output = UInt<N, T, F>;
 60
 61    #[tracing::instrument(target = "r1cs", skip(self, other))]
 62    fn bitor(self, other: UInt<N, T, F>) -> Self::Output {
 63        self._or(&other).unwrap()
 64    }
 65}
 66
 67impl<const N: usize, T: PrimUInt, F: PrimeField> BitOr<Self> for UInt<N, T, F> {
 68    type Output = Self;
 69
 70    #[tracing::instrument(target = "r1cs", skip(self, other))]
 71    fn bitor(self, other: Self) -> Self::Output {
 72        self._or(&other).unwrap()
 73    }
 74}
 75
 76impl<const N: usize, T: PrimUInt, F: PrimeField> BitOrAssign<Self> for UInt<N, T, F> {
 77    /// Sets `self = self | other`.
 78    ///
 79    /// If at least one of `self` and `other` are constants, then this method
 80    /// *does not* create any constraints or variables.
 81    ///
 82    /// ```
 83    /// # fn main() -> Result<(), ark_relations::r1cs::SynthesisError> {
 84    /// // We'll use the BLS12-381 scalar field for our constraints.
 85    /// use ark_test_curves::bls12_381::Fr;
 86    /// use ark_relations::r1cs::*;
 87    /// use ark_r1cs_std::prelude::*;
 88    ///
 89    /// let cs = ConstraintSystem::<Fr>::new_ref();
 90    /// let mut a = UInt8::new_witness(cs.clone(), || Ok(16))?;
 91    /// let b = UInt8::new_witness(cs.clone(), || Ok(17))?;
 92    /// let c = UInt8::new_witness(cs.clone(), || Ok(16 | 17))?;
 93    ///
 94    /// a |= b;
 95    /// a.enforce_equal(&c)?;
 96    /// assert!(cs.is_satisfied().unwrap());
 97    /// # Ok(())
 98    /// # }
 99    /// ```
100    #[tracing::instrument(target = "r1cs", skip(self, other))]
101    fn bitor_assign(&mut self, other: Self) {
102        let result = self._or(&other).unwrap();
103        *self = result;
104    }
105}
106
107impl<'a, const N: usize, T: PrimUInt, F: PrimeField> BitOrAssign<&'a Self> for UInt<N, T, F> {
108    #[tracing::instrument(target = "r1cs", skip(self, other))]
109    fn bitor_assign(&mut self, other: &'a Self) {
110        let result = self._or(other).unwrap();
111        *self = result;
112    }
113}
114
115#[cfg(test)]
116mod tests {
117    use super::*;
118    use crate::{
119        alloc::{AllocVar, AllocationMode},
120        prelude::EqGadget,
121        uint::test_utils::{run_binary_exhaustive, run_binary_random},
122        R1CSVar,
123    };
124    use ark_ff::PrimeField;
125    use ark_test_curves::bls12_381::Fr;
126
127    fn uint_or<T: PrimUInt, const N: usize, F: PrimeField>(
128        a: UInt<N, T, F>,
129        b: UInt<N, T, F>,
130    ) -> Result<(), SynthesisError> {
131        let cs = a.cs().or(b.cs());
132        let both_constant = a.is_constant() && b.is_constant();
133        let computed = &a | &b;
134        let expected_mode = if both_constant {
135            AllocationMode::Constant
136        } else {
137            AllocationMode::Witness
138        };
139        let expected = UInt::<N, T, F>::new_variable(
140            cs.clone(),
141            || Ok(a.value()? | b.value()?),
142            expected_mode,
143        )?;
144        assert_eq!(expected.value(), computed.value());
145        expected.enforce_equal(&computed)?;
146        if !both_constant {
147            assert!(cs.is_satisfied().unwrap());
148        }
149        Ok(())
150    }
151
152    #[test]
153    fn u8_or() {
154        run_binary_exhaustive(uint_or::<u8, 8, Fr>).unwrap()
155    }
156
157    #[test]
158    fn u16_or() {
159        run_binary_random::<1000, 16, _, _>(uint_or::<u16, 16, Fr>).unwrap()
160    }
161
162    #[test]
163    fn u32_or() {
164        run_binary_random::<1000, 32, _, _>(uint_or::<u32, 32, Fr>).unwrap()
165    }
166
167    #[test]
168    fn u64_or() {
169        run_binary_random::<1000, 64, _, _>(uint_or::<u64, 64, Fr>).unwrap()
170    }
171
172    #[test]
173    fn u128_or() {
174        run_binary_random::<1000, 128, _, _>(uint_or::<u128, 128, Fr>).unwrap()
175    }
176}