1use algebra::{Field, FpParameters, PrimeField};
  2
  3use r1cs_core::{ConstraintSystem, LinearCombination, SynthesisError};
  4
  5use crate::{
  6    boolean::{AllocatedBit, Boolean},
  7    prelude::*,
  8    Assignment,
  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    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        assert!(operands.len() >= 2); // Weird trivial cases that should never happen
171        assert!(operands.len() <= 10);
172
173        // Compute the maximum value of the sum so we allocate enough bits for
174        // the result
175        let mut max_value = (operands.len() as u64) * u64::from(u32::max_value());
176
177        // Keep track of the resulting value
178        let mut result_value = Some(0u64);
179
180        // This is a linear combination that we will enforce to be "zero"
181        let mut lc = LinearCombination::zero();
182
183        let mut all_constants = true;
184
185        // Iterate over the operands
186        for op in operands {
187            // Accumulate the value
188            match op.value {
189                Some(val) => {
190                    result_value.as_mut().map(|v| *v += u64::from(val));
191                },
192                None => {
193                    // If any of our operands have unknown value, we won't
194                    // know the value of the result
195                    result_value = None;
196                },
197            }
198
199            // Iterate over each bit_gadget of the operand and add the operand to
200            // the linear combination
201            let mut coeff = ConstraintF::one();
202            for bit in &op.bits {
203                match *bit {
204                    Boolean::Is(ref bit) => {
205                        all_constants = false;
206
207                        // Add coeff * bit_gadget
208                        lc += (coeff, bit.get_variable());
209                    },
210                    Boolean::Not(ref bit) => {
211                        all_constants = false;
212
213                        // Add coeff * (1 - bit_gadget) = coeff * ONE - coeff * bit_gadget
214                        lc = lc + (coeff, CS::one()) - (coeff, bit.get_variable());
215                    },
216                    Boolean::Constant(bit) => {
217                        if bit {
218                            lc += (coeff, CS::one());
219                        }
220                    },
221                }
222
223                coeff.double_in_place();
224            }
225        }
226
227        // The value of the actual result is modulo 2^32
228        let modular_value = result_value.map(|v| v as u32);
229
230        if all_constants && modular_value.is_some() {
231            // We can just return a constant, rather than
232            // unpacking the result into allocated bits.
233
234            return Ok(UInt32::constant(modular_value.unwrap()));
235        }
236
237        // Storage area for the resulting bits
238        let mut result_bits = vec![];
239
240        // Allocate each bit_gadget of the result
241        let mut coeff = ConstraintF::one();
242        let mut i = 0;
243        while max_value != 0 {
244            // Allocate the bit_gadget
245            let b = AllocatedBit::alloc(cs.ns(|| format!("result bit_gadget {}", i)), || {
246                result_value.map(|v| (v >> i) & 1 == 1).get()
247            })?;
248
249            // Subtract this bit_gadget from the linear combination to ensure the sums
250            // balance out
251            lc = lc - (coeff, b.get_variable());
252
253            result_bits.push(b.into());
254
255            max_value >>= 1;
256            i += 1;
257            coeff.double_in_place();
258        }
259
260        // Enforce that the linear combination equals zero
261        cs.enforce(|| "modular addition", |lc| lc, |lc| lc, |_| lc);
262
263        // Discard carry bits that we don't care about
264        result_bits.truncate(32);
265
266        Ok(UInt32 {
267            bits:  result_bits,
268            value: modular_value,
269        })
270    }
271}
272
273impl<ConstraintF: Field> ToBytesGadget<ConstraintF> for UInt32 {
274    #[inline]
275    fn to_bytes<CS: ConstraintSystem<ConstraintF>>(
276        &self,
277        _cs: CS,
278    ) -> Result<Vec<UInt8>, SynthesisError> {
279        let value_chunks = match self.value.map(|val| {
280            use algebra::bytes::ToBytes;
281            let mut bytes = [0u8; 4];
282            val.write(bytes.as_mut()).unwrap();
283            bytes
284        }) {
285            Some(chunks) => [
286                Some(chunks[0]),
287                Some(chunks[1]),
288                Some(chunks[2]),
289                Some(chunks[3]),
290            ],
291            None => [None, None, None, None],
292        };
293        let mut bytes = Vec::new();
294        for (i, chunk8) in self.to_bits_le().chunks(8).enumerate() {
295            let byte = UInt8 {
296                bits:  chunk8.to_vec(),
297                value: value_chunks[i],
298            };
299            bytes.push(byte);
300        }
301
302        Ok(bytes)
303    }
304
305    fn to_bytes_strict<CS: ConstraintSystem<ConstraintF>>(
306        &self,
307        cs: CS,
308    ) -> Result<Vec<UInt8>, SynthesisError> {
309        self.to_bytes(cs)
310    }
311}
312
313impl PartialEq for UInt32 {
314    fn eq(&self, other: &Self) -> bool {
315        self.value.is_some() && other.value.is_some() && self.value == other.value
316    }
317}
318
319impl Eq for UInt32 {}
320
321impl<ConstraintF: Field> ConditionalEqGadget<ConstraintF> for UInt32 {
322    fn conditional_enforce_equal<CS: ConstraintSystem<ConstraintF>>(
323        &self,
324        mut cs: CS,
325        other: &Self,
326        condition: &Boolean,
327    ) -> Result<(), SynthesisError> {
328        for (i, (a, b)) in self.bits.iter().zip(&other.bits).enumerate() {
329            a.conditional_enforce_equal(
330                &mut cs.ns(|| format!("uint32_equal_{}", i)),
331                b,
332                condition,
333            )?;
334        }
335        Ok(())
336    }
337
338    fn cost() -> usize {
339        32 * <Boolean as ConditionalEqGadget<ConstraintF>>::cost()
340    }
341}
342
343#[cfg(test)]
344mod test {
345    use super::UInt32;
346    use crate::{bits::boolean::Boolean, test_constraint_system::TestConstraintSystem};
347    use algebra::{fields::bls12_381::Fr, One, Zero};
348    use r1cs_core::ConstraintSystem;
349    use rand::{Rng, SeedableRng};
350    use rand_xorshift::XorShiftRng;
351
352    #[test]
353    fn test_uint32_from_bits() {
354        let mut rng = XorShiftRng::seed_from_u64(1231275789u64);
355
356        for _ in 0..1000 {
357            let v = (0..32)
358                .map(|_| Boolean::constant(rng.gen()))
359                .collect::<Vec<_>>();
360
361            let b = UInt32::from_bits_le(&v);
362
363            for (i, bit_gadget) in b.bits.iter().enumerate() {
364                match bit_gadget {
365                    &Boolean::Constant(bit_gadget) => {
366                        assert!(bit_gadget == ((b.value.unwrap() >> i) & 1 == 1));
367                    },
368                    _ => unreachable!(),
369                }
370            }
371
372            let expected_to_be_same = b.to_bits_le();
373
374            for x in v.iter().zip(expected_to_be_same.iter()) {
375                match x {
376                    (&Boolean::Constant(true), &Boolean::Constant(true)) => {},
377                    (&Boolean::Constant(false), &Boolean::Constant(false)) => {},
378                    _ => unreachable!(),
379                }
380            }
381        }
382    }
383
384    #[test]
385    fn test_uint32_xor() {
386        let mut rng = XorShiftRng::seed_from_u64(1231275789u64);
387
388        for _ in 0..1000 {
389            let mut cs = TestConstraintSystem::<Fr>::new();
390
391            let a: u32 = rng.gen();
392            let b: u32 = rng.gen();
393            let c: u32 = rng.gen();
394
395            let mut expected = a ^ b ^ c;
396
397            let a_bit = UInt32::alloc(cs.ns(|| "a_bit"), Some(a)).unwrap();
398            let b_bit = UInt32::constant(b);
399            let c_bit = UInt32::alloc(cs.ns(|| "c_bit"), Some(c)).unwrap();
400
401            let r = a_bit.xor(cs.ns(|| "first xor"), &b_bit).unwrap();
402            let r = r.xor(cs.ns(|| "second xor"), &c_bit).unwrap();
403
404            assert!(cs.is_satisfied());
405
406            assert!(r.value == Some(expected));
407
408            for b in r.bits.iter() {
409                match b {
410                    &Boolean::Is(ref b) => {
411                        assert!(b.get_value().unwrap() == (expected & 1 == 1));
412                    },
413                    &Boolean::Not(ref b) => {
414                        assert!(!b.get_value().unwrap() == (expected & 1 == 1));
415                    },
416                    &Boolean::Constant(b) => {
417                        assert!(b == (expected & 1 == 1));
418                    },
419                }
420
421                expected >>= 1;
422            }
423        }
424    }
425
426    #[test]
427    fn test_uint32_addmany_constants() {
428        let mut rng = XorShiftRng::seed_from_u64(1231275789u64);
429
430        for _ in 0..1000 {
431            let mut cs = TestConstraintSystem::<Fr>::new();
432
433            let a: u32 = rng.gen();
434            let b: u32 = rng.gen();
435            let c: u32 = rng.gen();
436
437            let a_bit = UInt32::constant(a);
438            let b_bit = UInt32::constant(b);
439            let c_bit = UInt32::constant(c);
440
441            let mut expected = a.wrapping_add(b).wrapping_add(c);
442
443            let r = UInt32::addmany(cs.ns(|| "addition"), &[a_bit, b_bit, c_bit]).unwrap();
444
445            assert!(r.value == Some(expected));
446
447            for b in r.bits.iter() {
448                match b {
449                    &Boolean::Is(_) => panic!(),
450                    &Boolean::Not(_) => panic!(),
451                    &Boolean::Constant(b) => {
452                        assert!(b == (expected & 1 == 1));
453                    },
454                }
455
456                expected >>= 1;
457            }
458        }
459    }
460
461    #[test]
462    fn test_uint32_addmany() {
463        let mut rng = XorShiftRng::seed_from_u64(1231275789u64);
464
465        for _ in 0..1000 {
466            let mut cs = TestConstraintSystem::<Fr>::new();
467
468            let a: u32 = rng.gen();
469            let b: u32 = rng.gen();
470            let c: u32 = rng.gen();
471            let d: u32 = rng.gen();
472
473            let mut expected = (a ^ b).wrapping_add(c).wrapping_add(d);
474
475            let a_bit = UInt32::alloc(cs.ns(|| "a_bit"), Some(a)).unwrap();
476            let b_bit = UInt32::constant(b);
477            let c_bit = UInt32::constant(c);
478            let d_bit = UInt32::alloc(cs.ns(|| "d_bit"), Some(d)).unwrap();
479
480            let r = a_bit.xor(cs.ns(|| "xor"), &b_bit).unwrap();
481            let r = UInt32::addmany(cs.ns(|| "addition"), &[r, c_bit, d_bit]).unwrap();
482
483            assert!(cs.is_satisfied());
484
485            assert!(r.value == Some(expected));
486
487            for b in r.bits.iter() {
488                match b {
489                    &Boolean::Is(ref b) => {
490                        assert!(b.get_value().unwrap() == (expected & 1 == 1));
491                    },
492                    &Boolean::Not(ref b) => {
493                        assert!(!b.get_value().unwrap() == (expected & 1 == 1));
494                    },
495                    &Boolean::Constant(_) => unreachable!(),
496                }
497
498                expected >>= 1;
499            }
500
501            // Flip a bit_gadget and see if the addition constraint still works
502            if cs.get("addition/result bit_gadget 0/boolean").is_zero() {
503                cs.set("addition/result bit_gadget 0/boolean", Fr::one());
504            } else {
505                cs.set("addition/result bit_gadget 0/boolean", Fr::zero());
506            }
507
508            assert!(!cs.is_satisfied());
509        }
510    }
511
512    #[test]
513    fn test_uint32_rotr() {
514        let mut rng = XorShiftRng::seed_from_u64(1231275789u64);
515
516        let mut num = rng.gen();
517
518        let a = UInt32::constant(num);
519
520        for i in 0..32 {
521            let b = a.rotr(i);
522
523            assert!(b.value.unwrap() == num);
524
525            let mut tmp = num;
526            for b in &b.bits {
527                match b {
528                    &Boolean::Constant(b) => {
529                        assert_eq!(b, tmp & 1 == 1);
530                    },
531                    _ => unreachable!(),
532                }
533
534                tmp >>= 1;
535            }
536
537            num = num.rotate_right(1);
538        }
539    }
540}