Add BLS signature cache

rusk-abi/CHANGELOG.md

@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Memoize the `verify_bls` function
 - Memoize the `verify_proof` function [#1228]
 - New ABIs: `owner_raw`, `self_owner_raw` [#1710]
 

rusk-abi/Cargo.toml

@@ -25,7 +25,7 @@ execution-core = { version = "0.1.0", path = "../execution-core" }
 # These are patches since these crates don't seem to like semver.
 rkyv = { version = "=0.7.39", default-features = false, features = ["size_32"] }
 
-lru = "0.12"
+lru = { version = "0.12", optional = true }
 
 [dev-dependencies]
 rand_core = { version = "0.6", default-features = false, features = ["getrandom"] }
@@ -44,8 +44,8 @@ dlmalloc = ["piecrust-uplink/dlmalloc"]
 
 # These are the features available for when one wishes to use `rusk-abi` as a
 # host.
-host = ["piecrust"]
-host_debug = ["piecrust/debug"]
+host = ["piecrust", "lru"]
+host_debug = ["piecrust/debug", "lru"]
 
 [[test]]
 name = "test-rusk-abi"

rusk-abi/src/host.rs

@@ -5,11 +5,7 @@
 // Copyright (c) DUSK NETWORK. All rights reserved.
 
 use alloc::vec::Vec;
-use blake2b_simd::Params;
-use std::env;
-use std::num::NonZeroUsize;
 use std::path::{Path, PathBuf};
-use std::sync::{Mutex, MutexGuard, OnceLock};
 
 use dusk_bytes::DeserializableSlice;
 use dusk_plonk::prelude::{Proof, Verifier};
@@ -18,10 +14,11 @@ use execution_core::{
     BlsAggPublicKey, BlsPublicKey, BlsScalar, BlsSignature, SchnorrPublicKey,
     SchnorrSignature,
 };
-use lru::LruCache;
 use rkyv::ser::serializers::AllocSerializer;
 use rkyv::{Archive, Deserialize, Serialize};
 
+mod cache;
+
 pub use piecrust::*;
 
 use crate::hash::Hasher;
@@ -104,8 +101,13 @@ fn host_poseidon_hash(arg_buf: &mut [u8], arg_len: u32) -> u32 {
 }
 
 fn host_verify_proof(arg_buf: &mut [u8], arg_len: u32) -> u32 {
+    let hash = *blake2b_simd::blake2b(&arg_buf[..arg_len as usize]).as_array();
+    let cached = cache::get_plonk_verification(hash);
+
     wrap_host_query(arg_buf, arg_len, |(vd, proof, pis)| {
-        verify_proof(vd, proof, pis)
+        let is_valid = cached.unwrap_or_else(|| verify_proof(vd, proof, pis));
+        cache::put_plonk_verification(hash, is_valid);
+        is_valid
     })
 }
 
@@ -116,7 +118,14 @@ fn host_verify_schnorr(arg_buf: &mut [u8], arg_len: u32) -> u32 {
 }
 
 fn host_verify_bls(arg_buf: &mut [u8], arg_len: u32) -> u32 {
-    wrap_host_query(arg_buf, arg_len, |(msg, pk, sig)| verify_bls(msg, pk, sig))
+    let hash = *blake2b_simd::blake2b(&arg_buf[..arg_len as usize]).as_array();
+    let cached = cache::get_bls_verification(hash);
+
+    wrap_host_query(arg_buf, arg_len, |(msg, pk, sig)| {
+        let is_valid = cached.unwrap_or_else(|| verify_bls(msg, pk, sig));
+        cache::put_bls_verification(hash, is_valid);
+        is_valid
+    })
 }
 
 /// Compute the blake2b hash of the given scalars, returning the resulting
@@ -131,54 +140,6 @@ pub fn poseidon_hash(scalars: Vec<BlsScalar>) -> BlsScalar {
     PoseidonHash::digest(Domain::Other, &scalars)[0]
 }
 
-/// A simple LRU cache for plonk verification.
-///
-/// # Safety
-/// `f` should not panic.
-unsafe fn with_verification_cache<T, F>(f: F) -> T
-where
-    F: FnOnce(MutexGuard<LruCache<[u8; blake2b_simd::OUTBYTES], bool>>) -> T,
-{
-    const VERIFICATION_CACHE_SIZE: usize = 512;
-
-    static CACHE: OnceLock<
-        Mutex<LruCache<[u8; blake2b_simd::OUTBYTES], bool>>,
-    > = OnceLock::new();
-
-    CACHE
-        .get_or_init(|| {
-            let mut cache_size = None;
-
-            if let Ok(s) = env::var("RUSK_ABI_PREFERIFY_CACHE_SIZE") {
-                cache_size = s.parse().ok();
-            }
-
-            let mut cache_size = cache_size.unwrap_or(VERIFICATION_CACHE_SIZE);
-            if cache_size == 0 {
-                cache_size = VERIFICATION_CACHE_SIZE;
-            }
-
-            Mutex::new(LruCache::new(NonZeroUsize::new(cache_size).unwrap()))
-        })
-        .lock()
-        .map(f)
-        .unwrap()
-}
-
-fn get_cache(hash: [u8; blake2b_simd::OUTBYTES]) -> Option<bool> {
-    // SAFETY: The cache never panics
-    unsafe { with_verification_cache(|mut cache| cache.get(&hash).copied()) }
-}
-
-fn put_cache(hash: [u8; blake2b_simd::OUTBYTES], verified: bool) {
-    // SAFETY: The cache never panics
-    unsafe {
-        with_verification_cache(|mut cache| {
-            cache.put(hash, verified);
-        });
-    }
-}
-
 /// Verify a proof is valid for a given circuit type and public inputs
 ///
 /// # Panics
@@ -188,22 +149,6 @@ pub fn verify_proof(
     proof: Vec<u8>,
     public_inputs: Vec<PublicInput>,
 ) -> bool {
-    let mut hasher = Params::default().to_state();
-
-    hasher.update(&verifier_data);
-    hasher.update(&proof);
-    public_inputs
-        .iter()
-        .for_each(|pi| pi.update_hasher(&mut hasher));
-
-    let hash = *hasher.finalize().as_array();
-
-    // If the proof verification has been memoized with the same arguments,
-    // return the result
-    if let Some(v) = get_cache(hash) {
-        return v;
-    }
-
     let verifier = Verifier::try_from_bytes(verifier_data)
         .expect("Verifier data coming from the contract should be valid");
     let proof = Proof::from_slice(&proof).expect("Proof should be valid");
@@ -227,11 +172,7 @@ pub fn verify_proof(
         }
     });
 
-    let verified = verifier.verify(&proof, &pis[..]).is_ok();
-    if verified {
-        put_cache(hash, verified);
-    }
-    verified
+    verifier.verify(&proof, &pis[..]).is_ok()
 }
 
 /// Verify a schnorr signature is valid for the given public key and message

