use crate::{
curves::{Field, FpParameters, PrimeField},
gadgets::{
r1cs::{Assignment, ConstraintSystem, LinearCombination},
utilities::{
alloc::AllocGadget,
boolean::{AllocatedBit, Boolean},
eq::{ConditionalEqGadget, EqGadget, EvaluateEqGadget},
select::CondSelectGadget,
uint::unsigned_integer::{UInt, UInt8},
ToBytesGadget,
},
},
};
use snarkvm_errors::gadgets::SynthesisError;
use snarkvm_utilities::{
biginteger::{BigInteger, BigInteger256},
bytes::ToBytes,
};
use std::{borrow::Borrow, cmp::Ordering};
#[derive(Clone, Debug)]
pub struct UInt128 {
pub bits: Vec<Boolean>,
pub negated: bool,
pub value: Option<u128>,
}
impl UInt128 {
pub fn constant(value: u128) -> Self {
let mut bits = Vec::with_capacity(128);
let mut tmp = value;
for _ in 0..128 {
if tmp & 1 == 1 {
bits.push(Boolean::constant(true))
} else {
bits.push(Boolean::constant(false))
}
tmp >>= 1;
}
Self {
bits,
negated: false,
value: Some(value),
}
}
}
impl UInt for UInt128 {
fn negate(&self) -> Self {
Self {
bits: self.bits.clone(),
negated: true,
value: self.value,
}
}
fn is_constant(&self) -> bool {
let mut constant = true;
for bit in &self.bits {
match *bit {
Boolean::Is(ref _bit) => constant = false,
Boolean::Not(ref _bit) => constant = false,
Boolean::Constant(_bit) => {}
}
}
constant
}
fn to_bits_le(&self) -> Vec<Boolean> {
self.bits.clone()
}
fn from_bits_le(bits: &[Boolean]) -> Self {
assert_eq!(bits.len(), 128);
let bits = bits.to_vec();
let mut value = Some(0u128);
for b in bits.iter().rev() {
if let Some(v) = value.as_mut() {
*v <<= 1;
}
match *b {
Boolean::Constant(b) => {
if b {
if let Some(v) = value.as_mut() {
*v |= 1;
}
}
}
Boolean::Is(ref b) => match b.get_value() {
Some(true) => {
if let Some(v) = value.as_mut() {
*v |= 1;
}
}
Some(false) => {}
None => value = None,
},
Boolean::Not(ref b) => match b.get_value() {
Some(false) => {
if let Some(v) = value.as_mut() {
*v |= 1;
}
}
Some(true) => {}
None => value = None,
},
}
}
Self {
value,
negated: false,
bits,
}
}
fn rotr(&self, by: usize) -> Self {
let by = by % 128;
let new_bits = self
.bits
.iter()
.skip(by)
.chain(self.bits.iter())
.take(128)
.cloned()
.collect();
Self {
bits: new_bits,
negated: false,
value: self.value.map(|v| v.rotate_right(by as u32) as u128),
}
}
fn xor<F: Field, CS: ConstraintSystem<F>>(&self, mut cs: CS, other: &Self) -> Result<Self, SynthesisError> {
let new_value = match (self.value, other.value) {
(Some(a), Some(b)) => Some(a ^ b),
_ => None,
};
let bits = self
.bits
.iter()
.zip(other.bits.iter())
.enumerate()
.map(|(i, (a, b))| Boolean::xor(cs.ns(|| format!("xor of bit_gadget {}", i)), a, b))
.collect::<Result<_, _>>()?;
Ok(Self {
bits,
negated: false,
value: new_value,
})
}
fn addmany<F: PrimeField, CS: ConstraintSystem<F>>(mut cs: CS, operands: &[Self]) -> Result<Self, SynthesisError> {
assert!(F::Parameters::MODULUS_BITS <= 253);
assert!(operands.len() >= 2);
let mut max_value = BigInteger256::from_u128(u128::max_value());
max_value.muln(operands.len() as u32);
let mut big_result_value = Some(BigInteger256::default());
let mut lc = LinearCombination::zero();
let mut all_constants = true;
for op in operands {
match op.value {
Some(val) => {
if op.negated {
big_result_value
.as_mut()
.map(|v| v.sub_noborrow(&BigInteger256::from_u128(val)));
} else {
big_result_value
.as_mut()
.map(|v| v.add_nocarry(&BigInteger256::from_u128(val)));
}
}
None => {
big_result_value = None;
}
}
let mut coeff = F::one();
for bit in &op.bits {
match *bit {
Boolean::Is(ref bit) => {
all_constants = false;
if op.negated {
lc = lc - (coeff, bit.get_variable());
} else {
lc += (coeff, bit.get_variable());
}
}
Boolean::Not(ref bit) => {
all_constants = false;
if op.negated {
lc = lc - (coeff, CS::one()) + (coeff, bit.get_variable());
} else {
lc = lc + (coeff, CS::one()) - (coeff, bit.get_variable());
}
}
Boolean::Constant(bit) => {
if bit {
if op.negated {
lc = lc - (coeff, CS::one());
} else {
lc += (coeff, CS::one());
}
}
}
}
coeff.double_in_place();
}
}
let modular_value = big_result_value.map(|v| v.to_u128());
if all_constants {
if let Some(val) = modular_value {
return Ok(Self::constant(val));
}
}
let mut result_bits = vec![];
let mut coeff = F::one();
let mut i = 0;
while !max_value.is_zero() {
let b = AllocatedBit::alloc(cs.ns(|| format!("result bit_gadget {}", i)), || {
big_result_value.map(|v| v.get_bit(i)).get()
})?;
lc = lc - (coeff, b.get_variable());
result_bits.push(b.into());
max_value.div2();
i += 1;
coeff.double_in_place();
}
cs.enforce(|| "modular addition", |lc| lc, |lc| lc, |_| lc);
result_bits.truncate(128);
Ok(Self {
bits: result_bits,
negated: false,
value: modular_value,
})
}
fn sub<F: PrimeField, CS: ConstraintSystem<F>>(&self, mut cs: CS, other: &Self) -> Result<Self, SynthesisError> {
Self::addmany(&mut cs.ns(|| "add_not"), &[self.clone(), other.negate()])
}
fn sub_unsafe<F: PrimeField, CS: ConstraintSystem<F>>(
&self,
mut cs: CS,
other: &Self,
) -> Result<Self, SynthesisError> {
match (self.value, other.value) {
(Some(val1), Some(val2)) => {
if val1 < val2 {
if Self::result_is_constant(&self, &other) {
Ok(Self::constant(0u128))
} else {
let result_value = Some(0u128);
let modular_value = result_value.map(|v| v as u128);
let mut result_bits = Vec::with_capacity(128);
let mut lc = LinearCombination::zero();
let mut coeff = F::one();
for i in 0..128 {
let b = AllocatedBit::alloc(cs.ns(|| format!("result bit_gadget {}", i)), || {
result_value.map(|v| (v >> i) & 1 == 1).get()
})?;
lc = lc - (coeff, b.get_variable());
result_bits.push(b.into());
coeff.double_in_place();
}
cs.enforce(|| "unsafe subtraction", |lc| lc, |lc| lc, |_| lc);
result_bits.truncate(128);
Ok(Self {
bits: result_bits,
negated: false,
value: modular_value,
})
}
} else {
self.sub(&mut cs.ns(|| ""), &other)
}
}
(_, _) => {
Err(SynthesisError::AssignmentMissing)
}
}
}
fn mul<F: PrimeField, CS: ConstraintSystem<F>>(&self, mut cs: CS, other: &Self) -> Result<Self, SynthesisError> {
let is_constant = Boolean::constant(Self::result_is_constant(&self, &other));
let constant_result = Self::constant(0u128);
let allocated_result = Self::alloc(&mut cs.ns(|| "allocated_1u128"), || Ok(0u128))?;
let zero_result = Self::conditionally_select(
&mut cs.ns(|| "constant_or_allocated"),
&is_constant,
&constant_result,
&allocated_result,
)?;
let mut left_shift = self.clone();
let partial_products = other
.bits
.iter()
.enumerate()
.map(|(i, bit)| {
let current_left_shift = left_shift.clone();
left_shift = Self::addmany(&mut cs.ns(|| format!("shift_left_{}", i)), &[
left_shift.clone(),
left_shift.clone(),
])
.unwrap();
Self::conditionally_select(
&mut cs.ns(|| format!("calculate_product_{}", i)),
&bit,
¤t_left_shift,
&zero_result,
)
.unwrap()
})
.collect::<Vec<Self>>();
Self::addmany(&mut cs.ns(|| "partial_products"), &partial_products)
}
fn div<F: PrimeField, CS: ConstraintSystem<F>>(&self, mut cs: CS, other: &Self) -> Result<Self, SynthesisError> {
if other.eq(&Self::constant(0u128)) {
return Err(SynthesisError::DivisionByZero);
}
let is_constant = Boolean::constant(Self::result_is_constant(&self, &other));
let allocated_true = Boolean::from(AllocatedBit::alloc(&mut cs.ns(|| "true"), || Ok(true)).unwrap());
let true_bit = Boolean::conditionally_select(
&mut cs.ns(|| "constant_or_allocated_true"),
&is_constant,
&Boolean::constant(true),
&allocated_true,
)?;
let allocated_one = Self::alloc(&mut cs.ns(|| "one"), || Ok(1u128))?;
let one = Self::conditionally_select(
&mut cs.ns(|| "constant_or_allocated_1u128"),
&is_constant,
&Self::constant(1u128),
&allocated_one,
)?;
let allocated_zero = Self::alloc(&mut cs.ns(|| "zero"), || Ok(0u128))?;
let zero = Self::conditionally_select(
&mut cs.ns(|| "constant_or_allocated_0u128"),
&is_constant,
&Self::constant(0u128),
&allocated_zero,
)?;
let self_is_zero = Boolean::Constant(self.eq(&Self::constant(0u128)));
let mut quotient = zero.clone();
let mut remainder = zero;
for (i, bit) in self.bits.iter().rev().enumerate() {
remainder = Self::addmany(&mut cs.ns(|| format!("shift_left_{}", i)), &[
remainder.clone(),
remainder.clone(),
])?;
let bit_is_true = Boolean::constant(bit.eq(&Boolean::constant(true)));
let new_remainder = Self::addmany(&mut cs.ns(|| format!("set_remainder_bit_{}", i)), &[
remainder.clone(),
one.clone(),
])?;
remainder = Self::conditionally_select(
&mut cs.ns(|| format!("increment_or_remainder_{}", i)),
&bit_is_true,
&new_remainder,
&remainder,
)?;
let no_remainder = Boolean::constant(remainder.eq(&other));
let subtraction = remainder.sub_unsafe(&mut cs.ns(|| format!("subtract_divisor_{}", i)), &other)?;
let sub_is_zero = Boolean::constant(subtraction.eq(&Self::constant(0)));
let cond1 = Boolean::and(
&mut cs.ns(|| format!("cond_1_{}", i)),
&no_remainder.not(),
&sub_is_zero.not(),
)?;
let cond2 = Boolean::or(&mut cs.ns(|| format!("cond_2_{}", i)), &no_remainder, &cond1)?;
remainder = Self::conditionally_select(
&mut cs.ns(|| format!("subtract_or_same_{}", i)),
&cond2,
&subtraction,
&remainder,
)?;
let index = 127 - i as usize;
let bit_value = 1u128 << (index as u128);
let mut new_quotient = quotient.clone();
new_quotient.bits[index] = true_bit;
new_quotient.value = Some(new_quotient.value.unwrap() + bit_value);
quotient = Self::conditionally_select(
&mut cs.ns(|| format!("set_bit_or_same_{}", i)),
&cond2,
&new_quotient,
"ient,
)?;
}
Self::conditionally_select(&mut cs.ns(|| "self_or_quotient"), &self_is_zero, self, "ient)
}
fn pow<F: Field + PrimeField, CS: ConstraintSystem<F>>(
&self,
mut cs: CS,
other: &Self,
) -> Result<Self, SynthesisError> {
let is_constant = Boolean::constant(Self::result_is_constant(&self, &other));
let constant_result = Self::constant(1u128);
let allocated_result = Self::alloc(&mut cs.ns(|| "allocated_1u128"), || Ok(1u128))?;
let mut result = Self::conditionally_select(
&mut cs.ns(|| "constant_or_allocated"),
&is_constant,
&constant_result,
&allocated_result,
)?;
for (i, bit) in other.bits.iter().rev().enumerate() {
let found_one = Boolean::Constant(result.eq(&Self::constant(1u128)));
let cond1 = Boolean::and(cs.ns(|| format!("found_one_{}", i)), &bit.not(), &found_one)?;
let square = result.mul(cs.ns(|| format!("square_{}", i)), &result).unwrap();
result = Self::conditionally_select(
&mut cs.ns(|| format!("result_or_sqaure_{}", i)),
&cond1,
&result,
&square,
)?;
let mul_by_self = result.mul(cs.ns(|| format!("multiply_by_self_{}", i)), &self).unwrap();
result = Self::conditionally_select(
&mut cs.ns(|| format!("mul_by_self_or_result_{}", i)),
&bit,
&mul_by_self,
&result,
)?;
}
Ok(result)
}
}
impl PartialEq for UInt128 {
fn eq(&self, other: &Self) -> bool {
self.value.is_some() && other.value.is_some() && self.value == other.value
}
}
impl Eq for UInt128 {}
impl PartialOrd for UInt128 {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
Option::from(self.value.cmp(&other.value))
}
}
impl<F: PrimeField> EvaluateEqGadget<F> for UInt128 {
fn evaluate_equal<CS: ConstraintSystem<F>>(&self, mut cs: CS, other: &Self) -> Result<Boolean, SynthesisError> {
let mut result = Boolean::constant(true);
for (i, (a, b)) in self.bits.iter().zip(&other.bits).enumerate() {
let equal = a.evaluate_equal(&mut cs.ns(|| format!("u128 evaluate equality for {}-th bit", i)), b)?;
result = Boolean::and(
&mut cs.ns(|| format!("u128 and result for {}-th bit", i)),
&equal,
&result,
)?;
}
Ok(result)
}
}
impl<F: Field> EqGadget<F> for UInt128 {}
impl<F: Field> ConditionalEqGadget<F> for UInt128 {
fn conditional_enforce_equal<CS: ConstraintSystem<F>>(
&self,
mut cs: CS,
other: &Self,
condition: &Boolean,
) -> Result<(), SynthesisError> {
for (i, (a, b)) in self.bits.iter().zip(&other.bits).enumerate() {
a.conditional_enforce_equal(&mut cs.ns(|| format!("uint128_equal_{}", i)), b, condition)?;
}
Ok(())
}
fn cost() -> usize {
128 * <Boolean as ConditionalEqGadget<F>>::cost()
}
}
impl<F: PrimeField> CondSelectGadget<F> for UInt128 {
fn conditionally_select<CS: ConstraintSystem<F>>(
mut cs: CS,
cond: &Boolean,
first: &Self,
second: &Self,
) -> Result<Self, SynthesisError> {
if let Boolean::Constant(cond) = *cond {
if cond { Ok(first.clone()) } else { Ok(second.clone()) }
} else {
let mut is_negated = false;
let result_val = cond.get_value().and_then(|c| {
if c {
is_negated = first.negated;
first.value
} else {
is_negated = second.negated;
second.value
}
});
let mut result = Self::alloc(cs.ns(|| "cond_select_result"), || result_val.get())?;
result.negated = is_negated;
let expected_bits = first
.bits
.iter()
.zip(&second.bits)
.enumerate()
.map(|(i, (a, b))| {
Boolean::conditionally_select(&mut cs.ns(|| format!("uint128_cond_select_{}", i)), cond, a, b)
.unwrap()
})
.collect::<Vec<Boolean>>();
for (i, (actual, expected)) in result.to_bits_le().iter().zip(expected_bits.iter()).enumerate() {
actual.enforce_equal(&mut cs.ns(|| format!("selected_result_bit_{}", i)), expected)?;
}
Ok(result)
}
}
fn cost() -> usize {
128 * (<Boolean as ConditionalEqGadget<F>>::cost() + <Boolean as CondSelectGadget<F>>::cost())
}
}
impl<F: Field> AllocGadget<u128, F> for UInt128 {
fn alloc<Fn: FnOnce() -> Result<T, SynthesisError>, T: Borrow<u128>, CS: ConstraintSystem<F>>(
mut cs: CS,
value_gen: Fn,
) -> Result<Self, SynthesisError> {
let value = value_gen().map(|val| *val.borrow());
let values = match value {
Ok(mut val) => {
let mut v = Vec::with_capacity(128);
for _ in 0..128 {
v.push(Some(val & 1 == 1));
val >>= 1;
}
v
}
_ => vec![None; 128],
};
let bits = values
.into_iter()
.enumerate()
.map(|(i, v)| {
Ok(Boolean::from(AllocatedBit::alloc(
&mut cs.ns(|| format!("allocated bit_gadget {}", i)),
|| v.ok_or(SynthesisError::AssignmentMissing),
)?))
})
.collect::<Result<Vec<_>, SynthesisError>>()?;
Ok(Self {
bits,
negated: false,
value: value.ok(),
})
}
fn alloc_input<Fn, T, CS: ConstraintSystem<F>>(mut cs: CS, value_gen: Fn) -> Result<Self, SynthesisError>
where
Fn: FnOnce() -> Result<T, SynthesisError>,
T: Borrow<u128>,
{
let value = value_gen().map(|val| *val.borrow());
let values = match value {
Ok(mut val) => {
let mut v = Vec::with_capacity(128);
for _ in 0..128 {
v.push(Some(val & 1 == 1));
val >>= 1;
}
v
}
_ => vec![None; 128],
};
let bits = values
.into_iter()
.enumerate()
.map(|(i, v)| {
Ok(Boolean::from(AllocatedBit::alloc_input(
&mut cs.ns(|| format!("allocated bit_gadget {}", i)),
|| v.ok_or(SynthesisError::AssignmentMissing),
)?))
})
.collect::<Result<Vec<_>, SynthesisError>>()?;
Ok(Self {
bits,
negated: false,
value: value.ok(),
})
}
}
impl<F: Field> ToBytesGadget<F> for UInt128 {
#[inline]
fn to_bytes<CS: ConstraintSystem<F>>(&self, _cs: CS) -> Result<Vec<UInt8>, SynthesisError> {
let value_chunks = match self.value.map(|val| {
let mut bytes = [0u8; 16];
val.write(bytes.as_mut()).unwrap();
bytes
}) {
Some(chunks) => [Some(chunks[0]), Some(chunks[1]), Some(chunks[2]), Some(chunks[3])],
None => [None, None, None, None],
};
let mut bytes = Vec::new();
for (i, chunk8) in self.to_bits_le().chunks(8).enumerate() {
let byte = UInt8 {
bits: chunk8.to_vec(),
negated: false,
value: value_chunks[i],
};
bytes.push(byte);
}
Ok(bytes)
}
fn to_bytes_strict<CS: ConstraintSystem<F>>(&self, cs: CS) -> Result<Vec<UInt8>, SynthesisError> {
self.to_bytes(cs)
}
}