feat: added 'VisitedPcs' trait in blockifier

crates/blockifier/bench/blockifier_bench.rs

@@ -2,17 +2,42 @@
 //! various aspects related to transferring between accounts, including preparation
 //! and execution of transfers.
 //!
-//! The main benchmark function is `transfers_benchmark`, which measures the performance
-//! of transfers between randomly created accounts, which are iterated over round-robin.
+//! The benchmark function `transfers_benchmark` measures the performance of transfers between
+//! randomly created accounts, which are iterated over round-robin.
+//!
+//! The benchmark function `execution_benchmark` measures the performance of the method
+//! [`blockifier::transactions::transaction::ExecutableTransaction::execute`] by executing the entry
+//! point `advance_counter` of the test contract.
+//!
+//! The benchmark function `cached_state_benchmark` measures the performance of
+//! [`blockifier::state::cached_state::CachedState::add_visited_pcs`] method using a realistic size
+//! of data.
 //!
 //! Run the benchmarks using `cargo bench --bench blockifier_bench`.
 
+use std::time::Duration;
+
+use blockifier::context::BlockContext;
+use blockifier::state::cached_state::{CachedState, TransactionalState};
+use blockifier::state::state_api::State;
+use blockifier::state::visited_pcs::VisitedPcsSet;
+use blockifier::test_utils::contracts::FeatureContract;
+use blockifier::test_utils::dict_state_reader::DictStateReader;
+use blockifier::test_utils::initial_test_state::test_state;
 use blockifier::test_utils::transfers_generator::{
     RecipientGeneratorType,
     TransfersGenerator,
     TransfersGeneratorConfig,
 };
-use criterion::{criterion_group, criterion_main, Criterion};
+use blockifier::test_utils::{create_calldata, CairoVersion, BALANCE};
+use blockifier::transaction::account_transaction::AccountTransaction;
+use blockifier::transaction::test_utils::{account_invoke_tx, block_context, max_resource_bounds};
+use blockifier::transaction::transactions::ExecutableTransaction;
+use criterion::{criterion_group, criterion_main, BatchSize, Criterion};
+use starknet_api::core::ClassHash;
+use starknet_api::hash::StarkHash;
+use starknet_api::test_utils::NonceManager;
+use starknet_api::{felt, invoke_tx_args};
 
 pub fn transfers_benchmark(c: &mut Criterion) {
     let transfers_generator_config = TransfersGeneratorConfig {
@@ -29,5 +54,201 @@ pub fn transfers_benchmark(c: &mut Criterion) {
     });
 }
 
