lambdaclass · juanbono · Nov 17, 2023 · Aug 23, 2023 · Aug 23, 2023 · Aug 24, 2023
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -63,6 +63,7 @@ tracing = "0.1.37"
 [dev-dependencies]
 assert_matches = "1.5.0"
 coverage-helper = "0.2.0"
+lru = "0.11.0"
 pretty_assertions_sorted = "1.2.3"
 tracing-subscriber = { version = "0.3.17", features = ["env-filter"] }
 

diff --git a/README.md b/README.md
@@ -117,6 +117,54 @@ You can find an example on how to use the CLI [here](/docs/CLI_USAGE_EXAMPLE.md)
 
 ### Customization
 
+#### Contract class cache behavior
+
+`starknet_in_rust` supports caching contracts in memory. Caching the contracts is useful for
+avoiding excessive RPC API usage and keeping the contract class deserialization overhead to the
+minimum. The project provides two builtin cache policies: null and permanent. The null cache behaves
+as if there was no cache at all. The permanent cache caches everything in memory forever.
+
+In addition to those two, an example is provided that implements and uses an LRU cache policy.
+Long-running applications should ideally implement a cache algorithm suited to their needs or
+alternatively use our example's implementation to avoid spamming the API when using the null cache
+or blowing the memory usage when running with the permanent cache.
+
+Customized cache policies may be used by implementing the `ContractClassCache` trait. Check out our
+[LRU cache example](examples/lru_cache/main.rs) for more details. Updating the cache requires
+manually merging the local state cache into the shared cache manually. This can be done by calling
+the `drain_private_contract_class_cache` on the `CachedState` instance.
+
+```rs
+// To use the null cache (aka. no cache at all), create the state as follows:
+let cache = Arc::new(NullContractClassCache::default());
+let state1 = CachedState::new(state_reader.clone(), cache.clone());
+let state2 = CachedState::new(state_reader.clone(), cache.clone()); // Cache is reused.
+
+// Insert state usage here.
+
+// The null cache doesn't have any method to extend it since it has no data.
+```
+
+```rs
+// If the permanent cache is preferred, then use `PermanentContractClassCache` instead:
+let cache = Arc::new(PermanentContractClassCache::default());
+let state1 = CachedState::new(state_reader.clone(), cache.clone());
+let state2 = CachedState::new(state_reader.clone(), cache.clone()); // Cache is reused.
+
+// Insert state usage here.
+
+// Extend the shared cache with the states' contracts after using them.
+cache.extend(state1.state.drain_private_contract_class_cache());
+cache.extend(state2.state.drain_private_contract_class_cache());
+```
+
+#### Logging configuration
+
+This project uses the [`tracing`](https://crates.io/crates/tracing) crate as a library. Check out
+its documentation for more information.
+
+### Testing
+
 #### Logging configuration
 
 This project uses the [`tracing`](https://crates.io/crates/tracing) crate as a library. Check out

diff --git a/bench/internals.rs b/bench/internals.rs
@@ -15,12 +15,15 @@ use starknet_in_rust::{
     services::api::contract_classes::{
         compiled_class::CompiledClass, deprecated_contract_class::ContractClass,
     },
-    state::in_memory_state_reader::InMemoryStateReader,
     state::{cached_state::CachedState, state_api::State},
+    state::{
+        contract_class_cache::PermanentContractClassCache,
+        in_memory_state_reader::InMemoryStateReader,
+    },
     transaction::{declare::Declare, Deploy, DeployAccount, InvokeFunction},
     utils::{Address, ClassHash},
 };
-use std::{collections::HashMap, hint::black_box, sync::Arc};
+use std::{hint::black_box, sync::Arc};
 
 #[cfg(feature = "cairo-native")]
 use std::{cell::RefCell, rc::Rc};
@@ -74,7 +77,10 @@ fn deploy_account(
     const RUNS: usize = 500;
 
     let state_reader = Arc::new(InMemoryStateReader::default());
-    let mut state = CachedState::new(state_reader, HashMap::new());
+    let mut state = CachedState::new(
+        state_reader,
+        Arc::new(PermanentContractClassCache::default()),
+    );
 
     state
         .set_contract_class(
@@ -86,7 +92,7 @@ fn deploy_account(
     let block_context = &Default::default();
 
     for _ in 0..RUNS {
-        let mut state_copy = state.clone();
+        let mut state_copy = state.clone_for_testing();
         let class_hash = *CLASS_HASH;
         let signature = SIGNATURE.clone();
         scope(|| {
@@ -118,12 +124,15 @@ fn declare(#[cfg(feature = "cairo-native")] program_cache: Rc<RefCell<ProgramCac
     const RUNS: usize = 5;
 
     let state_reader = Arc::new(InMemoryStateReader::default());
-    let state = CachedState::new(state_reader, HashMap::new());
+    let state = CachedState::new(
+        state_reader,
+        Arc::new(PermanentContractClassCache::default()),
+    );
 
     let block_context = &Default::default();
 
     for _ in 0..RUNS {
-        let mut cloned_state = state.clone();
+        let mut cloned_state = state.clone_for_testing();
         let class = CONTRACT_CLASS.clone();
         let address = CONTRACT_ADDRESS.clone();
         scope(|| {
@@ -155,7 +164,10 @@ fn deploy(#[cfg(feature = "cairo-native")] program_cache: Rc<RefCell<ProgramCach
     const RUNS: usize = 8;
 
     let state_reader = Arc::new(InMemoryStateReader::default());
-    let mut state = CachedState::new(state_reader, HashMap::new());
+    let mut state = CachedState::new(
+        state_reader,
+        Arc::new(PermanentContractClassCache::default()),
+    );
 
     state
         .set_contract_class(
@@ -167,7 +179,7 @@ fn deploy(#[cfg(feature = "cairo-native")] program_cache: Rc<RefCell<ProgramCach
     let block_context = &Default::default();
 
     for _ in 0..RUNS {
-        let mut state_copy = state.clone();
+        let mut state_copy = state.clone_for_testing();
         let salt = felt_str!(
             "2669425616857739096022668060305620640217901643963991674344872184515580705509"
         );
@@ -198,7 +210,10 @@ fn invoke(#[cfg(feature = "cairo-native")] program_cache: Rc<RefCell<ProgramCach
     const RUNS: usize = 100;
 
     let state_reader = Arc::new(InMemoryStateReader::default());
-    let mut state = CachedState::new(state_reader, HashMap::new());
+    let mut state = CachedState::new(
+        state_reader,
+        Arc::new(PermanentContractClassCache::default()),
+    );
 
     state
         .set_contract_class(
@@ -231,7 +246,7 @@ fn invoke(#[cfg(feature = "cairo-native")] program_cache: Rc<RefCell<ProgramCach
         .unwrap();
 
     for _ in 0..RUNS {
-        let mut state_copy = state.clone();
+        let mut state_copy = state.clone_for_testing();
         let address = CONTRACT_ADDRESS.clone();
         let selector = VALIDATE_ENTRY_POINT_SELECTOR.clone();
         let signature = SIGNATURE.clone();

diff --git a/bench/native_bench.rs b/bench/native_bench.rs
@@ -16,6 +16,8 @@ use num_traits::Zero;
 use starknet_in_rust::definitions::block_context::BlockContext;
 use starknet_in_rust::definitions::block_context::StarknetChainId;
 use starknet_in_rust::services::api::contract_classes::compiled_class::CompiledClass;
+use starknet_in_rust::state::contract_class_cache::ContractClassCache;
+use starknet_in_rust::state::contract_class_cache::PermanentContractClassCache;
 use starknet_in_rust::state::state_api::State;
 use starknet_in_rust::transaction::DeployAccount;
 use starknet_in_rust::utils::calculate_sn_keccak;
@@ -32,7 +34,6 @@ use starknet_in_rust::{
     utils::{Address, ClassHash},
 };
 use std::cell::RefCell;
-use std::collections::HashMap;
 use std::rc::Rc;
 use std::sync::Arc;
 
@@ -62,7 +63,7 @@ pub fn main() {
 
 fn bench_fibo(executions: usize, native: bool) {
     // Create state reader with class hash data
-    let mut contract_class_cache = HashMap::new();
+    let contract_class_cache = PermanentContractClassCache::default();
     static CASM_CLASS_HASH: ClassHash = ClassHash([2; 32]);
 
     let (contract_class, constructor_selector) = match native {
@@ -96,7 +97,7 @@ fn bench_fibo(executions: usize, native: bool) {
 
     let caller_address = Address(123456789.into());
 
-    contract_class_cache.insert(CASM_CLASS_HASH, contract_class);
+    contract_class_cache.set_contract_class(CASM_CLASS_HASH, contract_class);
     let mut state_reader = InMemoryStateReader::default();
     let nonce = Felt252::zero();
 
@@ -109,7 +110,7 @@ fn bench_fibo(executions: usize, native: bool) {
 
     // Create state from the state_reader and contract cache.
     let state_reader = Arc::new(state_reader);
-    let state = CachedState::new(state_reader, contract_class_cache);
+    let state = CachedState::new(state_reader, Arc::new(contract_class_cache));
 
     /* f0, f1, N */
     let mut calldata = [1.into(), 1.into(), 2000000.into()];
@@ -120,7 +121,7 @@ fn bench_fibo(executions: usize, native: bool) {
     for _ in 0..executions {
         calldata[2] = &calldata[2] + 1usize;
         let result = execute(
-            &mut state.clone(),
+            &mut state.clone_for_testing(),
             &caller_address,
             &caller_address,
             &Felt252::new(constructor_selector.clone()),
@@ -136,7 +137,7 @@ fn bench_fibo(executions: usize, native: bool) {
 
 fn bench_fact(executions: usize, native: bool) {
     // Create state reader with class hash data
-    let mut contract_class_cache = HashMap::new();
+    let contract_class_cache = PermanentContractClassCache::default();
     static CASM_CLASS_HASH: ClassHash = ClassHash([2; 32]);
 
     let (contract_class, constructor_selector) = match native {
@@ -172,7 +173,7 @@ fn bench_fact(executions: usize, native: bool) {
     // FACT 1M
     // FIBO 2M
 
-    contract_class_cache.insert(CASM_CLASS_HASH, contract_class);
+    contract_class_cache.set_contract_class(CASM_CLASS_HASH, contract_class);
     let mut state_reader = InMemoryStateReader::default();
     let nonce = Felt252::zero();
 
@@ -185,7 +186,7 @@ fn bench_fact(executions: usize, native: bool) {
 
     // Create state from the state_reader and contract cache.
     let state_reader = Arc::new(state_reader);
-    let state = CachedState::new(state_reader, contract_class_cache);
+    let state = CachedState::new(state_reader, Arc::new(contract_class_cache));
 
     /* N */
     let mut calldata = [2000000.into()];
@@ -196,7 +197,7 @@ fn bench_fact(executions: usize, native: bool) {
     for _ in 0..executions {
         calldata[0] = &calldata[0] + 1usize;
         let result = execute(
-            &mut state.clone(),
+            &mut state.clone_for_testing(),
             &caller_address,
             &caller_address,
             &Felt252::new(constructor_selector.clone()),
@@ -213,7 +214,7 @@ fn bench_fact(executions: usize, native: bool) {
 fn bench_erc20(executions: usize, native: bool) {
     // 1. setup ERC20 contract and state.
     // Create state reader and preload the contract classes.
-    let mut contract_class_cache = HashMap::new();
+    let contract_class_cache = PermanentContractClassCache::default();
 
     lazy_static! {
         static ref ERC20_CLASS_HASH: ClassHash = ClassHash::from(felt_str!("2"));
@@ -239,7 +240,10 @@ fn bench_erc20(executions: usize, native: bool) {
     }
 
     let program_cache = Rc::new(RefCell::new(ProgramCache::new(get_native_context())));
-    let (erc20_address, mut state): (Address, CachedState<InMemoryStateReader>) = match native {
+    let (erc20_address, mut state): (
+        Address,
+        CachedState<InMemoryStateReader, PermanentContractClassCache>,
+    ) = match native {
         true => {
             let erc20_sierra_class = include_bytes!("../starknet_programs/cairo2/erc20.sierra");
             let sierra_contract_class: cairo_lang_starknet::contract_class::ContractClass =
@@ -259,11 +263,11 @@ fn bench_erc20(executions: usize, native: bool) {
             let deploy_entrypoint_selector = &entrypoints.external.get(0).unwrap().selector;
 
             // insert deployer and erc20 classes into the cache.
-            contract_class_cache.insert(
+            contract_class_cache.set_contract_class(
                 *DEPLOYER_CLASS_HASH,
                 CompiledClass::Casm(Arc::new(erc20_deployer_class)),
             );
-            contract_class_cache.insert(*ERC20_CLASS_HASH, erc20_contract_class);
+            contract_class_cache.set_contract_class(*ERC20_CLASS_HASH, erc20_contract_class);
 
             let mut state_reader = InMemoryStateReader::default();
             // setup deployer nonce and address into the state reader
@@ -275,7 +279,8 @@ fn bench_erc20(executions: usize, native: bool) {
                 .insert(DEPLOYER_ADDRESS.clone(), Felt252::zero());
 
             // Create state from the state_reader and contract cache.
-            let mut state = CachedState::new(Arc::new(state_reader), contract_class_cache);
+            let mut state =
+                CachedState::new(Arc::new(state_reader), Arc::new(contract_class_cache));
 
             // deploy the erc20 contract by calling the deployer contract.
 
@@ -338,11 +343,11 @@ fn bench_erc20(executions: usize, native: bool) {
             let deploy_entrypoint_selector = &entrypoints.external.get(0).unwrap().selector;
 
             // insert deployer and erc20 classes into the cache.
-            contract_class_cache.insert(
+            contract_class_cache.set_contract_class(
                 *DEPLOYER_CLASS_HASH,
                 CompiledClass::Casm(Arc::new(erc20_deployer_class)),
             );
-            contract_class_cache.insert(*ERC20_CLASS_HASH, erc20_contract_class);
+            contract_class_cache.set_contract_class(*ERC20_CLASS_HASH, erc20_contract_class);
 
             let mut state_reader = InMemoryStateReader::default();
             // setup deployer nonce and address into the state reader
@@ -354,7 +359,8 @@ fn bench_erc20(executions: usize, native: bool) {
                 .insert(DEPLOYER_ADDRESS.clone(), Felt252::zero());
 
             // Create state from the state_reader and contract cache.
-            let mut state = CachedState::new(Arc::new(state_reader), contract_class_cache);
+            let mut state =
+                CachedState::new(Arc::new(state_reader), Arc::new(contract_class_cache));
 
             // deploy the erc20 contract by calling the deployer contract.
 
@@ -493,7 +499,7 @@ fn bench_erc20(executions: usize, native: bool) {
 
     for _ in 0..executions {
         let result = execute(
-            &mut state.clone(),
+            &mut state.clone_for_testing(),
             &account1_address,
             &erc20_address,
             &transfer_entrypoint_selector.clone(),
@@ -510,7 +516,7 @@ fn bench_erc20(executions: usize, native: bool) {
 #[inline(never)]
 #[allow(clippy::too_many_arguments)]
 fn execute(
-    state: &mut CachedState<InMemoryStateReader>,
+    state: &mut CachedState<InMemoryStateReader, PermanentContractClassCache>,
     caller_address: &Address,
     callee_address: &Address,
     selector: &Felt252,