diff --git a/rusk-abi/CHANGELOG.md b/rusk-abi/CHANGELOG.md index 5bc2525730..2fd12cb248 100644 --- a/rusk-abi/CHANGELOG.md +++ b/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] diff --git a/rusk-abi/Cargo.toml b/rusk-abi/Cargo.toml index d790267746..726af67b48 100644 --- a/rusk-abi/Cargo.toml +++ b/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" diff --git a/rusk-abi/src/host.rs b/rusk-abi/src/host.rs index 75ce095a9a..5b3c651330 100644 --- a/rusk-abi/src/host.rs +++ b/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 { PoseidonHash::digest(Domain::Other, &scalars)[0] } -/// A simple LRU cache for plonk verification. -/// -/// # Safety -/// `f` should not panic. -unsafe fn with_verification_cache(f: F) -> T -where - F: FnOnce(MutexGuard>) -> T, -{ - const VERIFICATION_CACHE_SIZE: usize = 512; - - static CACHE: OnceLock< - Mutex>, - > = 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 { - // 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, public_inputs: Vec, ) -> 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 diff --git a/rusk-abi/src/host/cache.rs b/rusk-abi/src/host/cache.rs new file mode 100644 index 0000000000..dd52fb82b9 --- /dev/null +++ b/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 { + // 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(f: F) -> T + where + F: FnOnce( + MutexGuard>, + ) -> T, + { + const DEFAULT_SIZE: usize = $size; + + static CACHE: OnceLock< + Mutex>, + > = 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" +);