diff --git a/src/exe/cache.rs b/src/exe/cache.rs new file mode 100644 index 00000000..95112f8f --- /dev/null +++ b/src/exe/cache.rs @@ -0,0 +1,50 @@ +use std::{collections::HashMap, sync::{LazyLock, RwLock}}; +use starknet_api::{core::ContractAddress, state::StorageKey}; +use starknet_types_core::felt::Felt as StarkFelt; +use ethers::types::U256; + +use crate::gen; + +type Key = (U256, U256, U256); // block hash + contract address + storage key +type Value = U256; + +static CACHE: LazyLock>> = LazyLock::new(|| { + RwLock::new(HashMap::new()) +}); + +fn get(key: &Key) -> Option { + let guard = CACHE.read().expect("cache-rlock"); + guard.get(key).cloned() +} + +fn set(key: Key, value: Value) -> Option { + let mut guard = CACHE.write().expect("cache-wlock"); + guard.insert(key, value) +} + +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 trait StorageCache { + fn lookup(&self, block_hash: &gen::Felt, contract_address: &ContractAddress, storage_key: &StorageKey) -> Option; + fn insert(&self, block_hash: &gen::Felt, contract_address: &ContractAddress, storage_key: &StorageKey, val: &gen::Felt); +} + +pub struct NaiveUnboundedCache; + +impl StorageCache for NaiveUnboundedCache { + fn lookup(&self, block_hash: &gen::Felt, contract_address: &ContractAddress, storage_key: &StorageKey) -> Option { + get(&key(block_hash, contract_address, storage_key)) + .map(|value| StarkFelt::from_raw(value.0)) + } + + fn insert(&self, block_hash: &gen::Felt, contract_address: &ContractAddress, storage_key: &StorageKey, val: &gen::Felt) { + let val = val.as_ref().parse().unwrap(); + set(key(block_hash, contract_address, storage_key), val); + } +} diff --git a/src/exe/mod.rs b/src/exe/mod.rs index 8c682191..9609c443 100644 --- a/src/exe/mod.rs +++ b/src/exe/mod.rs @@ -42,6 +42,7 @@ use crate::{ pub mod err; pub mod map; +pub mod cache; use err::Error; @@ -129,7 +130,8 @@ pub fn call( initial_gas: u64::MAX, }; - let mut proxy: StateProxy = StateProxy { client, state }; + let cache = cache::NaiveUnboundedCache; + let mut proxy: StateProxy = StateProxy { client, state, cache }; let call_info = call_entry_point.execute( &mut proxy, @@ -141,23 +143,28 @@ pub fn call( Ok(call_info) } -struct StateProxy { +struct StateProxy { client: gen::client::blocking::Client, state: State, + cache: C, } -impl StateReader for StateProxy { +impl StateReader for StateProxy { fn get_storage_at( &self, contract_address: ContractAddress, - key: StarknetStorageKey, + storage_key: StarknetStorageKey, ) -> StateResult { - tracing::info!(?contract_address, ?key, "get_storage_at"); + 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 contract_address = gen::Address(felt); + let address = gen::Address(felt); - let key = gen::StorageKey::try_new(&key.0.to_string()) + let key = gen::StorageKey::try_new(&storage_key.0.to_string()) .map_err(Into::::into)?; let block_id = gen::BlockId::BlockHash { @@ -167,27 +174,28 @@ impl StateReader for StateProxy { let ret = self .client .getStorageAt( - contract_address.clone(), + address.clone(), key.clone(), block_id.clone(), ) .map_err(Into::::into)?; - tracing::info!(?contract_address, ?key, value=?ret, "get_storage_at"); + 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, &gen::Felt::try_new("0x0").unwrap()); tracing::info!("get_storage_at: skipping proof for zero value"); return Ok(ret.try_into()?); } let proof = self .client - .getProof(block_id, contract_address.clone(), vec![key.clone()]) + .getProof(block_id, address.clone(), vec![key.clone()]) .map_err(Into::::into)?; tracing::info!("get_storage_at: proof received"); let global_root = self.state.root.clone(); let value = ret.clone(); - proof.verify(global_root, contract_address, key, value).map_err( + proof.verify(global_root, address, key, value).map_err( |e| { StateError::StateReadError(format!( "Failed to verify merkle proof: {e:?}" @@ -196,6 +204,7 @@ impl StateReader for StateProxy { )?; tracing::info!("get_storage_at: proof verified"); + self.cache.insert(&self.state.block_hash, &contract_address, &storage_key, &ret); Ok(ret.try_into()?) } @@ -264,7 +273,7 @@ impl StateReader for StateProxy { } } -impl BlockifierState for StateProxy { +impl BlockifierState for StateProxy { fn set_storage_at( &mut self, contract_address: ContractAddress,