From 72c5cdbdd1d30c254baa51821d97db1a7103b13e Mon Sep 17 00:00:00 2001 From: Sergey Melnychuk Date: Fri, 1 Nov 2024 00:23:30 +0100 Subject: [PATCH] feat(exe): add storage & contracts caching --- src/exe/cache.rs | 287 ++++++++++++++++++++++++++++++++++++----------- src/exe/mod.rs | 59 ++++------ 2 files changed, 243 insertions(+), 103 deletions(-) diff --git a/src/exe/cache.rs b/src/exe/cache.rs index e8487a3b..4b5c9161 100644 --- a/src/exe/cache.rs +++ b/src/exe/cache.rs @@ -1,105 +1,256 @@ +use blockifier::state::state_api::{State as BlockifierState, StateReader}; use ethers::types::U256; use lru::LruCache; use starknet_api::{core::ContractAddress, state::StorageKey}; use starknet_types_core::felt::Felt as StarkFelt; +use std::marker::PhantomData; use std::num::NonZeroUsize; use std::sync::{LazyLock, Mutex}; use crate::gen; -type Key = (U256, U256, U256); // block hash + contract address + storage key -type Value = U256; +mod storage { + use super::*; -const CACHE_SIZE: usize = 1024; + type Key = (U256, U256, U256); // block hash + contract address + storage key + type Value = StarkFelt; -static CACHE: LazyLock>> = LazyLock::new(|| { - Mutex::new(LruCache::new(NonZeroUsize::new(CACHE_SIZE).unwrap())) -}); + const SIZE: usize = 1024; -fn get(key: &Key) -> Option { - let mut guard = CACHE.lock().expect("cache-lock"); - guard.get(key).cloned() -} + static CACHE: LazyLock>> = LazyLock::new(|| { + Mutex::new(LruCache::new(NonZeroUsize::new(SIZE).unwrap())) + }); -fn set(key: Key, value: Value) -> Option { - let mut guard = CACHE.lock().expect("cache-lock"); - guard.put(key, value) -} + pub fn get(key: &Key) -> Option { + let mut guard = CACHE.lock().expect("storage-cache-lock"); + guard.get(key).cloned() + } -fn key( - block_hash: &gen::Felt, - contract_address: &ContractAddress, - storage_key: &StorageKey, -) -> Key { - ( - block_hash.as_ref().parse().unwrap(), - contract_address.0.key().to_bytes_be().into(), - storage_key.0.key().to_bytes_be().into(), - ) -} + pub fn set(key: Key, value: Value) -> Option { + let mut guard = CACHE.lock().expect("storage-cache-lock"); + guard.put(key, value) + } -pub trait StorageCache { - fn lookup( - &self, + pub fn key( block_hash: &gen::Felt, contract_address: &ContractAddress, storage_key: &StorageKey, - ) -> Option; - fn insert( - &self, + ) -> Key { + ( + block_hash.as_ref().parse().unwrap(), + contract_address.0.key().to_bytes_be().into(), + storage_key.0.key().to_bytes_be().into(), + ) + } +} + +mod class_hash { + use super::*; + + type Key = (U256, U256); // block hash + contract address + type Value = starknet_api::core::ClassHash; + + const SIZE: usize = 256; + + static CACHE: LazyLock>> = LazyLock::new(|| { + Mutex::new(LruCache::new(NonZeroUsize::new(SIZE).unwrap())) + }); + + pub fn get(key: &Key) -> Option { + let mut guard = CACHE.lock().expect("classhash-cache-lock"); + guard.get(key).cloned() + } + + pub fn set(key: Key, value: Value) -> Option { + let mut guard = CACHE.lock().expect("classhash-cache-lock"); + guard.put(key, value) + } + + pub fn key( block_hash: &gen::Felt, contract_address: &ContractAddress, - storage_key: &StorageKey, - val: &gen::Felt, - ); + ) -> Key { + ( + block_hash.as_ref().parse().unwrap(), + contract_address.0.key().to_bytes_be().into(), + ) + } } -pub struct Empty; +mod contract_class { + use super::*; + + type Key = (U256, U256); // block hash + class hash + type Value = blockifier::execution::contract_class::ContractClass; + + const SIZE: usize = 256; + + static CACHE: LazyLock>> = LazyLock::new(|| { + Mutex::new(LruCache::new(NonZeroUsize::new(SIZE).unwrap())) + }); + + pub fn get(key: &Key) -> Option { + let mut guard = CACHE.lock().expect("contractclass-cache-lock"); + guard.get(key).cloned() + } + + pub fn set(key: Key, value: Value) -> Option { + let mut guard = CACHE.lock().expect("contractclass-cache-lock"); + guard.put(key, value) + } + + pub fn key( + block_hash: &gen::Felt, + class_hash: &starknet_api::core::ClassHash, + ) -> Key { + ( + block_hash.as_ref().parse().unwrap(), + class_hash.0.to_bytes_be().into(), + ) + } +} -impl StorageCache for Empty { - fn lookup( +pub trait HasBlockHash { + fn get_block_hash(&self) -> &gen::Felt; +} + +pub struct CachedState { + inner: T, + _marker: PhantomData, +} + +impl CachedState { + pub fn new(inner: T) -> Self { + Self { inner, _marker: PhantomData } + } +} + +impl StateReader + for CachedState +{ + fn get_storage_at( &self, - _block_hash: &gen::Felt, - _contract_address: &ContractAddress, - _storage_key: &StorageKey, - ) -> Option { - None + contract_address: ContractAddress, + storage_key: StorageKey, + ) -> blockifier::state::state_api::StateResult { + let block_hash = self.inner.get_block_hash(); + if let Some(ret) = storage::get(&storage::key( + block_hash, + &contract_address, + &storage_key, + )) { + return Ok(ret); + } + let ret = self.inner.get_storage_at(contract_address, storage_key)?; + storage::set( + storage::key(block_hash, &contract_address, &storage_key), + ret, + ); + Ok(ret) } - fn insert( + fn get_nonce_at( &self, - _block_hash: &gen::Felt, - _contract_address: &ContractAddress, - _storage_key: &StorageKey, - _val: &gen::Felt, - ) { + contract_address: ContractAddress, + ) -> blockifier::state::state_api::StateResult + { + self.inner.get_nonce_at(contract_address) } -} -pub struct LRU; + fn get_class_hash_at( + &self, + contract_address: ContractAddress, + ) -> blockifier::state::state_api::StateResult + { + let block_hash = self.inner.get_block_hash(); + if let Some(ret) = + class_hash::get(&class_hash::key(block_hash, &contract_address)) + { + return Ok(ret); + } + let ret = self.inner.get_class_hash_at(contract_address)?; + class_hash::set(class_hash::key(block_hash, &contract_address), ret); + Ok(ret) + } -impl StorageCache for LRU { - fn lookup( + fn get_compiled_contract_class( &self, - block_hash: &gen::Felt, - contract_address: &ContractAddress, - storage_key: &StorageKey, - ) -> Option { - get(&key(block_hash, contract_address, storage_key)).map(|value| { - let mut bytes = [0u8; 32]; - value.to_big_endian(&mut bytes); - StarkFelt::from_bytes_be(&bytes) - }) + class_hash: starknet_api::core::ClassHash, + ) -> blockifier::state::state_api::StateResult< + blockifier::execution::contract_class::ContractClass, + > { + let block_hash = self.inner.get_block_hash(); + if let Some(ret) = + contract_class::get(&contract_class::key(block_hash, &class_hash)) + { + return Ok(ret); + } + let ret = self.inner.get_compiled_contract_class(class_hash)?; + contract_class::set( + contract_class::key(block_hash, &class_hash), + ret.clone(), + ); + Ok(ret) } - fn insert( + fn get_compiled_class_hash( &self, - block_hash: &gen::Felt, - contract_address: &ContractAddress, - storage_key: &StorageKey, - val: &gen::Felt, + class_hash: starknet_api::core::ClassHash, + ) -> blockifier::state::state_api::StateResult< + starknet_api::core::CompiledClassHash, + > { + self.inner.get_compiled_class_hash(class_hash) + } +} + +impl BlockifierState + for CachedState +{ + fn set_storage_at( + &mut self, + contract_address: ContractAddress, + storage_key: StorageKey, + value: StarkFelt, + ) -> blockifier::state::state_api::StateResult<()> { + self.inner.set_storage_at(contract_address, storage_key, value) + } + + fn increment_nonce( + &mut self, + contract_address: ContractAddress, + ) -> blockifier::state::state_api::StateResult<()> { + self.inner.increment_nonce(contract_address) + } + + fn set_class_hash_at( + &mut self, + contract_address: ContractAddress, + class_hash: starknet_api::core::ClassHash, + ) -> blockifier::state::state_api::StateResult<()> { + self.inner.set_class_hash_at(contract_address, class_hash) + } + + fn set_contract_class( + &mut self, + class_hash: starknet_api::core::ClassHash, + contract_class: blockifier::execution::contract_class::ContractClass, + ) -> blockifier::state::state_api::StateResult<()> { + self.inner.set_contract_class(class_hash, contract_class) + } + + fn set_compiled_class_hash( + &mut self, + class_hash: starknet_api::core::ClassHash, + compiled_class_hash: starknet_api::core::CompiledClassHash, + ) -> blockifier::state::state_api::StateResult<()> { + self.inner.set_compiled_class_hash(class_hash, compiled_class_hash) + } + + fn add_visited_pcs( + &mut self, + class_hash: starknet_api::core::ClassHash, + pcs: &std::collections::HashSet, ) { - let val = val.as_ref().parse().unwrap(); - set(key(block_hash, contract_address, storage_key), val); + self.inner.add_visited_pcs(class_hash, pcs); } } diff --git a/src/exe/mod.rs b/src/exe/mod.rs index 04158cb1..68d0fb1b 100644 --- a/src/exe/mod.rs +++ b/src/exe/mod.rs @@ -130,12 +130,13 @@ pub fn call( initial_gas: u64::MAX, }; - let cache = cache::LRU; - let mut proxy: StateProxy = StateProxy { client, state, cache }; + let state_proxy: StateProxy = StateProxy { client, state }; + let mut state_proxy = cache::CachedState::new(state_proxy); + let mut resources = Default::default(); let call_info = call_entry_point.execute( - &mut proxy, - &mut Default::default(), + &mut state_proxy, + &mut resources, &mut context, )?; @@ -143,16 +144,20 @@ pub fn call( Ok(call_info) } -struct StateProxy -{ +struct StateProxy { client: gen::client::blocking::Client, state: State, - cache: C, } -impl StateReader - for StateProxy +impl cache::HasBlockHash + for StateProxy { + fn get_block_hash(&self) -> &gen::Felt { + &self.state.block_hash + } +} + +impl StateReader for StateProxy { fn get_storage_at( &self, contract_address: ContractAddress, @@ -160,14 +165,6 @@ impl StateReader ) -> StateResult { tracing::info!(?contract_address, ?storage_key, "get_storage_at"); - if let Some(ret) = self.cache.lookup( - &self.state.block_hash, - &contract_address, - &storage_key, - ) { - return Ok(ret); - } - let felt: gen::Felt = contract_address.0.key().try_into()?; let address = gen::Address(felt); @@ -185,12 +182,6 @@ impl StateReader tracing::info!(?address, ?key, value=?ret, "get_storage_at"); if ret.as_ref() == "0x0" { - self.cache.insert( - &self.state.block_hash, - &contract_address, - &storage_key, - &ret, - ); tracing::info!("get_storage_at: skipping proof for zero value"); return Ok(ret.try_into()?); } @@ -210,12 +201,6 @@ impl StateReader })?; tracing::info!("get_storage_at: proof verified"); - self.cache.insert( - &self.state.block_hash, - &contract_address, - &storage_key, - &ret, - ); Ok(ret.try_into()?) } @@ -225,7 +210,9 @@ impl StateReader ) -> StateResult { tracing::info!(?contract_address, "get_nonce_at"); - let block_id = gen::BlockId::BlockTag(gen::BlockTag::Latest); + let block_id = gen::BlockId::BlockHash { + block_hash: gen::BlockHash(self.state.block_hash.clone()), + }; let felt: gen::Felt = contract_address.0.key().try_into()?; let contract_address = gen::Address(felt); @@ -244,7 +231,9 @@ impl StateReader ) -> StateResult { tracing::info!(?contract_address, "get_class_hash_at"); - let block_id = gen::BlockId::BlockTag(gen::BlockTag::Latest); + let block_id = gen::BlockId::BlockHash { + block_hash: gen::BlockHash(self.state.block_hash.clone()), + }; let felt: gen::Felt = contract_address.0.key().try_into()?; let contract_address = gen::Address(felt); @@ -263,7 +252,9 @@ impl StateReader ) -> StateResult { tracing::info!(?class_hash, "get_compiled_contract_class"); - let block_id = gen::BlockId::BlockTag(gen::BlockTag::Latest); + let block_id = gen::BlockId::BlockHash { + block_hash: gen::BlockHash(self.state.block_hash.clone()), + }; let class_hash: gen::Felt = class_hash.0.try_into()?; @@ -284,9 +275,7 @@ impl StateReader } } -impl - BlockifierState for StateProxy -{ +impl BlockifierState for StateProxy { fn set_storage_at( &mut self, contract_address: ContractAddress,