Added VisitedPcsSet as abstraction layer to visited_pcs in

`CachedState`.

crates/blockifier/src/blockifier/transaction_executor.rs

@@ -219,7 +219,7 @@ impl<S: StateReader + Send + Sync> TransactionExecutor<S> {
         chunk: &[Transaction],
     ) -> Vec<TransactionExecutorResult<TransactionExecutionInfo>> {
         use crate::concurrency::utils::AbortIfPanic;
-        use crate::state::cached_state::VisitedPcs;
+        use crate::state::visited_pcs::VisitedPcsSet;
 
         let block_state = self.block_state.take().expect("The block state should be `Some`.");
 
@@ -263,7 +263,7 @@ 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: VisitedPcs = VisitedPcs::new();
+        let mut visited_pcs: VisitedPcsSet = VisitedPcsSet::new();
         for execution_output in worker_executor.execution_outputs.iter() {
             if tx_execution_results.len() >= n_committed_txs {
                 break;

crates/blockifier/src/concurrency/fee_utils.rs

@@ -10,6 +10,7 @@ use crate::execution::call_info::CallInfo;
 use crate::fee::fee_utils::get_sequencer_balance_keys;
 use crate::state::cached_state::{ContractClassMapping, StateMaps};
 use crate::state::state_api::UpdatableState;
+use crate::state::visited_pcs::VisitedPcsSet;
 use crate::transaction::objects::TransactionExecutionInfo;
 
 #[cfg(test)]
@@ -120,5 +121,5 @@ pub fn add_fee_to_sequencer_balance(
         ]),
         ..StateMaps::default()
     };
-    state.apply_writes(&writes, &ContractClassMapping::default(), &HashMap::default());
+    state.apply_writes(&writes, &ContractClassMapping::default(), &VisitedPcsSet::default());
 }

crates/blockifier/src/concurrency/flow_test.rs

