add caching for jwk.n hash

Author: joyqvq

fastcrypto-zkp/Cargo.toml

@@ -15,6 +15,7 @@ harness = false
 [[bench]]
 name = "zklogin"
 harness = false
+required-features = ["test-utils"]
 
 [[bench]]
 name = "poseidon"
@@ -64,3 +65,4 @@ proptest = "1.1.0"
 
 [features]
 e2e = []
+test-utils = []

fastcrypto-zkp/benches/zklogin.rs

@@ -8,12 +8,13 @@ mod zklogin_benches {
 
     use ark_std::rand::rngs::StdRng;
     use ark_std::rand::SeedableRng;
-    use criterion::Criterion;
+    use criterion::{BatchSize, Criterion};
     use fastcrypto::ed25519::Ed25519KeyPair;
     use fastcrypto::error::FastCryptoError;
     use fastcrypto::rsa::{Base64UrlUnpadded, Encoding};
     use fastcrypto::traits::KeyPair;
     use fastcrypto_zkp::bn254::utils::gen_address_seed;
+    use fastcrypto_zkp::bn254::zk_login::clear_cache_for_testing;
     use fastcrypto_zkp::bn254::zk_login::ZkLoginInputs;
     use fastcrypto_zkp::bn254::zk_login::JWK;
     use fastcrypto_zkp::bn254::zk_login::{JwkId, OIDCProvider};
@@ -106,18 +107,38 @@ mod zklogin_benches {
             },
         );
 
-        // Benchmark the entire `verify_zk_login` function.
-        c.bench_function("verify_zk_login", move |b| {
+        // Benchmark the entire `verify_zk_login` function (warm: modulus hash cache hit).
+        let input_warm = input.clone();
+        let eph_warm = eph_pubkey.clone();
+        let map_warm = map.clone();
+        c.bench_function("verify_zk_login/warm", move |b| {
             b.iter(|| {
                 fastcrypto_zkp::bn254::zk_login_api::verify_zk_login(
-                    &input,
+                    &input_warm,
                     max_epoch,
-                    &eph_pubkey,
-                    &map,
+                    &eph_warm,
+                    &map_warm,
                     &ZkLoginEnv::Test,
                 )
             })
         });
+
+        // Benchmark `verify_zk_login` on a cold cache (first call after a JWK refresh).
+        c.bench_function("verify_zk_login/cold", move |b| {
+            b.iter_batched(
+                clear_cache_for_testing,
+                |_| {
+                    fastcrypto_zkp::bn254::zk_login_api::verify_zk_login(
+                        &input,
+                        max_epoch,
+                        &eph_pubkey,
+                        &map,
+                        &ZkLoginEnv::Test,
+                    )
+                },
+                BatchSize::SmallInput,
+            )
+        });
     }
 
     /// Benchmark V2 proof verification for 8192-bit RSA keys
@@ -199,18 +220,38 @@ mod zklogin_benches {
             },
         );
 
-        // Benchmark the entire `verify_zk_login` function.
-        c.bench_function("verify_zk_login_v2", move |b| {
+        // Benchmark the entire `verify_zk_login` function (warm: modulus hash cache hit).
+        let input_warm = input.clone();
+        let eph_warm = eph_pubkey.clone();
+        let map_warm = map.clone();
+        c.bench_function("verify_zk_login_v2/warm", move |b| {
             b.iter(|| {
                 fastcrypto_zkp::bn254::zk_login_api::verify_zk_login(
-                    &input,
+                    &input_warm,
                     max_epoch,
-                    &eph_pubkey,
-                    &map,
+                    &eph_warm,
+                    &map_warm,
                     &ZkLoginEnv::Test,
                 )
             })
         });
+
+        // Benchmark `verify_zk_login` on a cold cache (first call after a JWK refresh).
+        c.bench_function("verify_zk_login_v2/cold", move |b| {
+            b.iter_batched(
+                clear_cache_for_testing,
+                |_| {
+                    fastcrypto_zkp::bn254::zk_login_api::verify_zk_login(
+                        &input,
+                        max_epoch,
+                        &eph_pubkey,
+                        &map,
+                        &ZkLoginEnv::Test,
+                    )
+                },
+                BatchSize::SmallInput,
+            )
+        });
     }
 
     criterion_group! {

fastcrypto-zkp/src/bn254/zk_login.rs

@@ -21,13 +21,47 @@ pub use ark_serialize::{CanonicalDeserialize, CanonicalSerialize};
 use fastcrypto::error::FastCryptoError;
 use itertools::Itertools;
 use num_bigint::BigUint;
+use once_cell::sync::Lazy;
 use regex::Regex;
 use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
 use std::cmp::Ordering::{Equal, Greater, Less};
+use std::collections::HashMap;
 use std::error::Error;
 use std::fmt::Display;
 use std::str::FromStr;
+use std::sync::RwLock;
+
+/// Key for the modulus hash cache: (modulus bytes, max_rsa_bits).
+type ModulusHashKey = (Vec<u8>, u16);
+
+/// JWKs rotate occasionally, so caching by (modulus bytes, max_rsa_bits) avoids recomputing
+/// bit-packing + poseidon hash on every verification.
+static MODULUS_HASH_CACHE: Lazy<RwLock<HashMap<ModulusHashKey, Bn254Fr>>> =
+    Lazy::new(|| RwLock::new(HashMap::new()));
+
+fn cached_modulus_hash(modulus: &[u8], max_rsa_bits: u16) -> Result<Bn254Fr, FastCryptoError> {
+    if let Some(f) = MODULUS_HASH_CACHE
+        .read()
+        .ok()
+        .and_then(|m| m.get(&(modulus.to_vec(), max_rsa_bits)).copied())
+    {
+        return Ok(f);
+    }
+    let f = hash_to_field(&[BigUint::from_bytes_be(modulus)], max_rsa_bits, PACK_WIDTH)?;
+    if let Ok(mut m) = MODULUS_HASH_CACHE.write() {
+        m.insert((modulus.to_vec(), max_rsa_bits), f);
+    }
+    Ok(f)
+}
+
+/// Clear the modulus hash cache for testing only benchmark.
+#[cfg(any(test, feature = "test-utils"))]
+pub fn clear_cache_for_testing() {
+    if let Ok(mut m) = MODULUS_HASH_CACHE.write() {
+        m.clear();
+    }
+}
 
 #[cfg(test)]
 #[path = "unit_tests/zk_login_tests.rs"]
@@ -589,11 +623,7 @@ impl ZkLoginInputs {
         let iss_base64_f =
             hash_ascii_str_to_field(&self.iss_base64_details.value, config.max_iss_len_b64)?;
         let header_f = hash_ascii_str_to_field(&self.header_base64, config.max_header_len_b64)?;
-        let modulus_f = hash_to_field(
-            &[BigUint::from_bytes_be(modulus)],
-            config.max_rsa_bits,
-            PACK_WIDTH,
-        )?;
+        let modulus_f = cached_modulus_hash(modulus, config.max_rsa_bits)?;
         poseidon_zk_login(&[
             first,
             second,