1use algebra::Field;
  2use algebra::{FpParameters, PrimeField, ToConstraintField};
  3
  4use r1cs_core::{ConstraintSystemRef, Namespace, SynthesisError};
  5
  6use crate::{fields::fp::AllocatedFp, prelude::*, Assignment, Vec};
  7use core::borrow::Borrow;
  8
  9/// Represents an interpretation of 8 `Boolean` objects as an
 10/// unsigned integer.
 11#[derive(Clone, Debug)]
 12pub struct UInt8<F: Field> {
 13    /// Little-endian representation: least significant bit first
 14    pub(crate) bits: Vec<Boolean<F>>,
 15    /// Little-endian representation: least significant bit first
 16    pub(crate) value: Option<u8>,
 17}
 18
 19impl<F: Field> R1CSVar<F> for UInt8<F> {
 20    type Value = u8;
 21
 22    fn cs(&self) -> Option<ConstraintSystemRef<F>> {
 23        self.bits.as_slice().cs()
 24    }
 25
 26    fn value(&self) -> Result<Self::Value, SynthesisError> {
 27        let mut value = None;
 28        for (i, bit) in self.bits.iter().enumerate() {
 29            let b = u8::from(bit.value()?);
 30            value = match value {
 31                Some(value) => Some(value + (b << i)),
 32                None => Some(b << i),
 33            };
 34        }
 35        debug_assert_eq!(self.value, value);
 36        value.get()
 37    }
 38}
 39
 40impl<F: Field> UInt8<F> {
 41    /// Construct a constant vector of `UInt8` from a vector of `u8`
 42    pub fn constant_vec(values: &[u8]) -> Vec<Self> {
 43        let mut result = Vec::new();
 44        for value in values {
 45            result.push(UInt8::constant(*value));
 46        }
 47        result
 48    }
 49
 50    /// Construct a constant `UInt8` from a `u8`
 51    pub fn constant(value: u8) -> Self {
 52        let mut bits = Vec::with_capacity(8);
 53
 54        let mut tmp = value;
 55        for _ in 0..8 {
 56            // If last bit is one, push one.
 57            bits.push(Boolean::constant(tmp & 1 == 1));
 58            tmp >>= 1;
 59        }
 60
 61        Self {
 62            bits,
 63            value: Some(value),
 64        }
 65    }
 66
 67    pub fn new_witness_vec(
 68        cs: impl Into<Namespace<F>>,
 69        values: &[impl Into<Option<u8>> + Copy],
 70    ) -> Result<Vec<Self>, SynthesisError> {
 71        let ns = cs.into();
 72        let cs = ns.cs();
 73        let mut output_vec = Vec::with_capacity(values.len());
 74        for value in values {
 75            let byte: Option<u8> = Into::into(*value);
 76            output_vec.push(Self::new_witness(cs.clone(), || byte.get())?);
 77        }
 78        Ok(output_vec)
 79    }
 80
 81    /// Allocates a vector of `u8`'s by first converting (chunks of) them to
 82    /// `ConstraintF` elements, (thus reducing the number of input allocations),
 83    /// and then converts this list of `ConstraintF` gadgets back into
 84    /// bytes.
 85    pub fn new_input_vec(
 86        cs: impl Into<Namespace<F>>,
 87        values: &[u8],
 88    ) -> Result<Vec<Self>, SynthesisError>
 89    where
 90        F: PrimeField,
 91    {
 92        let ns = cs.into();
 93        let cs = ns.cs();
 94        let values_len = values.len();
 95        let field_elements: Vec<F> = ToConstraintField::<F>::to_field_elements(values).unwrap();
 96
 97        let max_size = 8 * (F::Params::CAPACITY / 8) as usize;
 98        let mut allocated_bits = Vec::new();
 99        for field_element in field_elements.into_iter() {
100            let fe = AllocatedFp::new_input(cs.clone(), || Ok(field_element))?;
101            let fe_bits = fe.to_bits_le()?;
102
103            // Remove the most significant bit, because we know it should be zero
104            // because `values.to_field_elements()` only
105            // packs field elements up to the penultimate bit.
106            // That is, the most significant bit (`ConstraintF::NUM_BITS`-th bit) is
107            // unset, so we can just pop it off.
108            allocated_bits.extend_from_slice(&fe_bits[0..max_size]);
109        }
110
111        // Chunk up slices of 8 bit into bytes.
112        Ok(allocated_bits[0..(8 * values_len)]
113            .chunks(8)
114            .map(Self::from_bits_le)
115            .collect())
116    }
117
118    /// Converts a little-endian byte order representation of bits into a
119    /// `UInt8`.
120    #[tracing::instrument(target = "r1cs")]
121    pub fn from_bits_le(bits: &[Boolean<F>]) -> Self {
122        assert_eq!(bits.len(), 8);
123
124        let bits = bits.to_vec();
125
126        let mut value = Some(0u8);
127        for (i, b) in bits.iter().enumerate() {
128            value = match b.value().ok() {
129                Some(b) => value.map(|v| v + (u8::from(b) << i)),
130                None => None,
131            }
132        }
133
134        Self { value, bits }
135    }
136
137    /// XOR this `UInt8` with another `UInt8`
138    #[tracing::instrument(target = "r1cs")]
139    pub fn xor(&self, other: &Self) -> Result<Self, SynthesisError> {
140        let new_value = match (self.value, other.value) {
141            (Some(a), Some(b)) => Some(a ^ b),
142            _ => None,
143        };
144
145        let bits = self
146            .bits
147            .iter()
148            .zip(other.bits.iter())
149            .map(|(a, b)| a.xor(b))
150            .collect::<Result<_, _>>()?;
151
152        Ok(Self {
153            bits,
154            value: new_value,
155        })
156    }
157}
158
159impl<ConstraintF: Field> EqGadget<ConstraintF> for UInt8<ConstraintF> {
160    #[tracing::instrument(target = "r1cs")]
161    fn is_eq(&self, other: &Self) -> Result<Boolean<ConstraintF>, SynthesisError> {
162        self.bits.as_slice().is_eq(&other.bits)
163    }
164
165    #[tracing::instrument(target = "r1cs")]
166    fn conditional_enforce_equal(
167        &self,
168        other: &Self,
169        condition: &Boolean<ConstraintF>,
170    ) -> Result<(), SynthesisError> {
171        self.bits.conditional_enforce_equal(&other.bits, condition)
172    }
173
174    #[tracing::instrument(target = "r1cs")]
175    fn conditional_enforce_not_equal(
176        &self,
177        other: &Self,
178        condition: &Boolean<ConstraintF>,
179    ) -> Result<(), SynthesisError> {
180        self.bits
181            .conditional_enforce_not_equal(&other.bits, condition)
182    }
183}
184
185impl<ConstraintF: Field> AllocVar<u8, ConstraintF> for UInt8<ConstraintF> {
186    fn new_variable<T: Borrow<u8>>(
187        cs: impl Into<Namespace<ConstraintF>>,
188        f: impl FnOnce() -> Result<T, SynthesisError>,
189        mode: AllocationMode,
190    ) -> Result<Self, SynthesisError> {
191        let ns = cs.into();
192        let cs = ns.cs();
193        let value = f().map(|f| *f.borrow());
194        let values = match value {
195            Ok(val) => (0..8).map(|i| Some((val >> i) & 1 == 1)).collect(),
196            _ => vec![None; 8],
197        };
198        let bits = values
199            .into_iter()
200            .map(|v| Boolean::new_variable(cs.clone(), || v.get(), mode))
201            .collect::<Result<Vec<_>, _>>()?;
202        Ok(Self {
203            bits,
204            value: value.ok(),
205        })
206    }
207}
208
209#[cfg(test)]
210mod test {
211    use super::UInt8;
212    use crate::{prelude::*, Vec};
213    use algebra::bls12_381::Fr;
214    use r1cs_core::{ConstraintSystem, SynthesisError};
215    use rand::{Rng, SeedableRng};
216    use rand_xorshift::XorShiftRng;
217
218    #[test]
219    fn test_uint8_from_bits_to_bits() -> Result<(), SynthesisError> {
220        let cs = ConstraintSystem::<Fr>::new_ref();
221        let byte_val = 0b01110001;
222        let byte = UInt8::new_witness(r1cs_core::ns!(cs, "alloc value"), || Ok(byte_val)).unwrap();
223        let bits = byte.to_bits_le()?;
224        for (i, bit) in bits.iter().enumerate() {
225            assert_eq!(bit.value()?, (byte_val >> i) & 1 == 1)
226        }
227        Ok(())
228    }
229
230    #[test]
231    fn test_uint8_new_input_vec() -> Result<(), SynthesisError> {
232        let cs = ConstraintSystem::<Fr>::new_ref();
233        let byte_vals = (64u8..128u8).collect::<Vec<_>>();
234        let bytes = UInt8::new_input_vec(r1cs_core::ns!(cs, "alloc value"), &byte_vals).unwrap();
235        dbg!(bytes.value())?;
236        for (native, variable) in byte_vals.into_iter().zip(bytes) {
237            let bits = variable.to_bits_le()?;
238            for (i, bit) in bits.iter().enumerate() {
239                assert_eq!(
240                    bit.value()?,
241                    (native >> i) & 1 == 1,
242                    "native value {}: bit {:?}",
243                    native,
244                    i
245                )
246            }
247        }
248        Ok(())
249    }
250
251    #[test]
252    fn test_uint8_from_bits() -> Result<(), SynthesisError> {
253        let mut rng = XorShiftRng::seed_from_u64(1231275789u64);
254
255        for _ in 0..1000 {
256            let v = (0..8)
257                .map(|_| Boolean::<Fr>::Constant(rng.gen()))
258                .collect::<Vec<_>>();
259
260            let val = UInt8::from_bits_le(&v);
261
262            for (i, bit) in val.bits.iter().enumerate() {
263                match bit {
264                    Boolean::Constant(b) => assert!(*b == ((val.value()? >> i) & 1 == 1)),
265                    _ => unreachable!(),
266                }
267            }
268
269            let expected_to_be_same = val.to_bits_le()?;
270
271            for x in v.iter().zip(expected_to_be_same.iter()) {
272                match x {
273                    (&Boolean::Constant(true), &Boolean::Constant(true)) => {}
274                    (&Boolean::Constant(false), &Boolean::Constant(false)) => {}
275                    _ => unreachable!(),
276                }
277            }
278        }
279        Ok(())
280    }
281
282    #[test]
283    fn test_uint8_xor() -> Result<(), SynthesisError> {
284        let mut rng = XorShiftRng::seed_from_u64(1231275789u64);
285
286        for _ in 0..1000 {
287            let cs = ConstraintSystem::<Fr>::new_ref();
288
289            let a: u8 = rng.gen();
290            let b: u8 = rng.gen();
291            let c: u8 = rng.gen();
292
293            let mut expected = a ^ b ^ c;
294
295            let a_bit = UInt8::new_witness(r1cs_core::ns!(cs, "a_bit"), || Ok(a)).unwrap();
296            let b_bit = UInt8::constant(b);
297            let c_bit = UInt8::new_witness(r1cs_core::ns!(cs, "c_bit"), || Ok(c)).unwrap();
298
299            let r = a_bit.xor(&b_bit).unwrap();
300            let r = r.xor(&c_bit).unwrap();
301
302            assert!(cs.is_satisfied().unwrap());
303
304            assert!(r.value == Some(expected));
305
306            for b in r.bits.iter() {
307                match b {
308                    Boolean::Is(b) => assert!(b.value()? == (expected & 1 == 1)),
309                    Boolean::Not(b) => assert!(!b.value()? == (expected & 1 == 1)),
310                    Boolean::Constant(b) => assert!(*b == (expected & 1 == 1)),
311                }
312
313                expected >>= 1;
314            }
315        }
316        Ok(())
317    }
318}