1use algebra::{FpParameters, PrimeField, Field};
  2
  3use r1cs_core::{ConstraintSystem, LinearCombination, SynthesisError};
  4
  5use crate::boolean::{AllocatedBit, Boolean};
  6use crate::Assignment;
  7use crate::prelude::*;
  8
  9/// Represents an interpretation of 32 `Boolean` objects as an
 10/// unsigned integer.
 11#[derive(Clone, Debug)]
 12pub struct UInt32 {
 13    // Least significant bit_gadget first
 14    bits:  Vec<Boolean>,
 15    value: Option<u32>,
 16}
 17
 18impl UInt32 {
 19    /// Construct a constant `UInt32` from a `u32`
 20    pub fn constant(value: u32) -> Self {
 21        let mut bits = Vec::with_capacity(32);
 22
 23        let mut tmp = value;
 24        for _ in 0..32 {
 25            if tmp & 1 == 1 {
 26                bits.push(Boolean::constant(true))
 27            } else {
 28                bits.push(Boolean::constant(false))
 29            }
 30
 31            tmp >>= 1;
 32        }
 33
 34        UInt32 {
 35            bits,
 36            value: Some(value),
 37        }
 38    }
 39
 40    /// Allocate a `UInt32` in the constraint system
 41    pub fn alloc<ConstraintF, CS>(mut cs: CS, value: Option<u32>) -> Result<Self, SynthesisError>
 42    where
 43        ConstraintF: Field,
 44        CS: ConstraintSystem<ConstraintF>,
 45    {
 46        let values = match value {
 47            Some(mut val) => {
 48                let mut v = Vec::with_capacity(32);
 49
 50                for _ in 0..32 {
 51                    v.push(Some(val & 1 == 1));
 52                    val >>= 1;
 53                }
 54
 55                v
 56            },
 57            None => vec![None; 32],
 58        };
 59
 60        let bits = values
 61            .into_iter()
 62            .enumerate()
 63            .map(|(i, v)| {
 64                Ok(Boolean::from(AllocatedBit::alloc(
 65                    cs.ns(|| format!("allocated bit_gadget {}", i)),
 66                    || v.get(),
 67                )?))
 68            })
 69            .collect::<Result<Vec<_>, SynthesisError>>()?;
 70
 71        Ok(UInt32 { bits, value })
 72    }
 73
 74    /// Turns this `UInt32` into its little-endian byte order representation.
 75    pub fn to_bits_le(&self) -> Vec<Boolean> {
 76        self.bits.clone()
 77    }
 78
 79    /// Converts a little-endian byte order representation of bits into a
 80    /// `UInt32`.
 81    pub fn from_bits_le(bits: &[Boolean]) -> Self {
 82        assert_eq!(bits.len(), 32);
 83
 84        let bits = bits.to_vec();
 85
 86        let mut value = Some(0u32);
 87        for b in bits.iter().rev() {
 88            value.as_mut().map(|v| *v <<= 1);
 89
 90            match b {
 91                &Boolean::Constant(b) => {
 92                    if b {
 93                        value.as_mut().map(|v| *v |= 1);
 94                    }
 95                },
 96                &Boolean::Is(ref b) => match b.get_value() {
 97                    Some(true) => {
 98                        value.as_mut().map(|v| *v |= 1);
 99                    },
100                    Some(false) => {},
101                    None => value = None,
102                },
103                &Boolean::Not(ref b) => match b.get_value() {
104                    Some(false) => {
105                        value.as_mut().map(|v| *v |= 1);
106                    },
107                    Some(true) => {},
108                    None => value = None,
109                },
110            }
111        }
112
113        Self { value, bits }
114    }
115
116    pub fn rotr(&self, by: usize) -> Self {
117        let by = by % 32;
118
119        let new_bits = self
120            .bits
121            .iter()
122            .skip(by)
123            .chain(self.bits.iter())
124            .take(32)
125            .cloned()
126            .collect();
127
128        UInt32 {
129            bits:  new_bits,
130            value: self.value.map(|v| v.rotate_right(by as u32)),
131        }
132    }
133
134    /// XOR this `UInt32` with another `UInt32`
135    pub fn xor<ConstraintF, CS>(&self, mut cs: CS, other: &Self) -> Result<Self, SynthesisError>
136    where
137        ConstraintF: Field,
138        CS: ConstraintSystem<ConstraintF>,
139    {
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            .enumerate()
150            .map(|(i, (a, b))| Boolean::xor(cs.ns(|| format!("xor of bit_gadget {}", i)), a, b))
151            .collect::<Result<_, _>>()?;
152
153        Ok(UInt32 {
154            bits,
155            value: new_value,
156        })
157    }
158
159    /// Perform modular addition of several `UInt32` objects.
160    pub fn addmany<ConstraintF, CS>(mut cs: CS, operands: &[Self]) -> Result<Self, SynthesisError>
161    where
162        ConstraintF: PrimeField,
163        CS: ConstraintSystem<ConstraintF>,
164    {
165        // Make some arbitrary bounds for ourselves to avoid overflows
166        // in the scalar field
167        assert!(ConstraintF::Params::MODULUS_BITS >= 64);
168        assert!(operands.len() >= 2); // Weird trivial cases that should never happen
169        assert!(operands.len() <= 10);
170
171        // Compute the maximum value of the sum so we allocate enough bits for
172        // the result
173        let mut max_value = (operands.len() as u64) * u64::from(u32::max_value());
174
175        // Keep track of the resulting value
176        let mut result_value = Some(0u64);
177
178        // This is a linear combination that we will enforce to be "zero"
179        let mut lc = LinearCombination::zero();
180
181        let mut all_constants = true;
182
183        // Iterate over the operands
184        for op in operands {
185            // Accumulate the value
186            match op.value {
187                Some(val) => {
188                    result_value.as_mut().map(|v| *v += u64::from(val));
189                },
190                None => {
191                    // If any of our operands have unknown value, we won't
192                    // know the value of the result
193                    result_value = None;
194                },
195            }
196
197            // Iterate over each bit_gadget of the operand and add the operand to
198            // the linear combination
199            let mut coeff = ConstraintF::one();
200            for bit in &op.bits {
201                match *bit {
202                    Boolean::Is(ref bit) => {
203                        all_constants = false;
204
205                        // Add coeff * bit_gadget
206                        lc = lc + (coeff, bit.get_variable());
207                    },
208                    Boolean::Not(ref bit) => {
209                        all_constants = false;
210
211                        // Add coeff * (1 - bit_gadget) = coeff * ONE - coeff * bit_gadget
212                        lc = lc + (coeff, CS::one()) - (coeff, bit.get_variable());
213                    },
214                    Boolean::Constant(bit) => {
215                        if bit {
216                            lc = lc + (coeff, CS::one());
217                        }
218                    },
219                }
220
221                coeff.double_in_place();
222            }
223        }
224
225        // The value of the actual result is modulo 2^32
226        let modular_value = result_value.map(|v| v as u32);
227
228        if all_constants && modular_value.is_some() {
229            // We can just return a constant, rather than
230            // unpacking the result into allocated bits.
231
232            return Ok(UInt32::constant(modular_value.unwrap()));
233        }
234
235        // Storage area for the resulting bits
236        let mut result_bits = vec![];
237
238        // Allocate each bit_gadget of the result
239        let mut coeff = ConstraintF::one();
240        let mut i = 0;
241        while max_value != 0 {
242            // Allocate the bit_gadget
243            let b = AllocatedBit::alloc(cs.ns(|| format!("result bit_gadget {}", i)), || {
244                result_value.map(|v| (v >> i) & 1 == 1).get()
245            })?;
246
247            // Subtract this bit_gadget from the linear combination to ensure the sums
248            // balance out
249            lc = lc - (coeff, b.get_variable());
250
251            result_bits.push(b.into());
252
253            max_value >>= 1;
254            i += 1;
255            coeff.double_in_place();
256        }
257
258        // Enforce that the linear combination equals zero
259        cs.enforce(|| "modular addition", |lc| lc, |lc| lc, |_| lc);
260
261        // Discard carry bits that we don't care about
262        result_bits.truncate(32);
263
264        Ok(UInt32 {
265            bits:  result_bits,
266            value: modular_value,
267        })
268    }
269}
270
271impl<ConstraintF: Field> ToBytesGadget<ConstraintF> for UInt32 {
272    #[inline]
273    fn to_bytes<CS: ConstraintSystem<ConstraintF>>(&self, _cs: CS) -> Result<Vec<UInt8>, SynthesisError> {
274        let value_chunks = match self.value.map(|val| {
275            use algebra::bytes::ToBytes;
276            let mut bytes = [0u8; 4];
277            val.write(bytes.as_mut()).unwrap();
278            bytes
279        }) {
280            Some(chunks) => [
281                Some(chunks[0]),
282                Some(chunks[1]),
283                Some(chunks[2]),
284                Some(chunks[3]),
285            ],
286            None => [None, None, None, None],
287        };
288        let mut bytes = Vec::new();
289        for (i, chunk8) in self.to_bits_le().chunks(8).into_iter().enumerate() {
290            let byte = UInt8 {
291                bits:  chunk8.to_vec(),
292                value: value_chunks[i],
293            };
294            bytes.push(byte);
295        }
296
297        Ok(bytes)
298    }
299
300    fn to_bytes_strict<CS: ConstraintSystem<ConstraintF>>(
301        &self,
302        cs: CS,
303    ) -> Result<Vec<UInt8>, SynthesisError> {
304        self.to_bytes(cs)
305    }
306}
307
308impl PartialEq for UInt32 {
309    fn eq(&self, other: &Self) -> bool {
310        !self.value.is_none() && !other.value.is_none() && self.value == other.value
311    }
312}
313
314impl Eq for UInt32 {}
315
316impl<ConstraintF: Field> ConditionalEqGadget<ConstraintF> for UInt32 {
317    fn conditional_enforce_equal<CS: ConstraintSystem<ConstraintF>>(
318        &self,
319        mut cs: CS,
320        other: &Self,
321        condition: &Boolean,
322    ) -> Result<(), SynthesisError> {
323        for (i, (a, b)) in self.bits.iter().zip(&other.bits).enumerate() {
324            a.conditional_enforce_equal(
325                &mut cs.ns(|| format!("uint32_equal_{}", i)),
326                b,
327                condition,
328            )?;
329        }
330        Ok(())
331    }
332
333    fn cost() -> usize {
334        32 * <Boolean as ConditionalEqGadget<ConstraintF>>::cost()
335    }
336}
337
338#[cfg(test)]
339mod test {
340    use super::UInt32;
341    use crate::{bits::boolean::Boolean, test_constraint_system::TestConstraintSystem};
342    use algebra::fields::{bls12_381::Fr, Field};
343    use rand::{Rng, SeedableRng, XorShiftRng};
344    use r1cs_core::ConstraintSystem;
345
346    #[test]
347    fn test_uint32_from_bits() {
348        let mut rng = XorShiftRng::from_seed([0x5dbe6259, 0x8d313d76, 0x3237db17, 0xe5bc0653]);
349
350        for _ in 0..1000 {
351            let v = (0..32)
352                .map(|_| Boolean::constant(rng.gen()))
353                .collect::<Vec<_>>();
354
355            let b = UInt32::from_bits_le(&v);
356
357            for (i, bit_gadget) in b.bits.iter().enumerate() {
358                match bit_gadget {
359                    &Boolean::Constant(bit_gadget) => {
360                        assert!(bit_gadget == ((b.value.unwrap() >> i) & 1 == 1));
361                    },
362                    _ => unreachable!(),
363                }
364            }
365
366            let expected_to_be_same = b.to_bits_le();
367
368            for x in v.iter().zip(expected_to_be_same.iter()) {
369                match x {
370                    (&Boolean::Constant(true), &Boolean::Constant(true)) => {},
371                    (&Boolean::Constant(false), &Boolean::Constant(false)) => {},
372                    _ => unreachable!(),
373                }
374            }
375        }
376    }
377
378    #[test]
379    fn test_uint32_xor() {
380        let mut rng = XorShiftRng::from_seed([0x5dbe6259, 0x8d313d76, 0x3237db17, 0xe5bc0653]);
381
382        for _ in 0..1000 {
383            let mut cs = TestConstraintSystem::<Fr>::new();
384
385            let a: u32 = rng.gen();
386            let b: u32 = rng.gen();
387            let c: u32 = rng.gen();
388
389            let mut expected = a ^ b ^ c;
390
391            let a_bit = UInt32::alloc(cs.ns(|| "a_bit"), Some(a)).unwrap();
392            let b_bit = UInt32::constant(b);
393            let c_bit = UInt32::alloc(cs.ns(|| "c_bit"), Some(c)).unwrap();
394
395            let r = a_bit.xor(cs.ns(|| "first xor"), &b_bit).unwrap();
396            let r = r.xor(cs.ns(|| "second xor"), &c_bit).unwrap();
397
398            assert!(cs.is_satisfied());
399
400            assert!(r.value == Some(expected));
401
402            for b in r.bits.iter() {
403                match b {
404                    &Boolean::Is(ref b) => {
405                        assert!(b.get_value().unwrap() == (expected & 1 == 1));
406                    },
407                    &Boolean::Not(ref b) => {
408                        assert!(!b.get_value().unwrap() == (expected & 1 == 1));
409                    },
410                    &Boolean::Constant(b) => {
411                        assert!(b == (expected & 1 == 1));
412                    },
413                }
414
415                expected >>= 1;
416            }
417        }
418    }
419
420    #[test]
421    fn test_uint32_addmany_constants() {
422        let mut rng = XorShiftRng::from_seed([0x5dbe6259, 0x8d313d76, 0x3237db17, 0xe5bc0654]);
423
424        for _ in 0..1000 {
425            let mut cs = TestConstraintSystem::<Fr>::new();
426
427            let a: u32 = rng.gen();
428            let b: u32 = rng.gen();
429            let c: u32 = rng.gen();
430
431            let a_bit = UInt32::constant(a);
432            let b_bit = UInt32::constant(b);
433            let c_bit = UInt32::constant(c);
434
435            let mut expected = a.wrapping_add(b).wrapping_add(c);
436
437            let r = UInt32::addmany(cs.ns(|| "addition"), &[a_bit, b_bit, c_bit]).unwrap();
438
439            assert!(r.value == Some(expected));
440
441            for b in r.bits.iter() {
442                match b {
443                    &Boolean::Is(_) => panic!(),
444                    &Boolean::Not(_) => panic!(),
445                    &Boolean::Constant(b) => {
446                        assert!(b == (expected & 1 == 1));
447                    },
448                }
449
450                expected >>= 1;
451            }
452        }
453    }
454
455    #[test]
456    fn test_uint32_addmany() {
457        let mut rng = XorShiftRng::from_seed([0x5dbe6259, 0x8d313d76, 0x3237db17, 0xe5bc0654]);
458
459        for _ in 0..1000 {
460            let mut cs = TestConstraintSystem::<Fr>::new();
461
462            let a: u32 = rng.gen();
463            let b: u32 = rng.gen();
464            let c: u32 = rng.gen();
465            let d: u32 = rng.gen();
466
467            let mut expected = (a ^ b).wrapping_add(c).wrapping_add(d);
468
469            let a_bit = UInt32::alloc(cs.ns(|| "a_bit"), Some(a)).unwrap();
470            let b_bit = UInt32::constant(b);
471            let c_bit = UInt32::constant(c);
472            let d_bit = UInt32::alloc(cs.ns(|| "d_bit"), Some(d)).unwrap();
473
474            let r = a_bit.xor(cs.ns(|| "xor"), &b_bit).unwrap();
475            let r = UInt32::addmany(cs.ns(|| "addition"), &[r, c_bit, d_bit]).unwrap();
476
477            assert!(cs.is_satisfied());
478
479            assert!(r.value == Some(expected));
480
481            for b in r.bits.iter() {
482                match b {
483                    &Boolean::Is(ref b) => {
484                        assert!(b.get_value().unwrap() == (expected & 1 == 1));
485                    },
486                    &Boolean::Not(ref b) => {
487                        assert!(!b.get_value().unwrap() == (expected & 1 == 1));
488                    },
489                    &Boolean::Constant(_) => unreachable!(),
490                }
491
492                expected >>= 1;
493            }
494
495            // Flip a bit_gadget and see if the addition constraint still works
496            if cs.get("addition/result bit_gadget 0/boolean").is_zero() {
497                cs.set("addition/result bit_gadget 0/boolean", Field::one());
498            } else {
499                cs.set("addition/result bit_gadget 0/boolean", Field::zero());
500            }
501
502            assert!(!cs.is_satisfied());
503        }
504    }
505
506    #[test]
507    fn test_uint32_rotr() {
508        let mut rng = XorShiftRng::from_seed([0x5dbe6259, 0x8d313d76, 0x3237db17, 0xe5bc0654]);
509
510        let mut num = rng.gen();
511
512        let a = UInt32::constant(num);
513
514        for i in 0..32 {
515            let b = a.rotr(i);
516
517            assert!(b.value.unwrap() == num);
518
519            let mut tmp = num;
520            for b in &b.bits {
521                match b {
522                    &Boolean::Constant(b) => {
523                        assert_eq!(b, tmp & 1 == 1);
524                    },
525                    _ => unreachable!(),
526                }
527
528                tmp >>= 1;
529            }
530
531            num = num.rotate_right(1);
532        }
533    }
534}