1use bellperson::{
  2  gadgets::{
  3    boolean::{AllocatedBit, Boolean},
  4    num::AllocatedNum,
  5    Assignment,
  6  },
  7  ConstraintSystem, LinearCombination, SynthesisError,
  8};
  9use bellperson_nonnative::mp::bignat::{nat_to_limbs, BigNat};
 10use ff::{PrimeField, PrimeFieldBits};
 11use rug::Integer;
 12
 13/// Gets as input the little indian representation of a number and spits out the number
 14#[allow(dead_code)]
 15pub fn le_bits_to_num<Scalar, CS>(
 16  mut cs: CS,
 17  bits: Vec<AllocatedBit>,
 18) -> Result<AllocatedNum<Scalar>, SynthesisError>
 19where
 20  Scalar: PrimeField + PrimeFieldBits,
 21  CS: ConstraintSystem<Scalar>,
 22{
 23  // We loop over the input bits and construct the constraint
 24  // and the field element that corresponds to the result
 25  let mut lc = LinearCombination::zero();
 26  let mut coeff = Scalar::one();
 27  let mut fe = Some(Scalar::zero());
 28  for bit in bits.iter() {
 29    lc = lc + (coeff, bit.get_variable());
 30    fe = bit.get_value().map(|val| {
 31      if val {
 32        fe.unwrap() + coeff
 33      } else {
 34        fe.unwrap()
 35      }
 36    });
 37    coeff = coeff.double();
 38  }
 39  let num = AllocatedNum::alloc(cs.namespace(|| "Field element"), || {
 40    fe.ok_or(SynthesisError::AssignmentMissing)
 41  })?;
 42  lc = lc - num.get_variable();
 43  cs.enforce(|| "compute number from bits", |lc| lc, |lc| lc, |_| lc);
 44  Ok(num)
 45}
 46
 47/// Allocate a variable that is set to zero
 48pub fn alloc_zero<F: PrimeField, CS: ConstraintSystem<F>>(
 49  mut cs: CS,
 50) -> Result<AllocatedNum<F>, SynthesisError> {
 51  let zero = AllocatedNum::alloc(cs.namespace(|| "alloc"), || Ok(F::zero()))?;
 52  cs.enforce(
 53    || "check zero is valid",
 54    |lc| lc,
 55    |lc| lc,
 56    |lc| lc + zero.get_variable(),
 57  );
 58  Ok(zero)
 59}
 60
 61/// Allocate a variable that is set to one
 62pub fn alloc_one<F: PrimeField, CS: ConstraintSystem<F>>(
 63  mut cs: CS,
 64) -> Result<AllocatedNum<F>, SynthesisError> {
 65  let one = AllocatedNum::alloc(cs.namespace(|| "alloc"), || Ok(F::one()))?;
 66  cs.enforce(
 67    || "check one is valid",
 68    |lc| lc + CS::one(),
 69    |lc| lc + CS::one(),
 70    |lc| lc + one.get_variable(),
 71  );
 72
 73  Ok(one)
 74}
 75
 76/// Allocate bignat a constant
 77pub fn alloc_bignat_constant<F: PrimeField, CS: ConstraintSystem<F>>(
 78  mut cs: CS,
 79  val: &Integer,
 80  limb_width: usize,
 81  n_limbs: usize,
 82) -> Result<BigNat<F>, SynthesisError> {
 83  let limbs = nat_to_limbs(val, limb_width, n_limbs).unwrap();
 84  let bignat = BigNat::alloc_from_limbs(
 85    cs.namespace(|| "alloc bignat"),
 86    || Ok(limbs.clone()),
 87    None,
 88    limb_width,
 89    n_limbs,
 90  )?;
 91  // Now enforce that the limbs are all equal to the constants
 92  #[allow(clippy::needless_range_loop)]
 93  for i in 0..n_limbs {
 94    cs.enforce(
 95      || format!("check limb {}", i),
 96      |lc| lc + &bignat.limbs[i],
 97      |lc| lc + CS::one(),
 98      |lc| lc + (limbs[i], CS::one()),
 99    );
100  }
101  Ok(bignat)
102}
103
104/// Check that two numbers are equal and return a bit
105pub fn alloc_num_equals<F: PrimeField, CS: ConstraintSystem<F>>(
106  mut cs: CS,
107  a: AllocatedNum<F>,
108  b: AllocatedNum<F>,
109) -> Result<AllocatedBit, SynthesisError> {
110  // Allocate and constrain `r`: result boolean bit.
111  // It equals `true` if `a` equals `b`, `false` otherwise
112
113  let r_value = match (a.get_value(), b.get_value()) {
114    (Some(a), Some(b)) => Some(a == b),
115    _ => None,
116  };
117
118  let r = AllocatedBit::alloc(cs.namespace(|| "r"), r_value)?;
119
120  let delta = AllocatedNum::alloc(cs.namespace(|| "delta"), || {
121    let a_value = *a.get_value().get()?;
122    let b_value = *b.get_value().get()?;
123
124    let mut delta = a_value;
125    delta.sub_assign(&b_value);
126
127    Ok(delta)
128  })?;
129
130  cs.enforce(
131    || "delta = (a - b)",
132    |lc| lc + a.get_variable() - b.get_variable(),
133    |lc| lc + CS::one(),
134    |lc| lc + delta.get_variable(),
135  );
136
137  let delta_inv = AllocatedNum::alloc(cs.namespace(|| "delta_inv"), || {
138    let delta = *delta.get_value().get()?;
139
140    if delta.is_zero().unwrap_u8() == 1 {
141      Ok(F::one()) // we can return any number here, it doesn't matter
142    } else {
143      Ok(delta.invert().unwrap())
144    }
145  })?;
146
147  // Allocate `t = delta * delta_inv`
148  // If `delta` is non-zero (a != b), `t` will equal 1
149  // If `delta` is zero (a == b), `t` cannot equal 1
150
151  let t = AllocatedNum::alloc(cs.namespace(|| "t"), || {
152    let mut tmp = *delta.get_value().get()?;
153    tmp.mul_assign(&(*delta_inv.get_value().get()?));
154
155    Ok(tmp)
156  })?;
157
158  // Constrain allocation:
159  // t = (a - b) * delta_inv
160  cs.enforce(
161    || "t = (a - b) * delta_inv",
162    |lc| lc + a.get_variable() - b.get_variable(),
163    |lc| lc + delta_inv.get_variable(),
164    |lc| lc + t.get_variable(),
165  );
166
167  // Constrain:
168  // (a - b) * (t - 1) == 0
169  // This enforces that correct `delta_inv` was provided,
170  // and thus `t` is 1 if `(a - b)` is non zero (a != b )
171  cs.enforce(
172    || "(a - b) * (t - 1) == 0",
173    |lc| lc + a.get_variable() - b.get_variable(),
174    |lc| lc + t.get_variable() - CS::one(),
175    |lc| lc,
176  );
177
178  // Constrain:
179  // (a - b) * r == 0
180  // This enforces that `r` is zero if `(a - b)` is non-zero (a != b)
181  cs.enforce(
182    || "(a - b) * r == 0",
183    |lc| lc + a.get_variable() - b.get_variable(),
184    |lc| lc + r.get_variable(),
185    |lc| lc,
186  );
187
188  // Constrain:
189  // (t - 1) * (r - 1) == 0
190  // This enforces that `r` is one if `t` is not one (a == b)
191  cs.enforce(
192    || "(t - 1) * (r - 1) == 0",
193    |lc| lc + t.get_variable() - CS::one(),
194    |lc| lc + r.get_variable() - CS::one(),
195    |lc| lc,
196  );
197
198  Ok(r)
199}
200
201/// If condition return a otherwise b
202pub fn conditionally_select<F: PrimeField, CS: ConstraintSystem<F>>(
203  mut cs: CS,
204  a: &AllocatedNum<F>,
205  b: &AllocatedNum<F>,
206  condition: &Boolean,
207) -> Result<AllocatedNum<F>, SynthesisError> {
208  let c = AllocatedNum::alloc(cs.namespace(|| "conditional select result"), || {
209    if *condition.get_value().get()? {
210      Ok(*a.get_value().get()?)
211    } else {
212      Ok(*b.get_value().get()?)
213    }
214  })?;
215
216  // a * condition + b*(1-condition) = c ->
217  // a * condition - b*condition = c - b
218
219  cs.enforce(
220    || "conditional select constraint",
221    |lc| lc + a.get_variable() - b.get_variable(),
222    |_| condition.lc(CS::one(), F::one()),
223    |lc| lc + c.get_variable() - b.get_variable(),
224  );
225
226  Ok(c)
227}
228
229/// If condition return a otherwise b where a and b are BigNats
230pub fn conditionally_select_bignat<F: PrimeField, CS: ConstraintSystem<F>>(
231  mut cs: CS,
232  a: &BigNat<F>,
233  b: &BigNat<F>,
234  condition: &Boolean,
235) -> Result<BigNat<F>, SynthesisError> {
236  assert!(a.limbs.len() == b.limbs.len());
237  let c = BigNat::alloc_from_nat(
238    cs.namespace(|| "conditional select result"),
239    || {
240      if *condition.get_value().get()? {
241        Ok(a.value.get()?.clone())
242      } else {
243        Ok(b.value.get()?.clone())
244      }
245    },
246    a.params.limb_width,
247    a.params.n_limbs,
248  )?;
249
250  // a * condition + b*(1-condition) = c ->
251  // a * condition - b*condition = c - b
252  for i in 0..c.limbs.len() {
253    cs.enforce(
254      || format!("conditional select constraint {}", i),
255      |lc| lc + &a.limbs[i] - &b.limbs[i],
256      |_| condition.lc(CS::one(), F::one()),
257      |lc| lc + &c.limbs[i] - &b.limbs[i],
258    );
259  }
260  Ok(c)
261}
262
263/// Same as the above but Condition is an AllocatedNum that needs to be
264/// 0 or 1. 1 => True, 0 => False
265pub fn conditionally_select2<F: PrimeField, CS: ConstraintSystem<F>>(
266  mut cs: CS,
267  a: &AllocatedNum<F>,
268  b: &AllocatedNum<F>,
269  condition: &AllocatedNum<F>,
270) -> Result<AllocatedNum<F>, SynthesisError> {
271  let c = AllocatedNum::alloc(cs.namespace(|| "conditional select result"), || {
272    if *condition.get_value().get()? == F::one() {
273      Ok(*a.get_value().get()?)
274    } else {
275      Ok(*b.get_value().get()?)
276    }
277  })?;
278
279  // a * condition + b*(1-condition) = c ->
280  // a * condition - b*condition = c - b
281
282  cs.enforce(
283    || "conditional select constraint",
284    |lc| lc + a.get_variable() - b.get_variable(),
285    |lc| lc + condition.get_variable(),
286    |lc| lc + c.get_variable() - b.get_variable(),
287  );
288
289  Ok(c)
290}
291
292/// If condition set to 0 otherwise a
293pub fn select_zero_or<F: PrimeField, CS: ConstraintSystem<F>>(
294  mut cs: CS,
295  a: &AllocatedNum<F>,
296  condition: &AllocatedNum<F>,
297) -> Result<AllocatedNum<F>, SynthesisError> {
298  let c = AllocatedNum::alloc(cs.namespace(|| "conditional select result"), || {
299    if *condition.get_value().get()? == F::one() {
300      Ok(F::zero())
301    } else {
302      Ok(*a.get_value().get()?)
303    }
304  })?;
305
306  // a * (1 - condition) = c
307  cs.enforce(
308    || "conditional select constraint",
309    |lc| lc + a.get_variable(),
310    |lc| lc + CS::one() - condition.get_variable(),
311    |lc| lc + c.get_variable(),
312  );
313
314  Ok(c)
315}
316
317/// If condition set to 1 otherwise a
318pub fn select_one_or<F: PrimeField, CS: ConstraintSystem<F>>(
319  mut cs: CS,
320  a: &AllocatedNum<F>,
321  condition: &AllocatedNum<F>,
322) -> Result<AllocatedNum<F>, SynthesisError> {
323  let c = AllocatedNum::alloc(cs.namespace(|| "conditional select result"), || {
324    if *condition.get_value().get()? == F::one() {
325      Ok(F::one())
326    } else {
327      Ok(*a.get_value().get()?)
328    }
329  })?;
330
331  cs.enforce(
332    || "conditional select constraint",
333    |lc| lc + CS::one() - a.get_variable(),
334    |lc| lc + condition.get_variable(),
335    |lc| lc + c.get_variable() - a.get_variable(),
336  );
337  Ok(c)
338}