MSM optimizations

halo2_proofs/src/poly/commitment/msm.rs

@@ -1,17 +1,19 @@
 use super::Params;
-use crate::arithmetic::{best_multiexp, parallelize, CurveAffine};
+use crate::arithmetic::{best_multiexp, CurveAffine};
 use ff::Field;
 use group::Group;
 
+use std::collections::BTreeMap;
+
 /// A multiscalar multiplication in the polynomial commitment scheme
 #[derive(Debug, Clone)]
 pub struct MSM<'a, C: CurveAffine> {
     pub(crate) params: &'a Params<C>,
     g_scalars: Option<Vec<C::Scalar>>,
     w_scalar: Option<C::Scalar>,
     u_scalar: Option<C::Scalar>,
-    other_scalars: Vec<C::Scalar>,
-    other_bases: Vec<C>,
+    // x-coordinate -> (scalar, y-coordinate)
+    other: BTreeMap<C::Base, (C::Scalar, C::Base)>,
 }
 
 impl<'a, C: CurveAffine> MSM<'a, C> {
@@ -20,23 +22,32 @@ impl<'a, C: CurveAffine> MSM<'a, C> {
         let g_scalars = None;
         let w_scalar = None;
         let u_scalar = None;
-        let other_scalars = vec![];
-        let other_bases = vec![];
+        let other = BTreeMap::new();
 
         MSM {
             params,
             g_scalars,
             w_scalar,
             u_scalar,
-            other_scalars,
-            other_bases,
+            other,
         }
     }
 
     /// Add another multiexp into this one
     pub fn add_msm(&mut self, other: &Self) {
-        self.other_scalars.extend(other.other_scalars.iter());
-        self.other_bases.extend(other.other_bases.iter());
+        for (x, (scalar, y)) in other.other.iter() {
+            self.other
+                .entry(*x)
+                .and_modify(|(our_scalar, our_y)| {
+                    if our_y == y {
+                        *our_scalar += *scalar;
+                    } else {
+                        assert!(*our_y == -*y);
+                        *our_scalar -= *scalar;
+                    }
+                })
+                .or_insert((*scalar, *y));
+        }
 
         if let Some(g_scalars) = &other.g_scalars {
             self.add_to_g_scalars(g_scalars);
@@ -53,8 +64,23 @@ impl<'a, C: CurveAffine> MSM<'a, C> {
 
     /// Add arbitrary term (the scalar and the point)
     pub fn append_term(&mut self, scalar: C::Scalar, point: C) {
-        self.other_scalars.push(scalar);
-        self.other_bases.push(point);
+        if !bool::from(point.is_identity()) {
+            let xy = point.coordinates().unwrap();
+            let x = *xy.x();
+            let y = *xy.y();
+
+            self.other
+                .entry(x)
+                .and_modify(|(our_scalar, our_y)| {
+                    if *our_y == y {
+                        *our_scalar += scalar;
+                    } else {
+                        assert!(*our_y == -y);
+                        *our_scalar -= scalar;
+                    }
+                })
+                .or_insert((scalar, y));
+        }
     }
 
     /// Add a value to the first entry of `g_scalars`.
@@ -73,11 +99,9 @@ impl<'a, C: CurveAffine> MSM<'a, C> {
     pub fn add_to_g_scalars(&mut self, scalars: &[C::Scalar]) {
         assert_eq!(scalars.len(), self.params.n as usize);
         if let Some(g_scalars) = &mut self.g_scalars {
-            parallelize(g_scalars, |g_scalars, start| {
-                for (g_scalar, scalar) in g_scalars.iter_mut().zip(scalars[start..].iter()) {
-                    *g_scalar += scalar;
-                }
-            })
+            for (g_scalar, scalar) in g_scalars.iter_mut().zip(scalars.iter()) {
+                *g_scalar += scalar;
+            }
         } else {
             self.g_scalars = Some(scalars.to_vec());
         }
@@ -96,19 +120,13 @@ impl<'a, C: CurveAffine> MSM<'a, C> {
     /// Scale all scalars in the MSM by some scaling factor
     pub fn scale(&mut self, factor: C::Scalar) {
         if let Some(g_scalars) = &mut self.g_scalars {
-            parallelize(g_scalars, |g_scalars, _| {
-                for g_scalar in g_scalars {
-                    *g_scalar *= &factor;
-                }
-            })
+            for g_scalar in g_scalars {
+                *g_scalar *= &factor;
+            }
         }
 
-        if !self.other_scalars.is_empty() {
-            parallelize(&mut self.other_scalars, |other_scalars, _| {
-                for other_scalar in other_scalars {
-                    *other_scalar *= &factor;
-                }
-            })
+        for other in self.other.values_mut() {
+            other.0 *= factor;
         }
 
         self.w_scalar = self.w_scalar.map(|a| a * &factor);
@@ -120,12 +138,16 @@ impl<'a, C: CurveAffine> MSM<'a, C> {
         let len = self.g_scalars.as_ref().map(|v| v.len()).unwrap_or(0)
             + self.w_scalar.map(|_| 1).unwrap_or(0)
             + self.u_scalar.map(|_| 1).unwrap_or(0)
-            + self.other_scalars.len();
+            + self.other.len();
         let mut scalars: Vec<C::Scalar> = Vec::with_capacity(len);
         let mut bases: Vec<C> = Vec::with_capacity(len);
 
-        scalars.extend(&self.other_scalars);
-        bases.extend(&self.other_bases);
+        scalars.extend(self.other.values().map(|(scalar, _)| scalar));
+        bases.extend(
+            self.other
+                .iter()
+                .map(|(x, (_, y))| C::from_xy(*x, *y).unwrap()),
+        );
 
         if let Some(w_scalar) = self.w_scalar {
             scalars.push(w_scalar);
@@ -147,3 +169,52 @@ impl<'a, C: CurveAffine> MSM<'a, C> {
         bool::from(best_multiexp(&scalars, &bases).is_identity())
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use crate::poly::commitment::{Params, MSM};
+    use group::Curve;
+    use pasta_curves::{arithmetic::CurveAffine, EpAffine, Fp, Fq};
+
+    #[test]
+    fn msm_arithmetic() {
+        let base = EpAffine::from_xy(-Fp::one(), Fp::from(2)).unwrap();
+        let base_viol = (base + base).to_affine();
+
+        let params = Params::new(4);
+        let mut a: MSM<EpAffine> = MSM::new(&params);
+        a.append_term(Fq::one(), base);
+        // a = [1] P
+        assert!(!a.clone().eval());
+        a.append_term(Fq::one(), base);
+        // a = [1+1] P
+        assert!(!a.clone().eval());
+        a.append_term(-Fq::one(), base_viol);
+        // a = [1+1] P + [-1] 2P
+        assert!(a.clone().eval());
+        let b = a.clone();
+
+        // Append a point that is the negation of an existing one.
+        a.append_term(Fq::from(4), -base);
+        // a = [1+1-4] P + [-1] 2P
+        assert!(!a.clone().eval());
+        a.append_term(Fq::from(2), base_viol);
+        // a = [1+1-4] P + [-1+2] 2P
+        assert!(a.clone().eval());
+
+        // Add two MSMs with common bases.
+        a.scale(Fq::from(3));
+        a.add_msm(&b);
+        // a = [3*(1+1)+(1+1-4)] P + [3*(-1)+(-1+2)] 2P
+        assert!(a.clone().eval());
+
+        let mut c: MSM<EpAffine> = MSM::new(&params);
+        c.append_term(Fq::from(2), base);
+        c.append_term(Fq::one(), -base_viol);
+        // c = [2] P + [1] (-2P)
+        assert!(c.clone().eval());
+        // Add two MSMs with bases that differ only in sign.
+        a.add_msm(&c);
+        assert!(a.eval());
+    }
+}