Generalised version of improved fp2

ff/src/fields/mod.rs

@@ -339,6 +339,33 @@ pub trait PrimeField:
     type Params: FpParameters<BigInt = Self::BigInt>;
     type BigInt: BigInteger;
 
+    #[inline(always)]
+    fn mul_by_i8(mut self, i: i8) -> Self {
+        self.mul_assign_by_i8(i);
+        self
+    }
+
+    #[inline(always)]
+    #[ark_ff_asm::unroll_for_loops]
+    fn mul_assign_by_i8(&mut self, i: i8) {
+        if i < 0 {
+            *self = self.neg();
+        }
+        let i = i.abs();
+
+        let tmp = *self;
+
+        for j in 1..7 {
+            if i >= 1 << (7 - j) {
+                self.double_in_place();
+                // if the next bit is 1
+                if (i >> (6 - j)) % 2 == 1 {
+                    *self += tmp;
+                }
+            }
+        }
+    }
+
     /// Returns a prime field element from its underlying representation.
     fn from_repr(repr: Self::BigInt) -> Option<Self>;
 

ff/src/fields/models/fp2.rs

@@ -1,12 +1,15 @@
 use super::quadratic_extension::*;
 use crate::fields::PrimeField;
 use core::marker::PhantomData;
+use num_traits::Zero;
 
 pub trait Fp2Parameters: 'static + Send + Sync {
     type Fp: PrimeField;
 
     const NONRESIDUE: Self::Fp;
 
+    const NONRESIDUE_SMALL: Option<i8> = None;
+
     const QUADRATIC_NONRESIDUE: (Self::Fp, Self::Fp);
 
     /// Coefficients for the Frobenius automorphism.
@@ -15,33 +18,11 @@ pub trait Fp2Parameters: 'static + Send + Sync {
     /// Return `fe * Self::NONRESIDUE`.
     #[inline(always)]
     fn mul_fp_by_nonresidue(fe: &Self::Fp) -> Self::Fp {
+        if let Some(i) = Self::NONRESIDUE_SMALL {
+            return fe.mul_by_i8(i);
+        }
         Self::NONRESIDUE * fe
     }
-
-    /// A specializable method for computing `x + mul_base_field_by_nonresidue(y)`
-    /// This allows for optimizations when the non-residue is
-    /// canonically negative in the field.
-    #[inline(always)]
-    fn add_and_mul_fp_by_nonresidue(x: &Self::Fp, y: &Self::Fp) -> Self::Fp {
-        *x + Self::mul_fp_by_nonresidue(y)
-    }
-
-    /// A specializable method for computing `x + y + mul_base_field_by_nonresidue(y)`
-    /// This allows for optimizations when the non-residue is not `-1`.
-    #[inline(always)]
-    fn add_and_mul_fp_by_nonresidue_plus_one(x: &Self::Fp, y: &Self::Fp) -> Self::Fp {
-        let mut tmp = *x;
-        tmp += y;
-        Self::add_and_mul_fp_by_nonresidue(&tmp, &y)
-    }
-
-    /// A specializable method for computing `x - mul_base_field_by_nonresidue(y)`
-    /// This allows for optimizations when the non-residue is
-    /// canonically negative in the field.
-    #[inline(always)]
-    fn sub_and_mul_fp_by_nonresidue(x: &Self::Fp, y: &Self::Fp) -> Self::Fp {
-        *x - Self::mul_fp_by_nonresidue(y)
-    }
 }
 
 pub struct Fp2ParamsWrapper<P: Fp2Parameters>(PhantomData<P>);
