1use algebra::{Field, FpParameters, PrimeField, ToConstraintField};
  2
  3use r1cs_core::{ConstraintSystem, SynthesisError};
  4
  5use crate::{boolean::AllocatedBit, fields::fp::FpGadget, prelude::*, Assignment, Vec};
  6use core::borrow::Borrow;
  7
  8/// Represents an interpretation of 8 `Boolean` objects as an
  9/// unsigned integer.
 10#[derive(Clone, Debug)]
 11pub struct UInt8 {
 12    // Least significant bit_gadget first
 13    pub(crate) bits: Vec<Boolean>,
 14    pub(crate) value: Option<u8>,
 15}
 16
 17impl UInt8 {
 18    pub fn get_value(&self) -> Option<u8> {
 19        self.value
 20    }
 21
 22    /// Construct a constant vector of `UInt8` from a vector of `u8`
 23    pub fn constant_vec(values: &[u8]) -> Vec<Self> {
 24        let mut result = Vec::new();
 25        for value in values {
 26            result.push(UInt8::constant(*value));
 27        }
 28        result
 29    }
 30
 31    /// Construct a constant `UInt8` from a `u8`
 32    pub fn constant(value: u8) -> Self {
 33        let mut bits = Vec::with_capacity(8);
 34
 35        let mut tmp = value;
 36        for _ in 0..8 {
 37            // If last bit is one, push one.
 38            if tmp & 1 == 1 {
 39                bits.push(Boolean::constant(true))
 40            } else {
 41                bits.push(Boolean::constant(false))
 42            }
 43
 44            tmp >>= 1;
 45        }
 46
 47        Self {
 48            bits,
 49            value: Some(value),
 50        }
 51    }
 52
 53    pub fn alloc_vec<ConstraintF, CS, T>(
 54        mut cs: CS,
 55        values: &[T],
 56    ) -> Result<Vec<Self>, SynthesisError>
 57    where
 58        ConstraintF: Field,
 59        CS: ConstraintSystem<ConstraintF>,
 60        T: Into<Option<u8>> + Copy,
 61    {
 62        let mut output_vec = Vec::with_capacity(values.len());
 63        for (i, value) in values.iter().enumerate() {
 64            let byte: Option<u8> = Into::into(*value);
 65            let alloc_byte = Self::alloc(&mut cs.ns(|| format!("byte_{}", i)), || byte.get())?;
 66            output_vec.push(alloc_byte);
 67        }
 68        Ok(output_vec)
 69    }
 70
 71    /// Allocates a vector of `u8`'s by first converting (chunks of) them to
 72    /// `ConstraintF` elements, (thus reducing the number of input allocations),
 73    /// and then converts this list of `ConstraintF` gadgets back into
 74    /// bytes.
 75    pub fn alloc_input_vec<ConstraintF, CS>(
 76        mut cs: CS,
 77        values: &[u8],
 78    ) -> Result<Vec<Self>, SynthesisError>
 79    where
 80        ConstraintF: PrimeField,
 81        CS: ConstraintSystem<ConstraintF>,
 82    {
 83        let values_len = values.len();
 84        let field_elements: Vec<ConstraintF> =
 85            ToConstraintField::<ConstraintF>::to_field_elements(values).unwrap();
 86
 87        let max_size = 8 * (ConstraintF::Params::CAPACITY / 8) as usize;
 88        let mut allocated_bits = Vec::new();
 89        for (i, field_element) in field_elements.into_iter().enumerate() {
 90            let fe = FpGadget::alloc_input(&mut cs.ns(|| format!("Field element {}", i)), || {
 91                Ok(field_element)
 92            })?;
 93            let mut fe_bits = fe.to_bits(cs.ns(|| format!("Convert fe to bits {}", i)))?;
 94            // FpGadget::to_bits outputs a big-endian binary representation of
 95            // fe_gadget's value, so we have to reverse it to get the little-endian
 96            // form.
 97            fe_bits.reverse();
 98
 99            // Remove the most significant bit, because we know it should be zero
100            // because `values.to_field_elements()` only
101            // packs field elements up to the penultimate bit.
102            // That is, the most significant bit (`ConstraintF::NUM_BITS`-th bit) is
103            // unset, so we can just pop it off.
104            allocated_bits.extend_from_slice(&fe_bits[0..max_size]);
105        }
106
107        // Chunk up slices of 8 bit into bytes.
108        Ok(allocated_bits[0..8 * values_len]
109            .chunks(8)
110            .map(Self::from_bits_le)
111            .collect())
112    }
113
114    /// Turns this `UInt8` into its little-endian byte order representation.
115    /// LSB-first means that we can easily get the corresponding field element
116    /// via double and add.
117    pub fn into_bits_le(&self) -> Vec<Boolean> {
118        self.bits.to_vec()
119    }
120
121    /// Converts a little-endian byte order representation of bits into a
122    /// `UInt8`.
123    pub fn from_bits_le(bits: &[Boolean]) -> Self {
124        assert_eq!(bits.len(), 8);
125
126        let bits = bits.to_vec();
127
128        let mut value = Some(0u8);
129        for b in bits.iter().rev() {
130            value.as_mut().map(|v| *v <<= 1);
131
132            match *b {
133                Boolean::Constant(b) => {
134                    if b {
135                        value.as_mut().map(|v| *v |= 1);
136                    }
137                }
138                Boolean::Is(ref b) => match b.get_value() {
139                    Some(true) => {
140                        value.as_mut().map(|v| *v |= 1);
141                    }
142                    Some(false) => {}
143                    None => value = None,
144                },
145                Boolean::Not(ref b) => match b.get_value() {
146                    Some(false) => {
147                        value.as_mut().map(|v| *v |= 1);
148                    }
149                    Some(true) => {}
150                    None => value = None,
151                },
152            }
153        }
154
155        Self { value, bits }
156    }
157
158    /// XOR this `UInt8` with another `UInt8`
159    pub fn xor<ConstraintF, CS>(&self, mut cs: CS, other: &Self) -> Result<Self, SynthesisError>
160    where
161        ConstraintF: Field,
162        CS: ConstraintSystem<ConstraintF>,
163    {
164        let new_value = match (self.value, other.value) {
165            (Some(a), Some(b)) => Some(a ^ b),
166            _ => None,
167        };
168
169        let bits = self
170            .bits
171            .iter()
172            .zip(other.bits.iter())
173            .enumerate()
174            .map(|(i, (a, b))| Boolean::xor(cs.ns(|| format!("xor of bit_gadget {}", i)), a, b))
175            .collect::<Result<_, _>>()?;
176
177        Ok(Self {
178            bits,
179            value: new_value,
180        })
181    }
182}
183
184impl PartialEq for UInt8 {
185    fn eq(&self, other: &Self) -> bool {
186        self.value.is_some() && other.value.is_some() && self.value == other.value
187    }
188}
189
190impl Eq for UInt8 {}
191
192impl<ConstraintF: Field> ConditionalEqGadget<ConstraintF> for UInt8 {
193    fn conditional_enforce_equal<CS: ConstraintSystem<ConstraintF>>(
194        &self,
195        mut cs: CS,
196        other: &Self,
197        condition: &Boolean,
198    ) -> Result<(), SynthesisError> {
199        for (i, (a, b)) in self.bits.iter().zip(&other.bits).enumerate() {
200            a.conditional_enforce_equal(
201                &mut cs.ns(|| format!("UInt8 equality check for {}-th bit", i)),
202                b,
203                condition,
204            )?;
205        }
206        Ok(())
207    }
208
209    fn cost() -> usize {
210        8 * <Boolean as ConditionalEqGadget<ConstraintF>>::cost()
211    }
212}
213
214impl<ConstraintF: Field> EqGadget<ConstraintF> for UInt8 {}
215
216impl<ConstraintF: Field> AllocGadget<u8, ConstraintF> for UInt8 {
217    fn alloc_constant<T, CS: ConstraintSystem<ConstraintF>>(
218        _cs: CS,
219        t: T,
220    ) -> Result<Self, SynthesisError>
221    where
222        T: Borrow<u8>,
223    {
224        Ok(UInt8::constant(*t.borrow()))
225    }
226
227    fn alloc<F, T, CS: ConstraintSystem<ConstraintF>>(
228        mut cs: CS,
229        value_gen: F,
230    ) -> Result<Self, SynthesisError>
231    where
232        F: FnOnce() -> Result<T, SynthesisError>,
233        T: Borrow<u8>,
234    {
235        let value = value_gen().map(|val| *val.borrow());
236        let values = match value {
237            Ok(mut val) => {
238                let mut v = Vec::with_capacity(8);
239
240                for _ in 0..8 {
241                    v.push(Some(val & 1 == 1));
242                    val >>= 1;
243                }
244
245                v
246            }
247            _ => vec![None; 8],
248        };
249
250        let bits = values
251            .into_iter()
252            .enumerate()
253            .map(|(i, v)| {
254                Ok(Boolean::from(AllocatedBit::alloc(
255                    &mut cs.ns(|| format!("allocated bit_gadget {}", i)),
256                    || v.ok_or(SynthesisError::AssignmentMissing),
257                )?))
258            })
259            .collect::<Result<Vec<_>, SynthesisError>>()?;
260
261        Ok(Self {
262            bits,
263            value: value.ok(),
264        })
265    }
266
267    fn alloc_input<F, T, CS: ConstraintSystem<ConstraintF>>(
268        mut cs: CS,
269        value_gen: F,
270    ) -> Result<Self, SynthesisError>
271    where
272        F: FnOnce() -> Result<T, SynthesisError>,
273        T: Borrow<u8>,
274    {
275        let value = value_gen().map(|val| *val.borrow());
276        let values = match value {
277            Ok(mut val) => {
278                let mut v = Vec::with_capacity(8);
279                for _ in 0..8 {
280                    v.push(Some(val & 1 == 1));
281                    val >>= 1;
282                }
283
284                v
285            }
286            _ => vec![None; 8],
287        };
288
289        let bits = values
290            .into_iter()
291            .enumerate()
292            .map(|(i, v)| {
293                Ok(Boolean::from(AllocatedBit::alloc_input(
294                    &mut cs.ns(|| format!("allocated bit_gadget {}", i)),
295                    || v.ok_or(SynthesisError::AssignmentMissing),
296                )?))
297            })
298            .collect::<Result<Vec<_>, SynthesisError>>()?;
299
300        Ok(Self {
301            bits,
302            value: value.ok(),
303        })
304    }
305}
306
307#[cfg(test)]
308mod test {
309    use super::UInt8;
310    use crate::{prelude::*, test_constraint_system::TestConstraintSystem, Vec};
311    use algebra::bls12_381::Fr;
312    use r1cs_core::ConstraintSystem;
313    use rand::{Rng, SeedableRng};
314    use rand_xorshift::XorShiftRng;
315
316    #[test]
317    fn test_uint8_from_bits_to_bits() {
318        let mut cs = TestConstraintSystem::<Fr>::new();
319        let byte_val = 0b01110001;
320        let byte = UInt8::alloc(cs.ns(|| "alloc value"), || Ok(byte_val)).unwrap();
321        let bits = byte.into_bits_le();
322        for (i, bit) in bits.iter().enumerate() {
323            assert_eq!(bit.get_value().unwrap(), (byte_val >> i) & 1 == 1)
324        }
325    }
326
327    #[test]
328    fn test_uint8_alloc_input_vec() {
329        let mut cs = TestConstraintSystem::<Fr>::new();
330        let byte_vals = (64u8..128u8).collect::<Vec<_>>();
331        let bytes = UInt8::alloc_input_vec(cs.ns(|| "alloc value"), &byte_vals).unwrap();
332        for (native_byte, gadget_byte) in byte_vals.into_iter().zip(bytes) {
333            let bits = gadget_byte.into_bits_le();
334            for (i, bit) in bits.iter().enumerate() {
335                assert_eq!(bit.get_value().unwrap(), (native_byte >> i) & 1 == 1)
336            }
337        }
338    }
339
340    #[test]
341    fn test_uint8_from_bits() {
342        let mut rng = XorShiftRng::seed_from_u64(1231275789u64);
343
344        for _ in 0..1000 {
345            let v = (0..8)
346                .map(|_| Boolean::constant(rng.gen()))
347                .collect::<Vec<_>>();
348
349            let b = UInt8::from_bits_le(&v);
350
351            for (i, bit_gadget) in b.bits.iter().enumerate() {
352                match bit_gadget {
353                    &Boolean::Constant(bit_gadget) => {
354                        assert!(bit_gadget == ((b.value.unwrap() >> i) & 1 == 1));
355                    }
356                    _ => unreachable!(),
357                }
358            }
359
360            let expected_to_be_same = b.into_bits_le();
361
362            for x in v.iter().zip(expected_to_be_same.iter()) {
363                match x {
364                    (&Boolean::Constant(true), &Boolean::Constant(true)) => {}
365                    (&Boolean::Constant(false), &Boolean::Constant(false)) => {}
366                    _ => unreachable!(),
367                }
368            }
369        }
370    }
371
372    #[test]
373    fn test_uint8_xor() {
374        let mut rng = XorShiftRng::seed_from_u64(1231275789u64);
375
376        for _ in 0..1000 {
377            let mut cs = TestConstraintSystem::<Fr>::new();
378
379            let a: u8 = rng.gen();
380            let b: u8 = rng.gen();
381            let c: u8 = rng.gen();
382
383            let mut expected = a ^ b ^ c;
384
385            let a_bit = UInt8::alloc(cs.ns(|| "a_bit"), || Ok(a)).unwrap();
386            let b_bit = UInt8::constant(b);
387            let c_bit = UInt8::alloc(cs.ns(|| "c_bit"), || Ok(c)).unwrap();
388
389            let r = a_bit.xor(cs.ns(|| "first xor"), &b_bit).unwrap();
390            let r = r.xor(cs.ns(|| "second xor"), &c_bit).unwrap();
391
392            assert!(cs.is_satisfied());
393
394            assert!(r.value == Some(expected));
395
396            for b in r.bits.iter() {
397                match b {
398                    &Boolean::Is(ref b) => {
399                        assert!(b.get_value().unwrap() == (expected & 1 == 1));
400                    }
401                    &Boolean::Not(ref b) => {
402                        assert!(!b.get_value().unwrap() == (expected & 1 == 1));
403                    }
404                    &Boolean::Constant(b) => {
405                        assert!(b == (expected & 1 == 1));
406                    }
407                }
408
409                expected >>= 1;
410            }
411        }
412    }
413}