feat(concurrency): delete writes

crates/blockifier/src/concurrency/test_utils.rs

@@ -1,3 +1,8 @@
+use rstest::fixture;
+use starknet_api::core::{ContractAddress, PatriciaKey};
+use starknet_api::hash::StarkHash;
+use starknet_api::{contract_address, patricia_key};
+
 use crate::concurrency::versioned_state_proxy::{ThreadSafeVersionedState, VersionedState};
 use crate::context::BlockContext;
 use crate::execution::call_info::CallInfo;
@@ -6,6 +11,16 @@ use crate::state::state_api::StateReader;
 use crate::test_utils::dict_state_reader::DictStateReader;
 use crate::transaction::account_transaction::AccountTransaction;
 use crate::transaction::transactions::ExecutableTransaction;
+// Fixtures.
+
+const TEST_CONTRACT_ADDRESS: &str = "0x18031991";
+
+#[fixture]
+pub fn contract_address() -> ContractAddress {
+    contract_address!(TEST_CONTRACT_ADDRESS)
+}
+
+// Macros.
 
 #[macro_export]
 macro_rules! default_scheduler {
@@ -35,13 +50,17 @@ macro_rules! default_scheduler {
     };
 }
 
+// Concurrency constructors.
+
 // TODO(meshi, 01/06/2024): Consider making this a macro.
 pub fn safe_versioned_state_for_testing(
     block_state: DictStateReader,
 ) -> ThreadSafeVersionedState<DictStateReader> {
     ThreadSafeVersionedState::new(VersionedState::new(block_state))
 }
 
+// Utils.
+
 // Note: this function does not mutate the state.
 pub fn create_fee_transfer_call_info<S: StateReader>(
     state: &mut CachedState<S>,

crates/blockifier/src/concurrency/versioned_state_proxy.rs

@@ -139,6 +139,14 @@ impl<S: StateReader> VersionedState<S> {
             self.compiled_contract_classes.write(tx_index, key, value.clone());
         }
     }
+
+    fn delete_writes(&mut self, tx_index: TxIndex) {
+        self.storage.delete_writes(tx_index);
+        self.nonces.delete_writes(tx_index);
+        self.class_hashes.delete_writes(tx_index);
+        self.compiled_class_hashes.delete_writes(tx_index);
+        self.compiled_contract_classes.delete_writes(tx_index);
+    }
 }
 
 pub struct ThreadSafeVersionedState<S: StateReader>(Arc<Mutex<VersionedState<S>>>);
@@ -177,6 +185,10 @@ impl<S: StateReader> VersionedStateProxy<S> {
     pub fn apply_writes(&self, writes: &StateMaps, class_hash_to_class: &ContractClassMapping) {
         self.state().apply_writes(self.tx_index, writes, class_hash_to_class)
     }
+
+    pub fn delete_writes(&self, tx_index: TxIndex) {
+        self.state().delete_writes(tx_index);
+    }
 }
 
 impl<S: StateReader> StateReader for VersionedStateProxy<S> {

crates/blockifier/src/concurrency/versioned_state_proxy_test.rs

@@ -9,7 +9,7 @@ use starknet_api::transaction::{Calldata, ContractAddressSalt, Fee, TransactionV
 use starknet_api::{calldata, class_hash, contract_address, patricia_key, stark_felt};
 
 use crate::abi::abi_utils::{get_fee_token_var_address, get_storage_var_address};
-use crate::concurrency::test_utils::safe_versioned_state_for_testing;
+use crate::concurrency::test_utils::{contract_address, safe_versioned_state_for_testing};
 use crate::concurrency::versioned_state_proxy::{
     ThreadSafeVersionedState, VersionedState, VersionedStateProxy,
 };
@@ -27,14 +27,8 @@ use crate::transaction::test_utils::l1_resource_bounds;
 use crate::transaction::transactions::ExecutableTransaction;
 use crate::{compiled_class_hash, deploy_account_tx_args, nonce, storage_key};
 
-const TEST_CONTRACT_ADDRESS: &str = "0x1";
 const TEST_CLASS_HASH: u8 = 27_u8;
 
-#[fixture]
-pub fn contract_address() -> ContractAddress {
-    contract_address!(TEST_CONTRACT_ADDRESS)
-}
-
 #[fixture]
 pub fn class_hash() -> ClassHash {
     class_hash!(TEST_CLASS_HASH)
@@ -371,3 +365,9 @@ fn test_apply_writes_reexecute_scenario(
     // The class hash should be updated.
     assert!(transactional_states[1].get_class_hash_at(contract_address).unwrap() == class_hash_0);
 }
+
+#[rstest]
+fn test_delete_writes() {
+    // TODO(barak, 01/07/2024): Complete the test.
+    assert_eq!(0, 0)
+}

crates/blockifier/src/concurrency/versioned_storage.rs

@@ -50,6 +50,12 @@ where
         cell.insert(tx_index, value);
     }
 
+    pub fn delete_writes(&mut self, tx_index: TxIndex) {
+        for (_, inner_map) in self.writes.iter_mut() {
+            inner_map.remove(&tx_index);
+        }
+    }
+
     /// This method inserts the provided key-value pair into the cached initial values map.
     /// It is typically used when reading a value that is not found in the versioned storage. In
     /// such a scenario, the value is retrieved from the initial storage and written to the

crates/blockifier/src/concurrency/versioned_storage_test.rs

@@ -1,7 +1,14 @@
+use std::collections::{BTreeMap, HashMap};
+
 use pretty_assertions::assert_eq;
+use rstest::rstest;
+use starknet_api::core::ContractAddress;
 
+use crate::concurrency::test_utils::contract_address;
 use crate::concurrency::versioned_storage::VersionedStorage;
+use crate::nonce;
 
+// TODO(barak, 01/07/2024): Split into test_read() and test_write().
 #[test]
 fn test_versioned_storage() {
     let mut storage = VersionedStorage::default();
@@ -37,3 +44,19 @@ fn test_versioned_storage() {
     // Test the write.
     assert_eq!(storage.read(50, 100).unwrap(), 194);
 }
+
+#[rstest]
+fn test_delete_writes(contract_address: ContractAddress) {
+    // TODO(barak, 01/07/2025): Create a macro versioned_storage!.
+    let mut storage = VersionedStorage {
+        cached_initial_values: HashMap::default(),
+        writes: HashMap::from([(
+            contract_address,
+            BTreeMap::from([(1, nonce!(18_u8)), (2, nonce!(19_u8))]),
+        )]),
+    };
+    storage.delete_writes(1);
+    let contract_address_writes = storage.writes.get(&contract_address).unwrap();
+    assert!(!contract_address_writes.contains_key(&1));
+    assert!(contract_address_writes.contains_key(&2));
+}