@@ -27,6 +27,8 @@ fn scheduler_flow_test(
     // transactions with multiple threads, where every transaction depends on its predecessor. Each
     // transaction sequentially advances a counter by reading the previous value and bumping it by
     // 1.
+
+    use crate::state::visited_pcs::VisitedPcsSet;
     let scheduler = Arc::new(Scheduler::new(DEFAULT_CHUNK_SIZE));
     let versioned_state =
         safe_versioned_state_for_testing(CachedState::from(DictStateReader::default()));
@@ -53,7 +55,7 @@ fn scheduler_flow_test(
                             state_proxy.apply_writes(
                                 &new_writes,
                                 &ContractClassMapping::default(),
-                                &HashMap::default(),
+                                &VisitedPcsSet::default(),
                             );
                             scheduler.finish_execution_during_commit(tx_index);
                         }
@@ -66,7 +68,7 @@ fn scheduler_flow_test(
                         versioned_state.pin_version(tx_index).apply_writes(
                             &writes,
                             &ContractClassMapping::default(),
-                            &HashMap::default(),
+                            &VisitedPcsSet::default(),
                         );
                         scheduler.finish_execution(tx_index);
                         Task::AskForTask

crates/blockifier/src/concurrency/versioned_state.rs

@@ -7,9 +7,10 @@ use starknet_types_core::felt::Felt;
 use crate::concurrency::versioned_storage::VersionedStorage;
 use crate::concurrency::TxIndex;
 use crate::execution::contract_class::ContractClass;
-use crate::state::cached_state::{ContractClassMapping, StateMaps, VisitedPcs};
+use crate::state::cached_state::{ContractClassMapping, StateMaps};
 use crate::state::errors::StateError;
 use crate::state::state_api::{StateReader, StateResult, UpdatableState};
+use crate::state::visited_pcs::VisitedPcsSet;
 
 #[cfg(test)]
 #[path = "versioned_state_test.rs"]
@@ -201,7 +202,7 @@ impl<U: UpdatableState> VersionedState<U> {
     pub fn commit_chunk_and_recover_block_state(
         mut self,
         n_committed_txs: usize,
-        visited_pcs: VisitedPcs,
+        visited_pcs: VisitedPcsSet,
     ) -> U {
         if n_committed_txs == 0 {
             return self.into_initial_state();
@@ -276,7 +277,7 @@ impl<S: StateReader> UpdatableState for VersionedStateProxy<S> {
         &mut self,
         writes: &StateMaps,
         class_hash_to_class: &ContractClassMapping,
-        _visited_pcs: &VisitedPcs,
+        _visited_pcs: &VisitedPcsSet,
     ) {
         self.state().apply_writes(self.tx_index, writes, class_hash_to_class)
     }

crates/blockifier/src/concurrency/versioned_state_test.rs

@@ -24,6 +24,7 @@ use crate::state::cached_state::{
 };
 use crate::state::errors::StateError;
 use crate::state::state_api::{State, StateReader, UpdatableState};
+use crate::state::visited_pcs::VisitedPcsSet;
 use crate::test_utils::contracts::FeatureContract;
 use crate::test_utils::deploy_account::deploy_account_tx;
 use crate::test_utils::dict_state_reader::DictStateReader;
@@ -419,7 +420,7 @@ fn test_apply_writes(
     safe_versioned_state.pin_version(0).apply_writes(
         &transactional_states[0].cache.borrow().writes,
         &transactional_states[0].class_hash_to_class.borrow().clone(),
-        &HashMap::default(),
+        &VisitedPcsSet::default(),
     );
     assert!(transactional_states[1].get_class_hash_at(contract_address).unwrap() == class_hash_0);
     assert!(
@@ -451,7 +452,7 @@ fn test_apply_writes_reexecute_scenario(
     safe_versioned_state.pin_version(0).apply_writes(
         &transactional_states[0].cache.borrow().writes,
         &transactional_states[0].class_hash_to_class.borrow().clone(),
-        &HashMap::default(),
+        &VisitedPcsSet::default(),
     );
     // Although transaction 0 wrote to the shared state, version 1 needs to be re-executed to see
     // the new value (its read value has already been cached).
@@ -496,7 +497,7 @@ fn test_delete_writes(
         safe_versioned_state.pin_version(i).apply_writes(
             &tx_state.cache.borrow().writes,
             &tx_state.class_hash_to_class.borrow(),
-            &HashMap::default(),
+            &VisitedPcsSet::default(),
         );
     }
 
@@ -558,7 +559,7 @@ fn test_delete_writes_completeness(
     versioned_state_proxy.apply_writes(
         &state_maps_writes,
         &class_hash_to_class_writes,
-        &HashMap::default(),
+        &VisitedPcsSet::default(),
     );
     assert_eq!(
         safe_versioned_state.0.lock().unwrap().get_writes_of_index(tx_index),
@@ -635,7 +636,7 @@ fn test_versioned_proxy_state_flow(
     }
     let modified_block_state = safe_versioned_state
         .into_inner_state()
-        .commit_chunk_and_recover_block_state(4, HashMap::new());
+        .commit_chunk_and_recover_block_state(4, VisitedPcsSet::new());
 
     assert!(modified_block_state.get_class_hash_at(contract_address).unwrap() == class_hash_3);
     assert!(

crates/blockifier/src/concurrency/worker_logic.rs

@@ -14,9 +14,10 @@ use crate::concurrency::versioned_state::ThreadSafeVersionedState;
 use crate::concurrency::TxIndex;
 use crate::context::BlockContext;
 use crate::state::cached_state::{
-    ContractClassMapping, StateChanges, StateMaps, TransactionalState, VisitedPcs,
+    ContractClassMapping, StateChanges, StateMaps, TransactionalState,
 };
 use crate::state::state_api::{StateReader, UpdatableState};
+use crate::state::visited_pcs::VisitedPcsSet;
 use crate::transaction::objects::{TransactionExecutionInfo, TransactionExecutionResult};
 use crate::transaction::transaction_execution::Transaction;
 use crate::transaction::transactions::{ExecutableTransaction, ExecutionFlags};
@@ -32,7 +33,7 @@ pub struct ExecutionTaskOutput {
     pub reads: StateMaps,
     pub writes: StateMaps,
     pub contract_classes: ContractClassMapping,
-    pub visited_pcs: VisitedPcs,
+    pub visited_pcs: VisitedPcsSet,
     pub result: TransactionExecutionResult<TransactionExecutionInfo>,
 }
 
@@ -135,7 +136,7 @@ impl<'a, S: StateReader> WorkerExecutor<'a, S> {
             self.state.pin_version(tx_index).apply_writes(
                 &transactional_state.cache.borrow().writes,
                 &transactional_state.class_hash_to_class.borrow(),
-                &HashMap::default(),
+                &VisitedPcsSet::default(),
             );
         }
 
@@ -145,7 +146,7 @@ impl<'a, S: StateReader> WorkerExecutor<'a, S> {
         // In case of a failed transaction, we don't record its writes and visited pcs.
         let (writes, contract_classes, visited_pcs) = match execution_result {
             Ok(_) => (tx_reads_writes.writes, class_hash_to_class, transactional_state.visited_pcs),
-            Err(_) => (StateMaps::default(), HashMap::default(), HashMap::default()),
+            Err(_) => (StateMaps::default(), HashMap::default(), VisitedPcsSet::default()),
         };
         let mut execution_output = lock_mutex_in_array(&self.execution_outputs, tx_index);
         *execution_output = Some(ExecutionTaskOutput {
@@ -262,7 +263,7 @@ impl<'a, U: UpdatableState> WorkerExecutor<'a, U> {
     pub fn commit_chunk_and_recover_block_state(
         self,
         n_committed_txs: usize,
-        visited_pcs: VisitedPcs,
+        visited_pcs: VisitedPcsSet,
     ) -> U {
         self.state
             .into_inner_state()

crates/blockifier/src/concurrency/worker_logic_test.rs

@@ -22,6 +22,7 @@ use crate::context::{BlockContext, TransactionContext};
 use crate::fee::fee_utils::get_sequencer_balance_keys;
 use crate::state::cached_state::StateMaps;
 use crate::state::state_api::StateReader;
+use crate::state::visited_pcs::VisitedPcsSet;
 use crate::test_utils::contracts::FeatureContract;
 use crate::test_utils::declare::declare_tx;
 use crate::test_utils::initial_test_state::test_state;
@@ -383,7 +384,7 @@ fn test_worker_execute(max_resource_bounds: ResourceBoundsMapping) {
 
     assert_eq!(execution_output.writes, writes);
     assert_eq!(execution_output.reads, reads);
-    assert_ne!(execution_output.visited_pcs, HashMap::default());
+    assert_ne!(execution_output.visited_pcs, VisitedPcsSet::default());
 
     // Failed execution.
     let tx_index = 1;
@@ -402,7 +403,7 @@ fn test_worker_execute(max_resource_bounds: ResourceBoundsMapping) {
     };
     assert_eq!(execution_output.reads, reads);
     assert_eq!(execution_output.writes, StateMaps::default());
-    assert_eq!(execution_output.visited_pcs, HashMap::default());
+    assert_eq!(execution_output.visited_pcs, VisitedPcsSet::default());
 
     // Reverted execution.
     let tx_index = 2;
@@ -416,7 +417,7 @@ fn test_worker_execute(max_resource_bounds: ResourceBoundsMapping) {
     let execution_output = execution_output.as_ref().unwrap();
     assert!(execution_output.result.as_ref().unwrap().is_reverted());
     assert_ne!(execution_output.writes, StateMaps::default());
-    assert_ne!(execution_output.visited_pcs, HashMap::default());
+    assert_ne!(execution_output.visited_pcs, VisitedPcsSet::default());
 
     // Validate status change.
     for tx_index in 0..3 {

crates/blockifier/src/execution/entry_point_execution.rs

@@ -1,5 +1,3 @@
-use std::collections::HashSet;
-
 use cairo_vm::types::builtin_name::BuiltinName;
 use cairo_vm::types::layout_name::LayoutName;
 use cairo_vm::types::relocatable::{MaybeRelocatable, Relocatable};
@@ -114,7 +112,7 @@ fn register_visited_pcs(
     program_segment_size: usize,
     bytecode_length: usize,
 ) -> EntryPointExecutionResult<()> {
-    let mut class_visited_pcs = HashSet::new();
+    let mut class_visited_pcs = Vec::new();
     // Relocate the trace, putting the program segment at address 1 and the execution segment right
     // after it.
     // TODO(lior): Avoid unnecessary relocation once the VM has a non-relocated `get_trace()`
@@ -131,7 +129,7 @@ fn register_visited_pcs(
         // Jumping to a PC that is not inside the bytecode is possible. For example, to obtain
         // the builtin costs. Filter out these values.
         if real_pc < bytecode_length {
-            class_visited_pcs.insert(real_pc);
+            class_visited_pcs.push(real_pc);
         }
     }
     state.add_visited_pcs(class_hash, &class_visited_pcs);

crates/blockifier/src/state.rs

@@ -4,3 +4,4 @@ pub mod error_format_test;
 pub mod errors;
 pub mod global_cache;
 pub mod state_api;
+pub mod visited_pcs;

crates/blockifier/src/state/cached_state.rs

@@ -7,6 +7,7 @@ use starknet_api::core::{ClassHash, CompiledClassHash, ContractAddress, Nonce};
 use starknet_api::state::StorageKey;
 use starknet_types_core::felt::Felt;
 
+use super::visited_pcs::{VisitedPcsSet, VisitedPcsTrait};
 use crate::abi::abi_utils::get_fee_token_var_address;
 use crate::context::TransactionContext;
 use crate::execution::contract_class::ContractClass;
@@ -21,8 +22,6 @@ mod test;
 
 pub type ContractClassMapping = HashMap<ClassHash, ContractClass>;
 
-pub type VisitedPcs = HashMap<ClassHash, HashSet<usize>>;
-
 /// Caches read and write requests.
 ///
 /// Writer functionality is builtin, whereas Reader functionality is injected through
@@ -35,7 +34,7 @@ pub struct CachedState<S: StateReader> {
     pub(crate) cache: RefCell<StateCache>,
     pub(crate) class_hash_to_class: RefCell<ContractClassMapping>,
     /// A map from class hash to the set of PC values that were visited in the class.
-    pub visited_pcs: VisitedPcs,
+    pub visited_pcs: VisitedPcsSet,
 }
 
 impl<S: StateReader> CachedState<S> {
@@ -44,7 +43,7 @@ impl<S: StateReader> CachedState<S> {
             state,
             cache: RefCell::new(StateCache::default()),
             class_hash_to_class: RefCell::new(HashMap::default()),
-            visited_pcs: VisitedPcs::default(),
+            visited_pcs: VisitedPcsSet::default(),
         }
     }
 
@@ -75,9 +74,10 @@ impl<S: StateReader> CachedState<S> {
         self.class_hash_to_class.get_mut().extend(local_contract_cache_updates);
     }
 
-    pub fn update_visited_pcs_cache(&mut self, visited_pcs: &VisitedPcs) {
+    pub fn update_visited_pcs_cache(&mut self, visited_pcs: &VisitedPcsSet) {
         for (class_hash, class_visited_pcs) in visited_pcs {
-            self.add_visited_pcs(*class_hash, class_visited_pcs);
+            let vec_visited_pcs = Vec::from_iter(class_visited_pcs.clone());
+            self.add_visited_pcs(*class_hash, &vec_visited_pcs);
         }
     }
 
@@ -114,7 +114,7 @@ impl<S: StateReader> UpdatableState for CachedState<S> {
         &mut self,
         writes: &StateMaps,
         class_hash_to_class: &ContractClassMapping,
-        visited_pcs: &VisitedPcs,
+        visited_pcs: &VisitedPcsSet,
     ) {
         // TODO(Noa,15/5/24): Reconsider the clone.
         self.update_cache(writes, class_hash_to_class.clone());
@@ -277,8 +277,8 @@ impl<S: StateReader> State for CachedState<S> {
         Ok(())
     }
 
-    fn add_visited_pcs(&mut self, class_hash: ClassHash, pcs: &HashSet<usize>) {
-        self.visited_pcs.entry(class_hash).or_default().extend(pcs);
+    fn add_visited_pcs(&mut self, class_hash: ClassHash, pcs: &Vec<usize>) {
+        self.visited_pcs.insert(&class_hash, pcs);
     }
 }
 

crates/blockifier/src/state/state_api.rs

@@ -1,10 +1,9 @@
-use std::collections::HashSet;
-
 use starknet_api::core::{ClassHash, CompiledClassHash, ContractAddress, Nonce};
 use starknet_api::state::StorageKey;
 use starknet_types_core::felt::Felt;
 
-use super::cached_state::{ContractClassMapping, StateMaps, VisitedPcs};
+use super::cached_state::{ContractClassMapping, StateMaps};
+use super::visited_pcs::VisitedPcsSet;
 use crate::abi::abi_utils::get_fee_token_var_address;
 use crate::abi::sierra_types::next_storage_key;
 use crate::execution::contract_class::ContractClass;
@@ -107,7 +106,7 @@ pub trait State: StateReader {
     /// Marks the given set of PC values as visited for the given class hash.
     // TODO(lior): Once we have a BlockResources object, move this logic there. Make sure reverted
     //   entry points do not affect the final set of PCs.
-    fn add_visited_pcs(&mut self, class_hash: ClassHash, pcs: &HashSet<usize>);
+    fn add_visited_pcs(&mut self, class_hash: ClassHash, pcs: &Vec<usize>);
 }
 
 /// A class defining the API for updating a state with transactions writes.
@@ -116,6 +115,6 @@ pub trait UpdatableState: StateReader {
         &mut self,
         writes: &StateMaps,
         class_hash_to_class: &ContractClassMapping,
-        visited_pcs: &VisitedPcs,
+        visited_pcs: &VisitedPcsSet,
     );
 }