From 2b3dd5db2b0cadb2b99e9eb973ccddc9cd79bfeb Mon Sep 17 00:00:00 2001
From: Jonathan Wang <31040440+jonathanpwang@users.noreply.github.com>
Date: Sun, 6 Aug 2023 16:48:01 -0400
Subject: [PATCH] feat(halo2-base): add `GateChip::pow_var` (#103)

---
 halo2-base/src/gates/flex_gate.rs             | 35 +++++++++++++++++++
 halo2-base/src/gates/tests/flex_gate.rs       | 12 +++++++
 halo2-base/src/gates/tests/pos_prop.rs        |  7 ++++
 .../src/keccak_packed_multi/tests.rs          |  4 ++-
 4 files changed, 57 insertions(+), 1 deletion(-)

diff --git a/halo2-base/src/gates/flex_gate.rs b/halo2-base/src/gates/flex_gate.rs
index 25b0da24..999cfd77 100644
--- a/halo2-base/src/gates/flex_gate.rs
+++ b/halo2-base/src/gates/flex_gate.rs
@@ -762,6 +762,17 @@ pub trait GateInstructions<F: ScalarField> {
         range_bits: usize,
     ) -> Vec<AssignedValue<F>>;
 
+    /// Constrains and computes `a`<sup>`exp`</sup> where both `a, exp` are witnesses. The exponent is computed in the native field `F`.
+    ///
+    /// Constrains that `exp` has at most `max_bits` bits.
+    fn pow_var(
+        &self,
+        ctx: &mut Context<F>,
+        a: AssignedValue<F>,
+        exp: AssignedValue<F>,
+        max_bits: usize,
+    ) -> AssignedValue<F>;
+
     /// Performs and constrains Lagrange interpolation on `coords` and evaluates the resulting polynomial at `x`.
     ///
     /// Given pairs `coords[i] = (x_i, y_i)`, let `f` be the unique degree `len(coords) - 1` polynomial such that `f(x_i) = y_i` for all `i`.
@@ -1137,4 +1148,28 @@ impl<F: ScalarField> GateInstructions<F> for GateChip<F> {
         }
         bit_cells
     }
+
+    /// Constrains and computes `a^exp` where both `a, exp` are witnesses. The exponent is computed in the native field `F`.
+    ///
+    /// Constrains that `exp` has at most `max_bits` bits.
+    fn pow_var(
+        &self,
+        ctx: &mut Context<F>,
+        a: AssignedValue<F>,
+        exp: AssignedValue<F>,
+        max_bits: usize,
+    ) -> AssignedValue<F> {
+        let exp_bits = self.num_to_bits(ctx, exp, max_bits);
+        // standard square-and-mul approach
+        let mut acc = ctx.load_constant(F::one());
+        for (i, bit) in exp_bits.into_iter().rev().enumerate() {
+            if i > 0 {
+                // square
+                acc = self.mul(ctx, acc, acc);
+            }
+            let mul = self.mul(ctx, acc, a);
+            acc = self.select(ctx, mul, acc, bit);
+        }
+        acc
+    }
 }
diff --git a/halo2-base/src/gates/tests/flex_gate.rs b/halo2-base/src/gates/tests/flex_gate.rs
index 8b047504..8a4a6e7a 100644
--- a/halo2-base/src/gates/tests/flex_gate.rs
+++ b/halo2-base/src/gates/tests/flex_gate.rs
@@ -1,8 +1,10 @@
 #![allow(clippy::type_complexity)]
 use super::*;
+use crate::utils::biguint_to_fe;
 use crate::utils::testing::base_test;
 use crate::QuantumCell::Witness;
 use crate::{gates::flex_gate::GateInstructions, QuantumCell};
+use num_bigint::BigUint;
 use test_case::test_case;
 
 #[test_case(&[10, 12].map(Fr::from).map(Witness)=> Fr::from(22); "add(): 10 + 12 == 22")]
@@ -172,3 +174,13 @@ pub fn test_num_to_bits(num: usize, bits: usize) -> Vec<Fr> {
         chip.num_to_bits(ctx, num, bits).iter().map(|a| *a.value()).collect()
     })
 }
+
+#[test_case(Fr::from(3), BigUint::from(3u32), 4 => Fr::from(27); "pow_var(): 3^3 = 27")]
+pub fn test_pow_var(a: Fr, exp: BigUint, max_bits: usize) -> Fr {
+    assert!(exp.bits() <= max_bits as u64);
+    base_test().run_gate(|ctx, chip| {
+        let a = ctx.load_witness(a);
+        let exp = ctx.load_witness(biguint_to_fe(&exp));
+        *chip.pow_var(ctx, a, exp, max_bits).value()
+    })
+}
diff --git a/halo2-base/src/gates/tests/pos_prop.rs b/halo2-base/src/gates/tests/pos_prop.rs
index 2d3a6cca..dc4e3702 100644
--- a/halo2-base/src/gates/tests/pos_prop.rs
+++ b/halo2-base/src/gates/tests/pos_prop.rs
@@ -266,6 +266,13 @@ proptest! {
         prop_assert_eq!(bits.into_iter().map(Fr::from).collect::<Vec<_>>(), result);
     }
 
+    #[test]
+    fn prop_test_pow_var(a in rand_fr(), num in any::<u64>()) {
+        let native_res = a.pow_vartime([num]);
+        let result = flex_gate::test_pow_var(a, BigUint::from(num), Fr::CAPACITY as usize);
+        prop_assert_eq!(result, native_res);
+    }
+
     /*
     #[test]
     fn prop_test_lagrange_eval(inputs in vec(rand_fr(), 3)) {
diff --git a/hashes/zkevm-keccak/src/keccak_packed_multi/tests.rs b/hashes/zkevm-keccak/src/keccak_packed_multi/tests.rs
index a3df0b0e..0797ef13 100644
--- a/hashes/zkevm-keccak/src/keccak_packed_multi/tests.rs
+++ b/hashes/zkevm-keccak/src/keccak_packed_multi/tests.rs
@@ -101,7 +101,7 @@ impl<F: Field> KeccakCircuit<F> {
 }
 
 fn verify<F: Field>(k: u32, inputs: Vec<Vec<u8>>, _success: bool) {
-    let circuit = KeccakCircuit::new(Some(2usize.pow(k)), inputs);
+    let circuit = KeccakCircuit::new(Some(2usize.pow(k) - 109), inputs);
 
     let prover = MockProver::<F>::run(k, &circuit, vec![]).unwrap();
     prover.assert_satisfied();
@@ -150,6 +150,7 @@ fn packed_multi_keccak_prover(k: u32, rows_per_round: usize) {
     let verifier_params: ParamsVerifierKZG<Bn256> = params.verifier_params().clone();
     let mut transcript = Blake2bWrite::<_, G1Affine, Challenge255<_>>::init(vec![]);
 
+    let start = std::time::Instant::now();
     create_proof::<
         KZGCommitmentScheme<Bn256>,
         ProverSHPLONK<'_, Bn256>,
@@ -160,6 +161,7 @@ fn packed_multi_keccak_prover(k: u32, rows_per_round: usize) {
     >(&params, &pk, &[circuit], &[&[]], OsRng, &mut transcript)
     .expect("proof generation should not fail");
     let proof = transcript.finalize();
+    dbg!(start.elapsed());
 
     let mut verifier_transcript = Blake2bRead::<_, G1Affine, Challenge255<_>>::init(&proof[..]);
     let strategy = SingleStrategy::new(&params);