@@ -55,39 +36,46 @@ impl<P: Fp2Parameters> QuadExtParameters for Fp2ParamsWrapper<P> {
 
     const NONRESIDUE: Self::BaseField = P::NONRESIDUE;
 
-    const FROBENIUS_COEFF_C1: &'static [Self::FrobCoeff] = P::FROBENIUS_COEFF_FP2_C1;
-
-    #[inline(always)]
-    fn mul_base_field_by_nonresidue(fe: &Self::BaseField) -> Self::BaseField {
-        P::mul_fp_by_nonresidue(fe)
-    }
+    const NONRESIDUE_SMALL: Option<i8> = P::NONRESIDUE_SMALL;
 
-    #[inline(always)]
-    fn add_and_mul_base_field_by_nonresidue(
-        x: &Self::BaseField,
-        y: &Self::BaseField,
-    ) -> Self::BaseField {
-        P::add_and_mul_fp_by_nonresidue(x, y)
-    }
+    const FROBENIUS_COEFF_C1: &'static [Self::FrobCoeff] = P::FROBENIUS_COEFF_FP2_C1;
 
-    #[inline(always)]
-    fn add_and_mul_base_field_by_nonresidue_plus_one(
-        x: &Self::BaseField,
-        y: &Self::BaseField,
-    ) -> Self::BaseField {
-        P::add_and_mul_fp_by_nonresidue_plus_one(x, y)
+    fn mul_base_field_by_frob_coeff(fe: &mut Self::BaseField, power: usize) {
+        *fe *= &Self::FROBENIUS_COEFF_C1[power % Self::DEGREE_OVER_BASE_PRIME_FIELD];
     }
 
     #[inline(always)]
-    fn sub_and_mul_base_field_by_nonresidue(
-        x: &Self::BaseField,
-        y: &Self::BaseField,
+    #[ark_ff_asm::unroll_for_loops]
+    fn op_and_mul_base_field_by_nonresidue(
+        other: &Self::BaseField,
+        fe: &Self::BaseField,
+        mode: MulNonResidueMode,
     ) -> Self::BaseField {
-        P::sub_and_mul_fp_by_nonresidue(x, y)
-    }
-
-    fn mul_base_field_by_frob_coeff(fe: &mut Self::BaseField, power: usize) {
-        *fe *= &Self::FROBENIUS_COEFF_C1[power % Self::DEGREE_OVER_BASE_PRIME_FIELD];
+        if let Some(mut i) = Self::NONRESIDUE_SMALL {
+            if mode == MulNonResidueMode::PlusAddOne {
+                i += 1;
+            }
+            if i.abs() < (1 << 5) {
+                if i == 0 {
+                    return Self::BaseField::zero();
+                }
+                let neg = mode == MulNonResidueMode::Minus;
+                let add = i >= 0 && !neg || i < 0 && neg;
+
+                let res = fe.mul_by_i8(i.abs());
+
+                if add {
+                    return *other + res;
+                } else {
+                    return *other - res;
+                }
+            }
+        }
+        match mode {
+            MulNonResidueMode::Minus => *other - P::mul_fp_by_nonresidue(fe),
+            MulNonResidueMode::Plus => *other + P::mul_fp_by_nonresidue(fe),
+            MulNonResidueMode::PlusAddOne => *other + P::mul_fp_by_nonresidue(fe) + fe,
+        }
     }
 }
 

ff/src/fields/models/fp3.rs

@@ -7,6 +7,8 @@ pub trait Fp3Parameters: 'static + Send + Sync {
 
     const NONRESIDUE: Self::Fp;
 
+    const NONRESIDUE_SMALL: Option<i8> = None;
+
     const FROBENIUS_COEFF_FP3_C1: &'static [Self::Fp];
     const FROBENIUS_COEFF_FP3_C2: &'static [Self::Fp];
 
@@ -18,6 +20,9 @@ pub trait Fp3Parameters: 'static + Send + Sync {
 
     #[inline(always)]
     fn mul_fp_by_nonresidue(fe: &Self::Fp) -> Self::Fp {
+        if let Some(i) = Self::NONRESIDUE_SMALL {
+            return fe.mul_by_i8(i);
+        }
         Self::NONRESIDUE * fe
     }
 }

ff/src/fields/models/fp4.rs

@@ -18,7 +18,7 @@ pub trait Fp4Parameters: 'static + Send + Sync {
     fn mul_fp2_by_nonresidue(fe: &Fp2<Self::Fp2Params>) -> Fp2<Self::Fp2Params> {
         // see [[DESD06, Section 5.1]](https://eprint.iacr.org/2006/471.pdf).
         Fp2::new(
-            <Self::Fp2Params as Fp2Parameters>::NONRESIDUE * &fe.c1,
+            <Self::Fp2Params as Fp2Parameters>::mul_fp_by_nonresidue(&fe.c1),
             fe.c0,
         )
     }

ff/src/fields/models/fp6_2over3.rs

@@ -14,12 +14,11 @@ pub trait Fp6Parameters: 'static + Send + Sync {
 
     #[inline(always)]
     fn mul_fp3_by_nonresidue(fe: &Fp3<Self::Fp3Params>) -> Fp3<Self::Fp3Params> {
-        let mut res = *fe;
-        res.c0 = fe.c2;
-        res.c1 = fe.c0;
-        res.c2 = fe.c1;
-        res.c0 = <Self::Fp3Params as Fp3Parameters>::mul_fp_by_nonresidue(&res.c0);
-        res
+        Fp3::<Self::Fp3Params>::new(
+             <Self::Fp3Params as Fp3Parameters>::mul_fp_by_nonresidue(&fe.c2),
+             fe.c0,
+             fe.c1,
+        )
     }
 }
 

