1use ark_ff::{Field, PrimeField};
  2use core::{borrow::Borrow, convert::TryFrom, fmt::Debug};
  3
  4use ark_relations::r1cs::{ConstraintSystemRef, Namespace, SynthesisError};
  5
  6use crate::{boolean::Boolean, prelude::*, Assignment, Vec};
  7
  8mod add;
  9mod and;
 10mod cmp;
 11mod convert;
 12mod eq;
 13mod not;
 14mod or;
 15mod rotate;
 16mod select;
 17mod shl;
 18mod shr;
 19mod xor;
 20
 21#[doc(hidden)]
 22pub mod prim_uint;
 23pub use prim_uint::*;
 24
 25#[cfg(test)]
 26pub(crate) mod test_utils;
 27
 28/// This struct represent an unsigned `N` bit integer as a sequence of `N` [`Boolean`]s.
 29#[derive(Clone, Debug)]
 30pub struct UInt<const N: usize, T: PrimUInt, F: Field> {
 31    #[doc(hidden)]
 32    pub bits: [Boolean<F>; N],
 33    #[doc(hidden)]
 34    pub value: Option<T>,
 35}
 36
 37impl<const N: usize, T: PrimUInt, F: Field> R1CSVar<F> for UInt<N, T, F> {
 38    type Value = T;
 39
 40    fn cs(&self) -> ConstraintSystemRef<F> {
 41        self.bits.as_ref().cs()
 42    }
 43
 44    fn value(&self) -> Result<Self::Value, SynthesisError> {
 45        let mut value = T::zero();
 46        for (i, bit) in self.bits.iter().enumerate() {
 47            value = value + (T::from(bit.value()? as u8).unwrap() << i);
 48        }
 49        debug_assert_eq!(self.value, Some(value));
 50        Ok(value)
 51    }
 52}
 53
 54impl<const N: usize, T: PrimUInt, F: Field> UInt<N, T, F> {
 55    pub const MAX: Self = Self {
 56        bits: [Boolean::TRUE; N],
 57        value: Some(T::MAX),
 58    };
 59
 60    /// Construct a constant [`UInt`] from the native unsigned integer type.
 61    ///
 62    /// This *does not* create new variables or constraints.
 63    ///
 64    /// ```
 65    /// # fn main() -> Result<(), ark_relations::r1cs::SynthesisError> {
 66    /// // We'll use the BLS12-381 scalar field for our constraints.
 67    /// use ark_test_curves::bls12_381::Fr;
 68    /// use ark_relations::r1cs::*;
 69    /// use ark_r1cs_std::prelude::*;
 70    ///
 71    /// let cs = ConstraintSystem::<Fr>::new_ref();
 72    /// let var = UInt8::new_witness(cs.clone(), || Ok(2))?;
 73    ///
 74    /// let constant = UInt8::constant(2);
 75    /// var.enforce_equal(&constant)?;
 76    /// assert!(cs.is_satisfied().unwrap());
 77    /// # Ok(())
 78    /// # }
 79    /// ```
 80    pub fn constant(value: T) -> Self {
 81        let mut bits = [Boolean::FALSE; N];
 82        let mut bit_values = value;
 83
 84        for i in 0..N {
 85            bits[i] = Boolean::constant((bit_values & T::one()) == T::one());
 86            bit_values = bit_values >> 1u8;
 87        }
 88
 89        Self {
 90            bits,
 91            value: Some(value),
 92        }
 93    }
 94
 95    /// Construct a constant vector of [`UInt`] from a vector of the native type
 96    ///
 97    /// This *does not* create any new variables or constraints.
 98    /// ```
 99    /// # fn main() -> Result<(), ark_relations::r1cs::SynthesisError> {
100    /// // We'll use the BLS12-381 scalar field for our constraints.
101    /// use ark_test_curves::bls12_381::Fr;
102    /// use ark_relations::r1cs::*;
103    /// use ark_r1cs_std::prelude::*;
104    ///
105    /// let cs = ConstraintSystem::<Fr>::new_ref();
106    /// let var = vec![UInt8::new_witness(cs.clone(), || Ok(2))?];
107    ///
108    /// let constant = UInt8::constant_vec(&[2]);
109    /// var.enforce_equal(&constant)?;
110    /// assert!(cs.is_satisfied().unwrap());
111    /// # Ok(())
112    /// # }
113    /// ```
114    pub fn constant_vec(values: &[T]) -> Vec<Self> {
115        values.iter().map(|v| Self::constant(*v)).collect()
116    }
117
118    /// Allocates a slice of `uN`'s as private witnesses.
119    pub fn new_witness_vec(
120        cs: impl Into<Namespace<F>>,
121        values: &[impl Into<Option<T>> + Copy],
122    ) -> Result<Vec<Self>, SynthesisError> {
123        let ns = cs.into();
124        let cs = ns.cs();
125        let mut output_vec = Vec::with_capacity(values.len());
126        for value in values {
127            let byte: Option<T> = Into::into(*value);
128            output_vec.push(Self::new_witness(cs.clone(), || byte.get())?);
129        }
130        Ok(output_vec)
131    }
132}
133
134impl<const N: usize, T: PrimUInt, ConstraintF: Field> AllocVar<T, ConstraintF>
135    for UInt<N, T, ConstraintF>
136{
137    fn new_variable<S: Borrow<T>>(
138        cs: impl Into<Namespace<ConstraintF>>,
139        f: impl FnOnce() -> Result<S, SynthesisError>,
140        mode: AllocationMode,
141    ) -> Result<Self, SynthesisError> {
142        let ns = cs.into();
143        let cs = ns.cs();
144        let value = f().map(|f| *f.borrow()).ok();
145
146        let mut values = [None; N];
147        if let Some(val) = value {
148            values
149                .iter_mut()
150                .enumerate()
151                .for_each(|(i, v)| *v = Some(((val >> i) & T::one()) == T::one()));
152        }
153
154        let mut bits = [Boolean::FALSE; N];
155        for (b, v) in bits.iter_mut().zip(&values) {
156            *b = Boolean::new_variable(cs.clone(), || v.get(), mode)?;
157        }
158        Ok(Self { bits, value })
159    }
160}