-criterion_group!(benches, transfers_benchmark);
+pub fn cached_state_benchmark(c: &mut Criterion) {
+    fn get_random_array(size: usize) -> Vec<usize> {
+        let mut vec: Vec<usize> = Vec::with_capacity(size);
+        for _ in 0..vec.capacity() {
+            vec.push(rand::random());
+        }
+        vec
+    }
+
+    fn create_class_hash(class_hash: &str) -> ClassHash {
+        ClassHash(StarkHash::from_hex_unchecked(class_hash))
+    }
+
+    // The state shared across all iterations.
+    let mut cached_state: CachedState<DictStateReader, VisitedPcsSet> = CachedState::default();
+
+    c.bench_function("cached_state", move |benchmark| {
+        benchmark.iter_batched(
+            || {
+                // This anonymous function creates the simulated visited program counters to add in
+                // `cached_state`.
+                // The numbers are taken from tx hash
+                // 0x0177C9365875CAA840EA8F03F97B0E3A8EE8851A8B952BF157B5DBD4FECCB060. This
+                // transaction has been chosen randomly, but it may not be representative of the
+                // average transaction on Starknet.
+
+                let mut class_hashes = Vec::new();
+                let mut random_arrays = Vec::new();
+
+                let class_hash = create_class_hash("a");
+                let random_array = get_random_array(11393);
+                class_hashes.push(class_hash);
+                random_arrays.push(random_array);
+
+                let class_hash = create_class_hash("a");
+                let random_array = get_random_array(453);
+                class_hashes.push(class_hash);
+                random_arrays.push(random_array);
+
+                let class_hash = create_class_hash("a");
+                let random_array = get_random_array(604);
+                class_hashes.push(class_hash);
+                random_arrays.push(random_array);
+
+                let class_hash = create_class_hash("a");
+                let random_array = get_random_array(806);
+                class_hashes.push(class_hash);
+                random_arrays.push(random_array);
+
+                let class_hash = create_class_hash("b");
+                let random_array = get_random_array(1327);
+                class_hashes.push(class_hash);
+                random_arrays.push(random_array);
+
+                let class_hash = create_class_hash("b");
+                let random_array = get_random_array(1135);
+                class_hashes.push(class_hash);
+                random_arrays.push(random_array);
+
+                let class_hash = create_class_hash("b");
+                let random_array = get_random_array(213);
+                class_hashes.push(class_hash);
+                random_arrays.push(random_array);
+
+                let class_hash = create_class_hash("b");
+                let random_array = get_random_array(135);
+                class_hashes.push(class_hash);
+                random_arrays.push(random_array);
+
+                let class_hash = create_class_hash("c");
+                let random_array = get_random_array(348);
+                class_hashes.push(class_hash);
+                random_arrays.push(random_array);
+
+                let class_hash = create_class_hash("c");
+                let random_array = get_random_array(88);
+                class_hashes.push(class_hash);
+                random_arrays.push(random_array);
+
+                let class_hash = create_class_hash("c");
+                let random_array = get_random_array(348);
+                class_hashes.push(class_hash);
+                random_arrays.push(random_array);
+
+                let class_hash = create_class_hash("c");
+                let random_array = get_random_array(348);
+                class_hashes.push(class_hash);
+                random_arrays.push(random_array);
+
+                let class_hash = create_class_hash("d");
+                let random_array = get_random_array(875);
+                class_hashes.push(class_hash);
+                random_arrays.push(random_array);
+
+                let class_hash = create_class_hash("d");
+                let random_array = get_random_array(450);
+                class_hashes.push(class_hash);
+                random_arrays.push(random_array);
+
+                let class_hash = create_class_hash("d");
+                let random_array = get_random_array(255);
+                class_hashes.push(class_hash);
+                random_arrays.push(random_array);
+
+                let class_hash = create_class_hash("d");
+                let random_array = get_random_array(210);
+                class_hashes.push(class_hash);
+                random_arrays.push(random_array);
+
+                let class_hash = create_class_hash("d");
+                let random_array = get_random_array(1403);
+                class_hashes.push(class_hash);
+                random_arrays.push(random_array);
+
+                let class_hash = create_class_hash("d");
+                let random_array = get_random_array(210);
+                class_hashes.push(class_hash);
+                random_arrays.push(random_array);
+
+                let class_hash = create_class_hash("d");
+                let random_array = get_random_array(210);
+                class_hashes.push(class_hash);
+                random_arrays.push(random_array);
+
+                let class_hash = create_class_hash("e");
+                let random_array = get_random_array(2386);
+                class_hashes.push(class_hash);
+                random_arrays.push(random_array);
+
+                let class_hash = create_class_hash("e");
+                let random_array = get_random_array(3602);
+                class_hashes.push(class_hash);
+                random_arrays.push(random_array);
+
+                (class_hashes, random_arrays)
+            },
+            |input_data| {
+                let mut transactional_state =
+                    TransactionalState::create_transactional(&mut cached_state);
+                for (class_hash, random_array) in input_data.0.into_iter().zip(input_data.1) {
+                    transactional_state.add_visited_pcs(class_hash, &random_array);
+                }
+                transactional_state.commit();
+            },
+            BatchSize::SmallInput,
+        )
+    });
+}
+
+pub fn execution_benchmark(c: &mut Criterion) {
+    /// This function sets up and returns all the objects required to execute an invoke transaction.
+    fn prepare_account_tx()
+    -> (AccountTransaction, CachedState<DictStateReader, VisitedPcsSet>, BlockContext) {
+        let block_context = block_context();
+        let max_resource_bounds = max_resource_bounds();
+        let cairo_version = CairoVersion::Cairo1;
+        let account = FeatureContract::AccountWithoutValidations(cairo_version);
+        let test_contract = FeatureContract::TestContract(cairo_version);
+        let state =
+            test_state(block_context.chain_info(), BALANCE, &[(account, 1), (test_contract, 1)]);
+        let account_address = account.get_instance_address(0);
+        let contract_address = test_contract.get_instance_address(0);
+        let index = felt!(123_u32);
+        let base_tx_args = invoke_tx_args! {
+            resource_bounds: max_resource_bounds,
+            sender_address: account_address,
+        };
+
+        let mut nonce_manager = NonceManager::default();
+        let counter_diffs = [101_u32, 102_u32];
+        let initial_counters = [felt!(counter_diffs[0]), felt!(counter_diffs[1])];
+        let calldata_args = vec![index, initial_counters[0], initial_counters[1]];
+
+        let account_tx = account_invoke_tx(invoke_tx_args! {
+            nonce: nonce_manager.next(account_address),
+            calldata:
+                create_calldata(contract_address, "advance_counter", &calldata_args),
+            ..base_tx_args
+        });
+        (account_tx, state, block_context)
+    }
+    c.bench_function("execution", move |benchmark| {
+        benchmark.iter_batched(
+            prepare_account_tx,
+            |(account_tx, mut state, block_context)| {
+                account_tx.execute(&mut state, &block_context, true, true).unwrap()
+            },
+            BatchSize::SmallInput,
+        )
+    });
+}
+
+criterion_group! {
+    name = benches;
+    config = Criterion::default().measurement_time(Duration::from_secs(20));
+    targets = transfers_benchmark, execution_benchmark, cached_state_benchmark
+}
 criterion_main!(benches);