ff/src/fields/models/fp6_3over2.rs

@@ -7,13 +7,39 @@ pub trait Fp6Parameters: 'static + Send + Sync + Copy {
 
     const NONRESIDUE: Fp2<Self::Fp2Params>;
 
+    const NONRESIDUE_SMALL: Option<(i8, i8, Option<i8>)> = None;
+
     /// Coefficients for the Frobenius automorphism.
     const FROBENIUS_COEFF_FP6_C1: &'static [Fp2<Self::Fp2Params>];
     const FROBENIUS_COEFF_FP6_C2: &'static [Fp2<Self::Fp2Params>];
 
     #[inline(always)]
+    #[ark_ff_asm::unroll_for_loops]
+    // (x + y \alpha) * (a + b \alpha) = (xa + ybU) + (xb + ya)\alpha
     fn mul_fp2_by_nonresidue(fe: &Fp2<Self::Fp2Params>) -> Fp2<Self::Fp2Params> {
-        Self::NONRESIDUE * fe
+        if let Some((x, y, f_res)) = Self::NONRESIDUE_SMALL {
+            if x == 0 && y == 0 {
+                return Fp2::<Self::Fp2Params>::zero();
+            } else if y == 0 {
+                return Fp2::<Self::Fp2Params>::new(fe.c0.mul_by_i8(x), fe.c1.mul_by_i8(x));
+            } else {
+                let t0 = if let Some(res) = f_res {
+                    fe.c1.mul_by_i8(y * res)
+                } else {
+                    Self::Fp2Params::mul_fp_by_nonresidue(&fe.c1.mul_by_i8(y))
+                };
+                if x == 0 {
+                    return Fp2::<Self::Fp2Params>::new(t0, fe.c0.mul_by_i8(y));
+                } else {
+                    return Fp2::<Self::Fp2Params>::new(
+                        fe.c0.mul_by_i8(x) + t0,
+                        fe.c0.mul_by_i8(y) + fe.c1.mul_by_i8(x),
+                    );
+                }
+            }
+        } else {
+            Self::NONRESIDUE * fe
+        }
     }
 }
 

ff/src/fields/models/quadratic_extension.rs

@@ -25,6 +25,13 @@ use crate::{
     ToConstraintField, UniformRand,
 };
 
+#[derive(PartialEq)]
+pub enum MulNonResidueMode {
+    Minus,
+    Plus,
+    PlusAddOne,
+}
+
 /// Defines a Quadratic extension field from a quadratic non-residue.
 pub trait QuadExtParameters: 'static + Send + Sync + Sized {
     /// The prime field that this quadratic extension is eventually an extension of.
@@ -41,6 +48,8 @@ pub trait QuadExtParameters: 'static + Send + Sync + Sized {
     /// The quadratic non-residue used to construct the extension.
     const NONRESIDUE: Self::BaseField;
 
+    const NONRESIDUE_SMALL: Option<i8> = None;
+
     /// Coefficients for the Frobenius automorphism.
     const FROBENIUS_COEFF_C1: &'static [Self::FrobCoeff];
 
@@ -52,38 +61,19 @@ pub trait QuadExtParameters: 'static + Send + Sync + Sized {
         Self::NONRESIDUE * fe
     }
 
-    /// A specializable method for computing x + mul_base_field_by_nonresidue(y)
-    /// This allows for optimizations when the non-residue is
-    /// canonically negative in the field.
-    #[inline(always)]
-    fn add_and_mul_base_field_by_nonresidue(
-        x: &Self::BaseField,
-        y: &Self::BaseField,
-    ) -> Self::BaseField {
-        *x + Self::mul_base_field_by_nonresidue(y)
-    }
-
-    /// A specializable method for computing x + mul_base_field_by_nonresidue(y) + y
-    /// This allows for optimizations when the non-residue is not -1.
-    #[inline(always)]
-    fn add_and_mul_base_field_by_nonresidue_plus_one(
-        x: &Self::BaseField,
-        y: &Self::BaseField,
-    ) -> Self::BaseField {
-        let mut tmp = *x;
-        tmp += y;
-        Self::add_and_mul_base_field_by_nonresidue(&tmp, &y)
-    }
-
-    /// A specializable method for computing x - mul_base_field_by_nonresidue(y)
-    /// This allows for optimizations when the non-residue is
-    /// canonically negative in the field.
     #[inline(always)]
