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