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