1use algebra::{Field, FpParameters, PrimeField, ToConstraintField};
  2
  3use r1cs_core::{ConstraintSystem, SynthesisError};
  4
  5use crate::{boolean::AllocatedBit, fields::fp::FpGadget, prelude::*, Assignment};
  6use std::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<F, T, CS: ConstraintSystem<ConstraintF>>(
218        mut cs: CS,
219        value_gen: F,
220    ) -> Result<Self, SynthesisError>
221    where
222        F: FnOnce() -> Result<T, SynthesisError>,
223        T: Borrow<u8>,
224    {
225        let value = value_gen().map(|val| *val.borrow());
226        let values = match value {
227            Ok(mut val) => {
228                let mut v = Vec::with_capacity(8);
229
230                for _ in 0..8 {
231                    v.push(Some(val & 1 == 1));
232                    val >>= 1;
233                }
234
235                v
236            },
237            _ => vec![None; 8],
238        };
239
240        let bits = values
241            .into_iter()
242            .enumerate()
243            .map(|(i, v)| {
244                Ok(Boolean::from(AllocatedBit::alloc(
245                    &mut cs.ns(|| format!("allocated bit_gadget {}", i)),
246                    || v.ok_or(SynthesisError::AssignmentMissing),
247                )?))
248            })
249            .collect::<Result<Vec<_>, SynthesisError>>()?;
250
251        Ok(Self {
252            bits,
253            value: value.ok(),
254        })
255    }
256
257    fn alloc_input<F, T, CS: ConstraintSystem<ConstraintF>>(
258        mut cs: CS,
259        value_gen: F,
260    ) -> Result<Self, SynthesisError>
261    where
262        F: FnOnce() -> Result<T, SynthesisError>,
263        T: Borrow<u8>,
264    {
265        let value = value_gen().map(|val| *val.borrow());
266        let values = match value {
267            Ok(mut val) => {
268                let mut v = Vec::with_capacity(8);
269                for _ in 0..8 {
270                    v.push(Some(val & 1 == 1));
271                    val >>= 1;
272                }
273
274                v
275            },
276            _ => vec![None; 8],
277        };
278
279        let bits = values
280            .into_iter()
281            .enumerate()
282            .map(|(i, v)| {
283                Ok(Boolean::from(AllocatedBit::alloc_input(
284                    &mut cs.ns(|| format!("allocated bit_gadget {}", i)),
285                    || v.ok_or(SynthesisError::AssignmentMissing),
286                )?))
287            })
288            .collect::<Result<Vec<_>, SynthesisError>>()?;
289
290        Ok(Self {
291            bits,
292            value: value.ok(),
293        })
294    }
295}
296
297#[cfg(test)]
298mod test {
299    use super::UInt8;
300    use crate::{prelude::*, test_constraint_system::TestConstraintSystem};
301    use algebra::fields::bls12_381::Fr;
302    use r1cs_core::ConstraintSystem;
303    use rand::{Rng, SeedableRng};
304    use rand_xorshift::XorShiftRng;
305
306    #[test]
307    fn test_uint8_from_bits_to_bits() {
308        let mut cs = TestConstraintSystem::<Fr>::new();
309        let byte_val = 0b01110001;
310        let byte = UInt8::alloc(cs.ns(|| "alloc value"), || Ok(byte_val)).unwrap();
311        let bits = byte.into_bits_le();
312        for (i, bit) in bits.iter().enumerate() {
313            assert_eq!(bit.get_value().unwrap(), (byte_val >> i) & 1 == 1)
314        }
315    }
316
317    #[test]
318    fn test_uint8_alloc_input_vec() {
319        let mut cs = TestConstraintSystem::<Fr>::new();
320        let byte_vals = (64u8..128u8).collect::<Vec<_>>();
321        let bytes = UInt8::alloc_input_vec(cs.ns(|| "alloc value"), &byte_vals).unwrap();
322        for (native_byte, gadget_byte) in byte_vals.into_iter().zip(bytes) {
323            let bits = gadget_byte.into_bits_le();
324            for (i, bit) in bits.iter().enumerate() {
325                assert_eq!(bit.get_value().unwrap(), (native_byte >> i) & 1 == 1)
326            }
327        }
328    }
329
330    #[test]
331    fn test_uint8_from_bits() {
332        let mut rng = XorShiftRng::seed_from_u64(1231275789u64);
333
334        for _ in 0..1000 {
335            let v = (0..8)
336                .map(|_| Boolean::constant(rng.gen()))
337                .collect::<Vec<_>>();
338
339            let b = UInt8::from_bits_le(&v);
340
341            for (i, bit_gadget) in b.bits.iter().enumerate() {
342                match bit_gadget {
343                    &Boolean::Constant(bit_gadget) => {
344                        assert!(bit_gadget == ((b.value.unwrap() >> i) & 1 == 1));
345                    },
346                    _ => unreachable!(),
347                }
348            }
349
350            let expected_to_be_same = b.into_bits_le();
351
352            for x in v.iter().zip(expected_to_be_same.iter()) {
353                match x {
354                    (&Boolean::Constant(true), &Boolean::Constant(true)) => {},
355                    (&Boolean::Constant(false), &Boolean::Constant(false)) => {},
356                    _ => unreachable!(),
357                }
358            }
359        }
360    }
361
362    #[test]
363    fn test_uint8_xor() {
364        let mut rng = XorShiftRng::seed_from_u64(1231275789u64);
365
366        for _ in 0..1000 {
367            let mut cs = TestConstraintSystem::<Fr>::new();
368
369            let a: u8 = rng.gen();
370            let b: u8 = rng.gen();
371            let c: u8 = rng.gen();
372
373            let mut expected = a ^ b ^ c;
374
375            let a_bit = UInt8::alloc(cs.ns(|| "a_bit"), || Ok(a)).unwrap();
376            let b_bit = UInt8::constant(b);
377            let c_bit = UInt8::alloc(cs.ns(|| "c_bit"), || Ok(c)).unwrap();
378
379            let r = a_bit.xor(cs.ns(|| "first xor"), &b_bit).unwrap();
380            let r = r.xor(cs.ns(|| "second xor"), &c_bit).unwrap();
381
382            assert!(cs.is_satisfied());
383
384            assert!(r.value == Some(expected));
385
386            for b in r.bits.iter() {
387                match b {
388                    &Boolean::Is(ref b) => {
389                        assert!(b.get_value().unwrap() == (expected & 1 == 1));
390                    },
391                    &Boolean::Not(ref b) => {
392                        assert!(!b.get_value().unwrap() == (expected & 1 == 1));
393                    },
394                    &Boolean::Constant(b) => {
395                        assert!(b == (expected & 1 == 1));
396                    },
397                }
398
399                expected >>= 1;
400            }
401        }
402    }
403}