Add `{mul,add,sub}_extended` and `rotate_extended` methods to `EvaluationDomain`.

src/poly/domain.rs

@@ -247,7 +247,7 @@ impl<G: Group> EvaluationDomain<G> {
         if rotation.0 == 0 {
             // In this special case, the powers of zeta repeat so we do not need
             // to compute them.
-            Self::distribute_powers_zeta(&mut a.values);
+            Self::distribute_powers_zeta(&mut a.values, true);
         } else {
             let mut g = self.g_coset;
             if rotation.0 > 0 {
@@ -268,6 +268,102 @@ impl<G: Group> EvaluationDomain<G> {
         }
     }
 
+    fn op_extended(
+        &self,
+        mut left: Polynomial<G, ExtendedLagrangeCoeff>,
+        right: &Polynomial<G, ExtendedLagrangeCoeff>,
+        rotation: Rotation,
+        op: impl Fn(&mut G, &G) + Send + Sync + 'static,
+    ) -> Polynomial<G, ExtendedLagrangeCoeff>
+    where
+        G: Field,
+    {
+        // Rotation while in the extended domain is simply exaggerated by how
+        // much larger the extended domain is compared to the original domain.
+        let rotation = (1 << (self.extended_k - self.k)) * rotation.0;
+
+        parallelize(&mut left.values, |lhs, start| {
+            let start = ((((start + self.extended_len()) as i32) + rotation) as usize)
+                % self.extended_len();
+
+            for (lhs, rhs) in lhs
+                .iter_mut()
+                .zip(right.values[start..].iter().chain(right.values.iter()))
+            {
+                op(lhs, rhs);
+            }
+        });
+
+        left
+    }
+
+    /// Multiply two polynomials in the extended domain, rotating the latter
+    /// polynomial over the original domain first.
+    pub fn mul_extended(
+        &self,
+        left: Polynomial<G, ExtendedLagrangeCoeff>,
+        right: &Polynomial<G, ExtendedLagrangeCoeff>,
+        rotation: Rotation,
+    ) -> Polynomial<G, ExtendedLagrangeCoeff>
+    where
+        G: Field,
+    {
+        self.op_extended(left, right, rotation, |lhs, rhs| {
+            *lhs *= *rhs;
+        })
+    }
+
+    /// Add two polynomials in the extended domain, rotating the latter
+    /// polynomial over the original domain first.
+    pub fn add_extended(
+        &self,
+        left: Polynomial<G, ExtendedLagrangeCoeff>,
+        right: &Polynomial<G, ExtendedLagrangeCoeff>,
+        rotation: Rotation,
+    ) -> Polynomial<G, ExtendedLagrangeCoeff>
+    where
+        G: Field,
+    {
+        self.op_extended(left, right, rotation, |lhs, rhs| {
+            *lhs += *rhs;
+        })
+    }
+
+    /// Subtract a polynomial from another in the extended domain, rotating the
+    /// former polynomial over the original domain first.
+    pub fn sub_extended(
+        &self,
+        left: Polynomial<G, ExtendedLagrangeCoeff>,
+        right: &Polynomial<G, ExtendedLagrangeCoeff>,
+        rotation: Rotation,
+    ) -> Polynomial<G, ExtendedLagrangeCoeff>
+    where
+        G: Field,
+    {
+        self.op_extended(left, right, rotation, |lhs, rhs| {
+            *lhs -= *rhs;
+        })
+    }
+
+    /// Rotate the extended domain polynomial over the original domain.
+    pub fn rotate_extended(
+        &self,
+        poly: &Polynomial<G, ExtendedLagrangeCoeff>,
+        rotation: Rotation,
+    ) -> Polynomial<G, ExtendedLagrangeCoeff> {
+        let new_rotation = ((1 << (self.extended_k - self.k)) * rotation.0.abs()) as usize;
+
+        let mut poly = poly.clone();
+
+        if rotation.0 >= 0 {
+            poly.values.rotate_left(new_rotation);
+        } else {
+            poly.values.rotate_right(new_rotation);
+        }
+
+        poly
+    }
+
     /// This takes us from the extended evaluation domain and gets us the
     /// quotient polynomial coefficients.
     ///
@@ -287,7 +383,7 @@ impl<G: Group> EvaluationDomain<G> {
 
         // Distribute powers to move from coset; opposite from the
         // transformation we performed earlier.
-        Self::distribute_powers(&mut a.values, self.g_coset_inv);
+        Self::distribute_powers_zeta(&mut a.values, false);
 
         // Truncate it to match the size of the quotient polynomial; the
         // evaluation domain might be slightly larger than necessary because
@@ -325,11 +421,15 @@ impl<G: Group> EvaluationDomain<G> {
     // `[a_0, [zeta]a_1, [zeta^2]a_2, a_3, [zeta]a_4, [zeta^2]a_5, a_6, ...]`,
     // where zeta is a cube root of unity in the multiplicative subgroup with
     // order (p - 1), i.e. zeta^3 = 1.
-    fn distribute_powers_zeta(mut a: &mut [G]) {
+    fn distribute_powers_zeta(mut a: &mut [G], direction: bool) {
-        let coset_powers = [G::Scalar::ZETA, G::Scalar::ZETA.square()];
+        let coset_powers = if direction {
+            [G::Scalar::ZETA, G::Scalar::ZETA.square()]
+        } else {
+            [G::Scalar::ZETA.square(), G::Scalar::ZETA]
+        };
         parallelize(&mut a, |a, mut index| {
             for a in a {
-                // Distribute powers to move into coset
+                // Distribute powers to move into/from coset
                 let i = index % (coset_powers.len() + 1);
                 if i != 0 {
                     a.group_scale(&coset_powers[i - 1]);