ark_relations/r1cs/
impl_lc.rs

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/// Generate a `LinearCombination` from arithmetic expressions involving
12/// `Variable`s.
13#[macro_export]
14macro_rules! lc {
15    () => {
16        $crate::r1cs::LinearCombination::zero()
17    };
18}
19
20impl<F: Field> LinearCombination<F> {
21    /// Create a new empty linear combination.
22    pub fn new() -> Self {
23        Default::default()
24    }
25
26    /// Create a new empty linear combination.
27    pub fn zero() -> Self {
28        Self::new()
29    }
30
31    /// Deduplicate entries in `self`.
32    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    /// Negate the coefficients of all variables in `self`.
81    #[inline]
82    pub fn negate_in_place(&mut self) {
83        self.0.iter_mut().for_each(|(coeff, _)| *coeff = -(*coeff));
84    }
85
86    /// Get the location of a variable in `self`.
87    #[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}