chore: mitigate copies in final_exponentiation (#12)

argumentcomputer · May 16, 2024 · b052a66 · b052a66
1 parent 8a493f7
commit b052a66
Show file tree

Hide file tree

Showing 4 changed files with 390 additions and 72 deletions.
diff --git a/src/fp12.rs b/src/fp12.rs
@@ -162,6 +162,44 @@ impl Fp12 {
         }
     }
 
+    #[inline]
+    #[cfg(target_os = "zkvm")]
+    pub fn mul_inp(&mut self, other: &Fp12) {
+        let aa = self.c0 * other.c0;
+        let bb = self.c1 * other.c1;
+        self.c1.add_inp(&self.c0);
+        let mut o = other.c0;
+        o.add_inp(&other.c1);
+        self.c1 *= o;
+        self.c1.sub_inp(&aa);
+        self.c1.sub_inp(&bb);
+        self.c0 = bb.mul_by_nonresidue_owned();
+        self.c0.add_inp(&aa);
+    }
+
+    #[inline]
+    fn mul(&self, other: &Fp12) -> Self {
+        cfg_if::cfg_if! {
+            if #[cfg(target_os = "zkvm")] {
+                let mut out = self.clone();
+                out.mul_inp(other);
+                out
+            } else {
+                let aa = self.c0 * other.c0;
+                let bb = self.c1 * other.c1;
+                let c1 = self.c1 + self.c0;
+                let o = other.c0 + other.c1;
+                let c1 = c1 * o;
+                let c1 = c1 - aa;
+                let c1 = c1 - bb;
+                let c0 = bb.mul_by_nonresidue();
+                let c0 = c0 + aa;
+
+                Fp12 { c0, c1 }
+            }
+        }
+    }
+
     /// Raises this element to p.
     #[inline(always)]
     pub fn frobenius_map(&self) -> Self {
@@ -192,6 +230,40 @@ impl Fp12 {
         Fp12 { c0, c1 }
     }
 
+    /// Raises this element to p.
+    #[inline]
+    #[cfg(target_os = "zkvm")]
+    pub fn frobenius_map_inp(&mut self) {
+        self.c0.frobenius_map_inp();
+        self.c1.frobenius_map_inp();
+
+        const C1_MUL: Fp6 = Fp6 {
+            c0: Fp2 {
+                c0: Fp::from_raw_unchecked([
+                    0x0708_9552_b319_d465,
+                    0xc669_5f92_b50a_8313,
+                    0x97e8_3ccc_d117_228f,
+                    0xa35b_aeca_b2dc_29ee,
+                    0x1ce3_93ea_5daa_ce4d,
+                    0x08f2_220f_b0fb_66eb,
+                ]),
+                c1: Fp::from_raw_unchecked([
+                    0xb2f6_6aad_4ce5_d646,
+                    0x5842_a06b_fc49_7cec,
+                    0xcf48_95d4_2599_d394,
+                    0xc11b_9cba_40a8_e8d0,
+                    0x2e38_13cb_e5a0_de89,
+                    0x110e_efda_8884_7faf,
+                ]),
+            },
+            c1: Fp2::zero(),
+            c2: Fp2::zero(),
+        };
+
+        // c1 = c1 * (u + 1)^((p - 1) / 6)
+        self.c1 *= C1_MUL;
+    }
+
     #[inline]
     #[cfg(target_os = "zkvm")]
     pub fn square(&self) -> Self {
@@ -237,17 +309,7 @@ impl<'a, 'b> Mul<&'b Fp12> for &'a Fp12 {
 
     #[inline]
     fn mul(self, other: &'b Fp12) -> Self::Output {
-        let aa = self.c0 * other.c0;
-        let bb = self.c1 * other.c1;
-        let o = other.c0 + other.c1;
-        let c1 = self.c1 + self.c0;
-        let c1 = c1 * o;
-        let c1 = c1 - aa;
-        let c1 = c1 - bb;
-        let c0 = bb.mul_by_nonresidue();
-        let c0 = c0 + aa;
-
-        Fp12 { c0, c1 }
+        self.mul(other)
     }
 }
 

diff --git a/src/fp2.rs b/src/fp2.rs
@@ -153,6 +153,14 @@ impl Fp2 {
         self.conjugate()
     }
 
+    /// Raises this element to p.
+    #[inline]
+    pub fn frobenius_map_inp(&mut self) {
+        // This is always just a conjugation. If you're curious why, here's
+        // an article about it: https://alicebob.cryptoland.net/the-frobenius-endomorphism-with-finite-fields/
+        self.conjugate_inp()
+    }
+
     #[inline(always)]
     pub fn conjugate(&self) -> Self {
         Fp2 {
@@ -161,6 +169,11 @@ impl Fp2 {
         }
     }
 
+    #[inline]
+    pub fn conjugate_inp(&mut self) {
+        self.c1 = -self.c1;
+    }
+
     #[inline]
     #[cfg(target_os = "zkvm")]
     pub fn mul_by_nonresidue_inp(&mut self) {

diff --git a/src/fp6.rs b/src/fp6.rs
@@ -165,6 +165,23 @@ impl Fp6 {
         }
     }
 
+    /// Multiply by quadratic nonresidue v.
+    #[cfg(target_os = "zkvm")]
+    pub fn mul_by_nonresidue_owned(self) -> Self {
+        // Given a + bv + cv^2, this produces
+        //     av + bv^2 + cv^3
+        // but because v^3 = u + 1, we have
+        //     c(u + 1) + av + v^2
+        let Fp6 { c0, c1, mut c2 } = self;
+        c2.mul_by_nonresidue_inp();
+
+        Fp6 {
+            c0: c2,
+            c1: c0,
+            c2: c1,
+        }
+    }
+
     /// Raises this element to p.
     #[inline(always)]
     pub fn frobenius_map(&self) -> Self {
@@ -203,6 +220,45 @@ impl Fp6 {
         Fp6 { c0, c1, c2 }
     }
 
+    /// Raises this element to p.
+    #[inline]
+    #[cfg(target_os = "zkvm")]
+    pub fn frobenius_map_inp(&mut self) {
+        self.c0.frobenius_map_inp();
+        self.c1.frobenius_map_inp();
+        self.c2.frobenius_map_inp();
+
+        const C1_MUL: Fp2 = Fp2 {
+            c0: Fp::zero(),
+            c1: Fp::from_raw_unchecked([
+                0xcd03_c9e4_8671_f071,
+                0x5dab_2246_1fcd_a5d2,
+                0x5870_42af_d385_1b95,
+                0x8eb6_0ebe_01ba_cb9e,
+                0x03f9_7d6e_83d0_50d2,
+                0x18f0_2065_5463_8741,
+            ]),
+        };
+
+        const C2_MUL: Fp2 = Fp2 {
+            c0: Fp::from_raw_unchecked([
+                0x890d_c9e4_8675_45c3,
+                0x2af3_2253_3285_a5d5,
+                0x5088_0866_309b_7e2c,
+                0xa20d_1b8c_7e88_1024,
+                0x14e4_f04f_e2db_9068,
+                0x14e5_6d3f_1564_853a,
+            ]),
+            c1: Fp::zero(),
+        };
+
+        // c1 = c1 * (u + 1)^((p - 1) / 3)
+        self.c1.mul_inp(&C1_MUL);
+
+        // c2 = c2 * (u + 1)^((2p - 2) / 3)
+        self.c2.mul_inp(&C2_MUL);
+    }
+
     #[inline(always)]
     pub fn is_zero(&self) -> Choice {
         self.c0.is_zero() & self.c1.is_zero() & self.c2.is_zero()
@@ -250,10 +306,23 @@ impl Fp6 {
         // Each of these is a "sum of products", which we can compute efficiently.
 
         let a = self;
-        let b10_p_b11 = b.c1.c0 + b.c1.c1;
-        let b10_m_b11 = b.c1.c0 - b.c1.c1;
-        let b20_p_b21 = b.c2.c0 + b.c2.c1;
-        let b20_m_b21 = b.c2.c0 - b.c2.c1;
+        cfg_if::cfg_if! {
+            if #[cfg(target_os = "zkvm")] {
+                let mut b10_p_b11 = b.c1.c0;
+                b10_p_b11.add_inp(&b.c1.c1);
+                let mut b10_m_b11 = b.c1.c0;
+                b10_m_b11.sub_inp(&b.c1.c1);
+                let mut b20_p_b21 = b.c2.c0;
+                b20_p_b21.add_inp(&b.c2.c1);
+                let mut b20_m_b21 = b.c2.c0;
+                b20_m_b21.sub_inp(&b.c2.c1)
+            } else {
+                let b10_p_b11 = b.c1.c0 + b.c1.c1;
+                let b10_m_b11 = b.c1.c0 - b.c1.c1;
+                let b20_p_b21 = b.c2.c0 + b.c2.c1;
+                let b20_m_b21 = b.c2.c0 - b.c2.c1;
+            }
+        }
 
         Fp6 {
             c0: Fp2 {