-    fn sub_and_mul_base_field_by_nonresidue(
-        x: &Self::BaseField,
-        y: &Self::BaseField,
+    #[ark_ff_asm::unroll_for_loops]
+    fn op_and_mul_base_field_by_nonresidue(
+        other: &Self::BaseField,
+        fe: &Self::BaseField,
+        mode: MulNonResidueMode,
     ) -> Self::BaseField {
-        *x - Self::mul_base_field_by_nonresidue(y)
+        let t0 = Self::mul_base_field_by_nonresidue(fe);
+        match mode {
+            MulNonResidueMode::Minus => *other - t0,
+            MulNonResidueMode::Plus => *other + t0,
+            MulNonResidueMode::PlusAddOne => *other + fe + t0,
+        }
     }
 
     /// A specializable method for multiplying an element of the base field by
@@ -164,9 +154,8 @@ impl<P: QuadExtParameters> QuadExtField<P> {
     pub fn norm(&self) -> P::BaseField {
         let t0 = self.c0.square();
         // t1 = t0 - P::NON_RESIDUE * c1^2
-        let mut t1 = self.c1.square();
-        t1 = P::sub_and_mul_base_field_by_nonresidue(&t0, &t1);
-        t1
+        let t1 = self.c1.square();
+        P::op_and_mul_base_field_by_nonresidue(&t0, &t1, MulNonResidueMode::Minus)
     }
 
     pub fn mul_assign_by_basefield(&mut self, element: &P::BaseField) {
@@ -279,7 +268,11 @@ impl<P: QuadExtParameters> Field for QuadExtField<P> {
             // v0 = c0 - c1
             let mut v0 = self.c0 - &self.c1;
             // v3 = c0 - beta * c1
-            let v3 = P::sub_and_mul_base_field_by_nonresidue(&self.c0, &self.c1);
+            let v3 = P::op_and_mul_base_field_by_nonresidue(
+                &self.c0,
+                &self.c1,
+                MulNonResidueMode::Minus,
+            );
             // v2 = c0 * c1
             let v2 = self.c0 * &self.c1;
 
@@ -294,7 +287,8 @@ impl<P: QuadExtParameters> Field for QuadExtField<P> {
             // result.c0 = (c0^2 - beta * c0 * c1 - c0 * c1 + beta * c1^2) + ((beta + 1) c0 * c1)
             // result.c0 = (c0^2 - beta * c0 * c1 + beta * c1^2) + (beta * c0 * c1)
             // result.c0 = c0^2 + beta * c1^2
-            self.c0 = P::add_and_mul_base_field_by_nonresidue_plus_one(&v0, &v2);
+            self.c0 =
+                P::op_and_mul_base_field_by_nonresidue(&v0, &v2, MulNonResidueMode::PlusAddOne);
 
             self
         }
@@ -308,7 +302,11 @@ impl<P: QuadExtParameters> Field for QuadExtField<P> {
             // v1 = c1.square()
             let v1 = self.c1.square();
             // v0 = c0.square() - beta * v1
-            let v0 = P::sub_and_mul_base_field_by_nonresidue(&self.c0.square(), &v1);
+            let v0 = P::op_and_mul_base_field_by_nonresidue(
+                &self.c0.square(),
+                &v1,
+                MulNonResidueMode::Minus,
+            );
 
             v0.inverse().map(|v1| {
                 let c0 = self.c0 * &v1;
@@ -571,7 +569,7 @@ impl<'a, P: QuadExtParameters> MulAssign<&'a Self> for QuadExtField<P> {
         self.c1 *= &(other.c0 + &other.c1);
         self.c1 -= &v0;
         self.c1 -= &v1;
-        self.c0 = P::add_and_mul_base_field_by_nonresidue(&v0, &v1);
+        self.c0 = P::op_and_mul_base_field_by_nonresidue(&v0, &v1, MulNonResidueMode::Plus);
     }
 }