rusk-abi/src/host/cache.rs

@@ -0,0 +1,88 @@
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this
+// file, You can obtain one at http://mozilla.org/MPL/2.0/.
+//
+// Copyright (c) DUSK NETWORK. All rights reserved.
+
+use std::env;
+use std::num::NonZeroUsize;
+use std::sync::{Mutex, MutexGuard, OnceLock};
+
+use lru::LruCache;
+
+macro_rules! define_cache {
+    ($get_func:ident, $put_func:ident, $cache_func:ident, $type:ty, $size:literal, $var:literal) => {
+        /// Gets an entry out of the cache. Returns `None` if there is no
+        /// element in the cache. `Some` signifies that there is a
+        /// cache element.
+        pub fn $get_func(hash: [u8; blake2b_simd::OUTBYTES]) -> Option<bool> {
+            // SAFETY: the closure never panics
+            unsafe { $cache_func(|mut cache| cache.get(&hash).copied()) }
+        }
+
+        /// Put an entry into the cache.
+        pub fn $put_func(hash: [u8; blake2b_simd::OUTBYTES], is_valid: bool) {
+            // SAFETY: The closure never panics
+            unsafe {
+                $cache_func(|mut cache| {
+                    cache.put(hash, is_valid);
+                });
+            }
+        }
+
+        /// A simple LRU cache.
+        ///
+        /// # Safety
+        /// `f` should *never* panic, otherwise we poison the Mutex.
+        unsafe fn $cache_func<T, F>(f: F) -> T
+        where
+            F: FnOnce(
+                MutexGuard<LruCache<[u8; blake2b_simd::OUTBYTES], $type>>,
+            ) -> T,
+        {
+            const DEFAULT_SIZE: usize = $size;
+
+            static CACHE: OnceLock<
+                Mutex<LruCache<[u8; blake2b_simd::OUTBYTES], $type>>,
+            > = OnceLock::new();
+
+            CACHE
+                .get_or_init(|| {
+                    let mut cache_size = None;
+
+                    if let Ok(s) = env::var($var) {
+                        cache_size = s.parse().ok();
+                    }
+
+                    let mut cache_size = cache_size.unwrap_or(DEFAULT_SIZE);
+                    if cache_size == 0 {
+                        cache_size = DEFAULT_SIZE;
+                    }
+
+                    Mutex::new(LruCache::new(
+                        NonZeroUsize::new(cache_size).unwrap(),
+                    ))
+                })
+                .lock()
+                .map(f)
+                .unwrap()
+        }
+    };
+}
+
+define_cache!(
+    get_plonk_verification,
+    put_plonk_verification,
+    with_plonk_cache,
+    bool,
+    512,
+    "RUSK_ABI_PLONK_CACHE_SIZE"
+);
+define_cache!(
+    get_bls_verification,
+    put_bls_verification,
+    with_bls_cache,
+    bool,
+    512,
+    "RUSK_ABI_BLS_CACHE_SIZE"
+);