From df690f92bf0b49c71eb81e83eb7697b9707925ee Mon Sep 17 00:00:00 2001
From: wborgeaud <williamborgeaud@gmail.com>
Date: Mon, 2 Aug 2021 13:31:12 +0200
Subject: [PATCH] Hardcode small exponents in `exp_u64_extension`

---
 src/gadgets/arithmetic_extension.rs | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/src/gadgets/arithmetic_extension.rs b/src/gadgets/arithmetic_extension.rs
index c1bedb9d..812e7317 100644
--- a/src/gadgets/arithmetic_extension.rs
+++ b/src/gadgets/arithmetic_extension.rs
@@ -315,6 +315,11 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
         self.mul_extension(x, x)
     }
 
+    /// Computes `x^3`.
+    pub fn cube_extension(&mut self, x: ExtensionTarget<D>) -> ExtensionTarget<D> {
+        self.mul_three_extension(x, x, x)
+    }
+
     pub fn mul_ext_algebra(
         &mut self,
         a: ExtensionAlgebraTarget<D>,
@@ -450,6 +455,13 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
         base: ExtensionTarget<D>,
         exponent: u64,
     ) -> ExtensionTarget<D> {
+        match exponent {
+            0 => return self.one_extension(),
+            1 => return base,
+            2 => return self.square_extension(base),
+            3 => return self.cube_extension(base),
+            _ => (),
+        }
         let mut current = base;
         let mut product = self.one_extension();