From 6ade1ec5bb5ed2a4f1dca0bf9b38c33407f71976 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Eduardo=20Leegwater=20Sim=C3=B5es?= Date: Tue, 23 Jul 2024 17:45:40 +0200 Subject: [PATCH] rusk-abi: add BLS signature cache Cache definition placed in a different file, and abstracted away with a macro, allowing one to generate different cache types according to in-code parameters. This macro is then leveraged to generate both Plonk proof and BLS signature caches, which are then subsequently used in our code. To slightly improve the implementation of both, the hash used as a cache key is now computed using the argument buffer contents, as opposed to the deserialized parameters. This allows for a cleaner and more performant implementation. See-also: #1984 --- rusk-abi/CHANGELOG.md | 1 + rusk-abi/Cargo.toml | 6 +-- rusk-abi/src/host.rs | 93 +++++++------------------------------- rusk-abi/src/host/cache.rs | 88 ++++++++++++++++++++++++++++++++++++ 4 files changed, 109 insertions(+), 79 deletions(-) create mode 100644 rusk-abi/src/host/cache.rs 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" +);