1#![allow(clippy::suspicious_arithmetic_impl)]
2
3use crate::r1cs::{LinearCombination, Variable};
4use ark_ff::Field;
5use ark_std::{
6 ops::{Add, AddAssign, Deref, DerefMut, Mul, MulAssign, Neg, Sub},
7 vec,
8 vec::Vec,
9};
10
11#[macro_export]
14macro_rules! lc {
15 () => {
16 $crate::r1cs::LinearCombination::zero()
17 };
18}
19
20impl<F: Field> LinearCombination<F> {
21 pub fn new() -> Self {
23 Default::default()
24 }
25
26 pub fn zero() -> Self {
28 Self::new()
29 }
30
31 pub fn compactify(&mut self) {
33 self.0.sort_by_key(|e| e.1);
34 let mut current_var = None;
35 let mut current_var_first_index = 0;
36 for i in 0..self.0.len() {
37 let (f, v) = self.0[i];
38 if Some(v) == current_var {
39 self.0[current_var_first_index].0 += &f;
40 } else {
41 current_var = Some(v);
42 current_var_first_index = i;
43 }
44 }
45 self.0.dedup_by_key(|e| e.1);
46 }
47}
48
49impl<'a, F: Field> Deref for LinearCombination<F> {
50 type Target = Vec<(F, Variable)>;
51
52 #[inline]
53 fn deref(&self) -> &Vec<(F, Variable)> {
54 &self.0
55 }
56}
57
58impl<F: Field> DerefMut for LinearCombination<F> {
59 #[inline]
60 fn deref_mut(&mut self) -> &mut Self::Target {
61 &mut self.0
62 }
63}
64
65impl<F: Field> From<(F, Variable)> for LinearCombination<F> {
66 #[inline]
67 fn from(input: (F, Variable)) -> Self {
68 LinearCombination(vec![input])
69 }
70}
71
72impl<F: Field> From<Variable> for LinearCombination<F> {
73 #[inline]
74 fn from(var: Variable) -> Self {
75 LinearCombination(vec![(F::one(), var)])
76 }
77}
78
79impl<F: Field> LinearCombination<F> {
80 #[inline]
82 pub fn negate_in_place(&mut self) {
83 self.0.iter_mut().for_each(|(coeff, _)| *coeff = -(*coeff));
84 }
85
86 #[inline]
88 pub fn get_var_loc(&self, search_var: &Variable) -> Result<usize, usize> {
89 if self.0.len() < 6 {
90 let mut found_index = 0;
91 for (i, (_, var)) in self.iter().enumerate() {
92 if var >= search_var {
93 found_index = i;
94 break;
95 } else {
96 found_index += 1;
97 }
98 }
99 Err(found_index)
100 } else {
101 self.0
102 .binary_search_by_key(search_var, |&(_, cur_var)| cur_var)
103 }
104 }
105}
106
107impl<F: Field> Add<(F, Variable)> for LinearCombination<F> {
108 type Output = Self;
109
110 #[inline]
111 fn add(mut self, coeff_var: (F, Variable)) -> Self {
112 self += coeff_var;
113 self
114 }
115}
116
117impl<F: Field> AddAssign<(F, Variable)> for LinearCombination<F> {
118 #[inline]
119 fn add_assign(&mut self, (coeff, var): (F, Variable)) {
120 match self.get_var_loc(&var) {
121 Ok(found) => self.0[found].0 += &coeff,
122 Err(not_found) => self.0.insert(not_found, (coeff, var)),
123 }
124 }
125}
126
127impl<F: Field> Sub<(F, Variable)> for LinearCombination<F> {
128 type Output = Self;
129
130 #[inline]
131 fn sub(self, (coeff, var): (F, Variable)) -> Self {
132 self + (-coeff, var)
133 }
134}
135
136impl<F: Field> Neg for LinearCombination<F> {
137 type Output = Self;
138
139 #[inline]
140 fn neg(mut self) -> Self {
141 self.negate_in_place();
142 self
143 }
144}
145
146impl<F: Field> Mul<F> for LinearCombination<F> {
147 type Output = Self;
148
149 #[inline]
150 fn mul(mut self, scalar: F) -> Self {
151 self *= scalar;
152 self
153 }
154}
155
156impl<'a, F: Field> Mul<F> for &'a LinearCombination<F> {
157 type Output = LinearCombination<F>;
158
159 #[inline]
160 fn mul(self, scalar: F) -> LinearCombination<F> {
161 let mut cur = self.clone();
162 cur *= scalar;
163 cur
164 }
165}
166
167impl<F: Field> MulAssign<F> for LinearCombination<F> {
168 #[inline]
169 fn mul_assign(&mut self, scalar: F) {
170 self.0.iter_mut().for_each(|(coeff, _)| *coeff *= &scalar);
171 }
172}
173
174impl<F: Field> Add<Variable> for LinearCombination<F> {
175 type Output = Self;
176
177 #[inline]
178 fn add(self, other: Variable) -> LinearCombination<F> {
179 self + (F::one(), other)
180 }
181}
182
183impl<'a, F: Field> Add<&'a Variable> for LinearCombination<F> {
184 type Output = Self;
185
186 #[inline]
187 fn add(self, other: &'a Variable) -> LinearCombination<F> {
188 self + *other
189 }
190}
191
192impl<'a, F: Field> Sub<&'a Variable> for LinearCombination<F> {
193 type Output = Self;
194
195 #[inline]
196 fn sub(self, other: &'a Variable) -> LinearCombination<F> {
197 self - *other
198 }
199}
200
201impl<F: Field> Sub<Variable> for LinearCombination<F> {
202 type Output = LinearCombination<F>;
203
204 #[inline]
205 fn sub(self, other: Variable) -> LinearCombination<F> {
206 self - (F::one(), other)
207 }
208}
209
210fn op_impl<F: Field, F1, F2>(
211 cur: &LinearCombination<F>,
212 other: &LinearCombination<F>,
213 push_fn: F1,
214 combine_fn: F2,
215) -> LinearCombination<F>
216where
217 F1: Fn(F) -> F,
218 F2: Fn(F, F) -> F,
219{
220 let mut new_vec = Vec::new();
221 let mut i = 0;
222 let mut j = 0;
223 while i < cur.len() && j < other.len() {
224 let self_cur = &cur[i];
225 let other_cur = &other[j];
226 use core::cmp::Ordering;
227 match self_cur.1.cmp(&other_cur.1) {
228 Ordering::Greater => {
229 new_vec.push((push_fn(other[j].0), other[j].1));
230 j += 1;
231 },
232 Ordering::Less => {
233 new_vec.push(*self_cur);
234 i += 1;
235 },
236 Ordering::Equal => {
237 new_vec.push((combine_fn(self_cur.0, other_cur.0), self_cur.1));
238 i += 1;
239 j += 1;
240 },
241 };
242 }
243 new_vec.extend_from_slice(&cur[i..]);
244 while j < other.0.len() {
245 new_vec.push((push_fn(other[j].0), other[j].1));
246 j += 1;
247 }
248 LinearCombination(new_vec)
249}
250
251impl<F: Field> Add<&LinearCombination<F>> for &LinearCombination<F> {
252 type Output = LinearCombination<F>;
253
254 fn add(self, other: &LinearCombination<F>) -> LinearCombination<F> {
255 if other.0.is_empty() {
256 return self.clone();
257 } else if self.0.is_empty() {
258 return other.clone();
259 }
260 op_impl(
261 self,
262 other,
263 |coeff| coeff,
264 |cur_coeff, other_coeff| cur_coeff + other_coeff,
265 )
266 }
267}
268
269impl<F: Field> Add<LinearCombination<F>> for &LinearCombination<F> {
270 type Output = LinearCombination<F>;
271
272 fn add(self, other: LinearCombination<F>) -> LinearCombination<F> {
273 if self.0.is_empty() {
274 return other;
275 } else if other.0.is_empty() {
276 return self.clone();
277 }
278 op_impl(
279 self,
280 &other,
281 |coeff| coeff,
282 |cur_coeff, other_coeff| cur_coeff + other_coeff,
283 )
284 }
285}
286
287impl<'a, F: Field> Add<&'a LinearCombination<F>> for LinearCombination<F> {
288 type Output = LinearCombination<F>;
289
290 fn add(self, other: &'a LinearCombination<F>) -> LinearCombination<F> {
291 if other.0.is_empty() {
292 return self;
293 } else if self.0.is_empty() {
294 return other.clone();
295 }
296 op_impl(
297 &self,
298 other,
299 |coeff| coeff,
300 |cur_coeff, other_coeff| cur_coeff + other_coeff,
301 )
302 }
303}
304
305impl<F: Field> Add<LinearCombination<F>> for LinearCombination<F> {
306 type Output = Self;
307
308 fn add(self, other: Self) -> Self {
309 if other.0.is_empty() {
310 return self;
311 } else if self.0.is_empty() {
312 return other;
313 }
314 op_impl(
315 &self,
316 &other,
317 |coeff| coeff,
318 |cur_coeff, other_coeff| cur_coeff + other_coeff,
319 )
320 }
321}
322
323impl<F: Field> Sub<&LinearCombination<F>> for &LinearCombination<F> {
324 type Output = LinearCombination<F>;
325
326 fn sub(self, other: &LinearCombination<F>) -> LinearCombination<F> {
327 if other.0.is_empty() {
328 let cur = self.clone();
329 return cur;
330 } else if self.0.is_empty() {
331 let mut other = other.clone();
332 other.negate_in_place();
333 return other;
334 }
335
336 op_impl(
337 self,
338 other,
339 |coeff| -coeff,
340 |cur_coeff, other_coeff| cur_coeff - other_coeff,
341 )
342 }
343}
344
345impl<'a, F: Field> Sub<&'a LinearCombination<F>> for LinearCombination<F> {
346 type Output = LinearCombination<F>;
347
348 fn sub(self, other: &'a LinearCombination<F>) -> LinearCombination<F> {
349 if other.0.is_empty() {
350 return self;
351 } else if self.0.is_empty() {
352 let mut other = other.clone();
353 other.negate_in_place();
354 return other;
355 }
356 op_impl(
357 &self,
358 other,
359 |coeff| -coeff,
360 |cur_coeff, other_coeff| cur_coeff - other_coeff,
361 )
362 }
363}
364
365impl<F: Field> Sub<LinearCombination<F>> for &LinearCombination<F> {
366 type Output = LinearCombination<F>;
367
368 fn sub(self, mut other: LinearCombination<F>) -> LinearCombination<F> {
369 if self.0.is_empty() {
370 other.negate_in_place();
371 return other;
372 } else if other.0.is_empty() {
373 return self.clone();
374 }
375
376 op_impl(
377 self,
378 &other,
379 |coeff| -coeff,
380 |cur_coeff, other_coeff| cur_coeff - other_coeff,
381 )
382 }
383}
384
385impl<F: Field> Sub<LinearCombination<F>> for LinearCombination<F> {
386 type Output = LinearCombination<F>;
387
388 fn sub(self, mut other: LinearCombination<F>) -> LinearCombination<F> {
389 if other.0.is_empty() {
390 return self;
391 } else if self.0.is_empty() {
392 other.negate_in_place();
393 return other;
394 }
395 op_impl(
396 &self,
397 &other,
398 |coeff| -coeff,
399 |cur_coeff, other_coeff| cur_coeff - other_coeff,
400 )
401 }
402}
403
404impl<F: Field> Add<(F, &LinearCombination<F>)> for &LinearCombination<F> {
405 type Output = LinearCombination<F>;
406
407 fn add(self, (mul_coeff, other): (F, &LinearCombination<F>)) -> LinearCombination<F> {
408 if other.0.is_empty() {
409 return self.clone();
410 } else if self.0.is_empty() {
411 let mut other = other.clone();
412 other.mul_assign(mul_coeff);
413 return other;
414 }
415 op_impl(
416 self,
417 other,
418 |coeff| mul_coeff * coeff,
419 |cur_coeff, other_coeff| cur_coeff + mul_coeff * other_coeff,
420 )
421 }
422}
423
424impl<'a, F: Field> Add<(F, &'a LinearCombination<F>)> for LinearCombination<F> {
425 type Output = LinearCombination<F>;
426
427 fn add(self, (mul_coeff, other): (F, &'a LinearCombination<F>)) -> LinearCombination<F> {
428 if other.0.is_empty() {
429 return self;
430 } else if self.0.is_empty() {
431 let mut other = other.clone();
432 other.mul_assign(mul_coeff);
433 return other;
434 }
435 op_impl(
436 &self,
437 other,
438 |coeff| mul_coeff * coeff,
439 |cur_coeff, other_coeff| cur_coeff + mul_coeff * other_coeff,
440 )
441 }
442}
443
444impl<F: Field> Add<(F, LinearCombination<F>)> for &LinearCombination<F> {
445 type Output = LinearCombination<F>;
446
447 fn add(self, (mul_coeff, mut other): (F, LinearCombination<F>)) -> LinearCombination<F> {
448 if other.0.is_empty() {
449 return self.clone();
450 } else if self.0.is_empty() {
451 other.mul_assign(mul_coeff);
452 return other;
453 }
454 op_impl(
455 self,
456 &other,
457 |coeff| mul_coeff * coeff,
458 |cur_coeff, other_coeff| cur_coeff + mul_coeff * other_coeff,
459 )
460 }
461}
462
463impl<F: Field> Add<(F, Self)> for LinearCombination<F> {
464 type Output = Self;
465
466 fn add(self, (mul_coeff, other): (F, Self)) -> Self {
467 if other.0.is_empty() {
468 return self;
469 } else if self.0.is_empty() {
470 let mut other = other;
471 other.mul_assign(mul_coeff);
472 return other;
473 }
474 op_impl(
475 &self,
476 &other,
477 |coeff| mul_coeff * coeff,
478 |cur_coeff, other_coeff| cur_coeff + mul_coeff * other_coeff,
479 )
480 }
481}
482
483impl<F: Field> Sub<(F, &LinearCombination<F>)> for &LinearCombination<F> {
484 type Output = LinearCombination<F>;
485
486 fn sub(self, (coeff, other): (F, &LinearCombination<F>)) -> LinearCombination<F> {
487 self + (-coeff, other)
488 }
489}
490
491impl<'a, F: Field> Sub<(F, &'a LinearCombination<F>)> for LinearCombination<F> {
492 type Output = LinearCombination<F>;
493
494 fn sub(self, (coeff, other): (F, &'a LinearCombination<F>)) -> LinearCombination<F> {
495 self + (-coeff, other)
496 }
497}
498
499impl<F: Field> Sub<(F, LinearCombination<F>)> for &LinearCombination<F> {
500 type Output = LinearCombination<F>;
501
502 fn sub(self, (coeff, other): (F, LinearCombination<F>)) -> LinearCombination<F> {
503 self + (-coeff, other)
504 }
505}
506
507impl<'a, F: Field> Sub<(F, LinearCombination<F>)> for LinearCombination<F> {
508 type Output = LinearCombination<F>;
509
510 fn sub(self, (coeff, other): (F, LinearCombination<F>)) -> LinearCombination<F> {
511 self + (-coeff, other)
512 }
513}