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