Merge pull request #31 from lambdaclass/benchmark-feature

Benchmark feature

.github/workflows/ci.yml

@@ -34,7 +34,7 @@ jobs:
       - name: Install LLVM
         run: sudo apt-get install llvm-18 llvm-18-dev llvm-18-runtime clang-18 clang-tools-18 lld-18 libpolly-18-dev libmlir-18-dev mlir-18-tools
       - name: Clippy
-        run: cargo clippy --all-targets -- -D warnings
+        run: cargo clippy --all-targets --all-features -- -D warnings
 
   format:
     name: rustfmt
@@ -164,4 +164,4 @@ jobs:
           make runtime
           echo "CAIRO_NATIVE_RUNTIME_LIBRARY=$(pwd)/libcairo_native_runtime.a" > $GITHUB_ENV
       - name: Test
-        run: cargo test
+        run: cargo test --all-features

replay/src/benchmark.rs

@@ -0,0 +1,169 @@
+use blockifier::{
+    context::BlockContext,
+    state::{cached_state::CachedState, state_api::StateReader},
+};
+use rpc_state_reader::{
+    blockifier_state_reader::{execute_tx_with_blockifier, fetch_block_context, RpcStateReader},
+    rpc_state::{RpcChain, RpcState},
+};
+use starknet_api::{
+    block::BlockNumber,
+    hash::StarkFelt,
+    transaction::{Transaction as SNTransaction, TransactionHash},
+};
+use tracing::{error, info};
+
+pub type BlockCachedData = (
+    CachedState<OptionalStateReader<RpcStateReader>>,
+    BlockContext,
+    Vec<(TransactionHash, SNTransaction)>,
+);
+
+/// Fetches context data to execute the given block range
+///
+/// It does not actually execute the block range, so not all data required
+/// by blockifier will be cached. To ensure that all rpc data is cached,
+/// the block range must be executed once.
+///
+/// See `execute_block_range` to execute the block range
+pub fn fetch_block_range_data(
+    block_start: BlockNumber,
+    block_end: BlockNumber,
+    chain: RpcChain,
+) -> Vec<BlockCachedData> {
+    let mut block_caches = Vec::new();
+
+    for block_number in block_start.0..=block_end.0 {
+        // For each block
+        let block_number = BlockNumber(block_number);
+
+        let rpc_state = RpcState::new_rpc(chain, block_number.into()).unwrap();
+
+        // Fetch block context
+        let block_context = fetch_block_context(&rpc_state, block_number);
+
+        // Fetch transactions for the block
+        let transactions = rpc_state
+            .get_transaction_hashes()
+            .unwrap()
+            .into_iter()
+            .map(|transaction_hash| {
+                let transaction_hash = TransactionHash(
+                    StarkFelt::try_from(transaction_hash.strip_prefix("0x").unwrap()).unwrap(),
+                );
+
+                // Fetch transaction
+                let transaction = rpc_state.get_transaction(&transaction_hash).unwrap();
+
+                (transaction_hash, transaction)
+            })
+            .collect::<Vec<_>>();
+
+        // Create cached state
+        let previous_rpc_state =
+            RpcState::new_rpc(chain, block_number.prev().unwrap().into()).unwrap();
+        let previous_rpc_state_reader = RpcStateReader::new(previous_rpc_state);
+        let cached_state = CachedState::new(OptionalStateReader::new(previous_rpc_state_reader));
+
+        block_caches.push((cached_state, block_context, transactions));
+    }
+
+    block_caches
+}
+
+/// Executes the given block range, discarding any state changes applied to it
+///
+/// Can also be used to fill up the cache
+pub fn execute_block_range(block_range_data: &mut Vec<BlockCachedData>) {
+    for (state, block_context, transactions) in block_range_data {
+        // For each block
+
+        // The transactional state is used to execute a transaction while discarding state changes applied to it.
+        let mut transactional_state = CachedState::create_transactional(state);
+
+        for (transaction_hash, transaction) in transactions {
+            // Execute each transaction
+            let result = execute_tx_with_blockifier(
+                &mut transactional_state,
+                block_context.clone(),
+                transaction.to_owned(),
+                transaction_hash.to_owned(),
+            );
+
+            match result {
+                Ok(info) => {
+                    info!(
+                        transaction_hash = transaction_hash.to_string(),
+                        succeeded = info.revert_error.is_none(),
+                        "tx execution status"
+                    )
+                }
+                Err(_) => error!(
+                    transaction_hash = transaction_hash.to_string(),
+                    "tx execution failed"
+                ),
+            }
+        }
+    }
+}
+
+/// An implementation of StateReader that can be disabled, panicking if atempted to be read from
+///
+/// Used to ensure that no requests are made after disabling it.
+pub struct OptionalStateReader<S: StateReader>(pub Option<S>);
+
+impl<S: StateReader> OptionalStateReader<S> {
+    pub fn new(state_reader: S) -> Self {
+        Self(Some(state_reader))
+    }
+
+    pub fn get_inner(&self) -> &S {
+        self.0
+            .as_ref()
+            .expect("atempted to read from a disabled state reader")
+    }
+
+    pub fn disable(&mut self) {
+        self.0 = None;
+    }
+}
+
+impl<S: StateReader> StateReader for OptionalStateReader<S> {
+    fn get_storage_at(
+        &self,
+        contract_address: starknet_api::core::ContractAddress,
+        key: starknet_api::state::StorageKey,
+    ) -> blockifier::state::state_api::StateResult<StarkFelt> {
+        self.get_inner().get_storage_at(contract_address, key)
+    }
+
+    fn get_nonce_at(
+        &self,
+        contract_address: starknet_api::core::ContractAddress,
+    ) -> blockifier::state::state_api::StateResult<starknet_api::core::Nonce> {
+        self.get_inner().get_nonce_at(contract_address)
+    }
+
+    fn get_class_hash_at(
+        &self,
+        contract_address: starknet_api::core::ContractAddress,
+    ) -> blockifier::state::state_api::StateResult<starknet_api::core::ClassHash> {
+        self.get_inner().get_class_hash_at(contract_address)
+    }
+
+    fn get_compiled_contract_class(
+        &self,
+        class_hash: starknet_api::core::ClassHash,
+    ) -> blockifier::state::state_api::StateResult<
+        blockifier::execution::contract_class::ContractClass,
+    > {
+        self.get_inner().get_compiled_contract_class(class_hash)
+    }
+
+    fn get_compiled_class_hash(
+        &self,
+        class_hash: starknet_api::core::ClassHash,
+    ) -> blockifier::state::state_api::StateResult<starknet_api::core::CompiledClassHash> {
+        self.get_inner().get_compiled_class_hash(class_hash)
+    }
+}