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