crates/blockifier/src/blockifier/stateful_validator.rs

@@ -17,6 +17,7 @@ use crate::fee::receipt::TransactionReceipt;
 use crate::state::cached_state::CachedState;
 use crate::state::errors::StateError;
 use crate::state::state_api::StateReader;
+use crate::state::visited_pcs::VisitedPcs;
 use crate::transaction::account_transaction::AccountTransaction;
 use crate::transaction::errors::{TransactionExecutionError, TransactionPreValidationError};
 use crate::transaction::transaction_execution::Transaction;
@@ -41,12 +42,12 @@ pub enum StatefulValidatorError {
 pub type StatefulValidatorResult<T> = Result<T, StatefulValidatorError>;
 
 /// Manages state related transaction validations for pre-execution flows.
-pub struct StatefulValidator<S: StateReader> {
-    tx_executor: TransactionExecutor<S>,
+pub struct StatefulValidator<S: StateReader, V: VisitedPcs> {
+    tx_executor: TransactionExecutor<S, V>,
 }
 
-impl<S: StateReader> StatefulValidator<S> {
-    pub fn create(state: CachedState<S>, block_context: BlockContext) -> Self {
+impl<S: StateReader, V: VisitedPcs> StatefulValidator<S, V> {
+    pub fn create(state: CachedState<S, V>, block_context: BlockContext) -> Self {
         let tx_executor =
             TransactionExecutor::new(state, block_context, TransactionExecutorConfig::default());
         Self { tx_executor }

crates/blockifier/src/blockifier/transaction_executor.rs

@@ -1,6 +1,4 @@
 #[cfg(feature = "concurrency")]
-use std::collections::{HashMap, HashSet};
-#[cfg(feature = "concurrency")]
 use std::panic::{self, catch_unwind, AssertUnwindSafe};
 #[cfg(feature = "concurrency")]
 use std::sync::Arc;
@@ -20,6 +18,7 @@ use crate::context::BlockContext;
 use crate::state::cached_state::{CachedState, CommitmentStateDiff, TransactionalState};
 use crate::state::errors::StateError;
 use crate::state::state_api::StateReader;
+use crate::state::visited_pcs::VisitedPcs;
 use crate::transaction::errors::TransactionExecutionError;
 use crate::transaction::objects::TransactionExecutionInfo;
 use crate::transaction::transaction_execution::Transaction;
@@ -45,7 +44,7 @@ pub type TransactionExecutorResult<T> = Result<T, TransactionExecutorError>;
 pub type VisitedSegmentsMapping = Vec<(ClassHash, Vec<usize>)>;
 
 // TODO(Gilad): make this hold TransactionContext instead of BlockContext.
-pub struct TransactionExecutor<S: StateReader> {
+pub struct TransactionExecutor<S: StateReader, V: VisitedPcs> {
     pub block_context: BlockContext,
     pub bouncer: Bouncer,
     // Note: this config must not affect the execution result (e.g. state diff and traces).
@@ -56,12 +55,12 @@ pub struct TransactionExecutor<S: StateReader> {
     // block state to the worker executor - operating at the chunk level - and gets it back after
     // committing the chunk. The block state is wrapped with an Option<_> to allow setting it to
     // `None` while it is moved to the worker executor.
-    pub block_state: Option<CachedState<S>>,
+    pub block_state: Option<CachedState<S, V>>,
 }
 
-impl<S: StateReader> TransactionExecutor<S> {
+impl<S: StateReader, V: VisitedPcs> TransactionExecutor<S, V> {
     pub fn new(
-        block_state: CachedState<S>,
+        block_state: CachedState<S, V>,
         block_context: BlockContext,
         config: TransactionExecutorConfig,
     ) -> Self {
@@ -159,7 +158,8 @@ impl<S: StateReader> TransactionExecutor<S> {
                     .as_ref()
                     .expect(BLOCK_STATE_ACCESS_ERR)
                     .get_compiled_contract_class(*class_hash)?;
-                Ok((*class_hash, contract_class.get_visited_segments(class_visited_pcs)?))
+                let class_visited_pcs = V::to_set(class_visited_pcs.clone());
+                Ok((*class_hash, contract_class.get_visited_segments(&class_visited_pcs)?))
             })
             .collect::<TransactionExecutorResult<_>>()?;
 
@@ -172,7 +172,11 @@ impl<S: StateReader> TransactionExecutor<S> {
     }
 }
 
-impl<S: StateReader + Send + Sync> TransactionExecutor<S> {
+impl<S, V> TransactionExecutor<S, V>
+where
+    S: StateReader + Send + Sync,
+    V: VisitedPcs + Send + Sync,
+{
     /// Executes the given transactions on the state maintained by the executor.
     /// Stops if and when there is no more room in the block, and returns the executed transactions'
     /// results.
@@ -221,6 +225,7 @@ impl<S: StateReader + Send + Sync> TransactionExecutor<S> {
         chunk: &[Transaction],
     ) -> Vec<TransactionExecutorResult<TransactionExecutionInfo>> {
         use crate::concurrency::utils::AbortIfPanic;
+        use crate::concurrency::worker_logic::ExecutionTaskOutput;
 
         let block_state = self.block_state.take().expect("The block state should be `Some`.");
 
@@ -264,20 +269,20 @@ impl<S: StateReader + Send + Sync> TransactionExecutor<S> {
 
         let n_committed_txs = worker_executor.scheduler.get_n_committed_txs();
         let mut tx_execution_results = Vec::new();
-        let mut visited_pcs: HashMap<ClassHash, HashSet<usize>> = HashMap::new();
+        let mut visited_pcs: V = V::new();
         for execution_output in worker_executor.execution_outputs.iter() {
             if tx_execution_results.len() >= n_committed_txs {
                 break;
             }
-            let locked_execution_output = execution_output
+            let locked_execution_output: ExecutionTaskOutput<V> = execution_output
                 .lock()
                 .expect("Failed to lock execution output.")
                 .take()
                 .expect("Output must be ready.");
             tx_execution_results
                 .push(locked_execution_output.result.map_err(TransactionExecutorError::from));
-            for (class_hash, class_visited_pcs) in locked_execution_output.visited_pcs {
-                visited_pcs.entry(class_hash).or_default().extend(class_visited_pcs);
+            for (class_hash, class_visited_pcs) in locked_execution_output.visited_pcs.iter() {
+                visited_pcs.extend(class_hash, class_visited_pcs);
             }
         }
 

crates/blockifier/src/blockifier/transaction_executor_test.rs

@@ -16,6 +16,7 @@ use crate::bouncer::{Bouncer, BouncerWeights};
 use crate::context::BlockContext;
 use crate::state::cached_state::CachedState;
 use crate::state::state_api::StateReader;
+use crate::state::visited_pcs::VisitedPcs;
 use crate::test_utils::contracts::FeatureContract;
 use crate::test_utils::declare::declare_tx;
 use crate::test_utils::deploy_account::deploy_account_tx;
@@ -35,8 +36,8 @@ use crate::transaction::test_utils::{
 use crate::transaction::transaction_execution::Transaction;
 use crate::transaction::transactions::L1HandlerTransaction;
 
-fn tx_executor_test_body<S: StateReader>(
-    state: CachedState<S>,
+fn tx_executor_test_body<S: StateReader, V: VisitedPcs>(
+    state: CachedState<S, V>,
     block_context: BlockContext,
     tx: Transaction,
     expected_bouncer_weights: BouncerWeights,