refactor InMemoryCustomRefundStorage to leverage InMemoryStorage as a…

… field

src/tests/complex_tests/mod.rs

@@ -176,7 +176,7 @@ pub(crate) fn generate_base_layer(
     use crate::external_calls::run;
     use crate::toolset::GeometryConfig;
 
-    let mut storage_impl = InMemoryCustomRefundStorage::new(None);
+    let mut storage_impl = InMemoryStorage::new();
     let mut tree = ZKSyncTestingTree::empty();
 
     test_artifact.entry_point_address =

src/tests/mod.rs

@@ -5,6 +5,9 @@ pub mod complex_tests;
 pub mod run_manually;
 #[cfg(test)]
 pub mod simple_tests;
+#[cfg(test)]
+pub(crate) mod storage;
+#[cfg(test)]
 pub(crate) mod utils;
 
 use crate::blake2::Blake2s256;
@@ -22,7 +25,6 @@ use circuit_definitions::circuit_definitions::recursion_layer::ZkSyncRecursiveLa
 use circuit_definitions::ZkSyncDefaultRoundFunction;
 use std::alloc::Global;
 use std::collections::HashMap;
-use utils::storage::InMemoryCustomRefundStorage;
 
 const ACCOUNT_CODE_STORAGE_ADDRESS: Address = H160([
     0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
@@ -35,7 +37,7 @@ const KNOWN_CODE_HASHES_ADDRESS: Address = H160([
 ]);
 
 pub(crate) fn save_predeployed_contracts(
-    storage: &mut InMemoryCustomRefundStorage,
+    storage: &mut InMemoryStorage,
     tree: &mut impl BinarySparseStorageTree<256, 32, 32, 8, 32, Blake2s256, ZkSyncStorageLeaf>,
     contracts: &HashMap<Address, Vec<[u8; 32]>>,
 ) {

src/tests/run_manually.rs

@@ -24,8 +24,7 @@ use crate::zkevm_circuits::base_structures::vm_state::GlobalContextWitness;
 use crate::zkevm_circuits::main_vm::main_vm_entry_point;
 use circuit_definitions::aux_definitions::witness_oracle::VmWitnessOracle;
 use circuit_definitions::zk_evm::vm_state::cycle;
-use utils::storage::InMemoryCustomRefundStorage;
-use utils::StorageRefund;
+use storage::{InMemoryCustomRefundStorage, StorageRefund};
 use zkevm_assembly::Assembly;
 
 #[test]
@@ -261,7 +260,7 @@ pub(crate) fn run_with_options(entry_point_bytecode: Vec<[u8; 32]>, options: Opt
     let mut known_contracts = HashMap::new();
     known_contracts.extend(options.other_contracts.iter().cloned());
 
-    save_predeployed_contracts(&mut storage_impl, &mut tree, &known_contracts);
+    save_predeployed_contracts(&mut storage_impl.storage, &mut tree, &known_contracts);
 
     let mut basic_block_circuits = vec![];
 

src/tests/storage.rs

@@ -0,0 +1,94 @@
+use std::{
+    collections::{HashMap, HashSet},
+    sync::{Arc, Mutex},
+};
+
+use circuit_definitions::{
+    ethereum_types::{Address, H160},
+    zk_evm::{
+        abstractions::{Storage, StorageAccessRefund},
+        aux_structures::{LogQuery, PubdataCost, Timestamp},
+        reference_impls::{event_sink::ApplicationData, memory::SimpleMemory},
+        testing::{storage::InMemoryStorage, NUM_SHARDS},
+        tracing::{
+            AfterDecodingData, AfterExecutionData, BeforeExecutionData, Tracer, VmLocalStateData,
+        },
+        vm_state::PrimitiveValue,
+    },
+};
+use zkevm_assembly::zkevm_opcode_defs::{
+    decoding::{AllowedPcOrImm, EncodingModeProduction, VmEncodingMode},
+    AddOpcode, DecodedOpcode, NopOpcode, Opcode, PtrOpcode, RetOpcode, MAX_PUBDATA_COST_PER_QUERY,
+    STORAGE_ACCESS_COLD_READ_COST, STORAGE_ACCESS_COLD_WRITE_COST, STORAGE_ACCESS_WARM_READ_COST,
+    STORAGE_ACCESS_WARM_WRITE_COST, STORAGE_AUX_BYTE, TRANSIENT_STORAGE_AUX_BYTE,
+};
+
+use crate::ethereum_types::U256;
+
+/// Enum holding the types of storage refunds
+#[derive(Debug, Copy, Clone)]
+pub(crate) enum StorageRefund {
+    Cold,
+    Warm,
+}
+
+#[derive(Debug, Clone)]
+pub struct InMemoryCustomRefundStorage {
+    pub storage: InMemoryStorage,
+    pub slot_refund: Option<Arc<Mutex<(StorageRefund, u32)>>>,
+}
+
+impl InMemoryCustomRefundStorage {
+    pub fn new(slot_refund: Option<Arc<Mutex<(StorageRefund, u32)>>>) -> Self {
+        Self {
+            storage: InMemoryStorage::new(),
+            slot_refund,
+        }
+    }
+}
+
+impl Storage for InMemoryCustomRefundStorage {
+    #[track_caller]
+    fn get_access_refund(
+        &mut self, // to avoid any hacks inside, like prefetch
+        _monotonic_cycle_counter: u32,
+        _partial_query: &LogQuery,
+    ) -> StorageAccessRefund {
+        match &self.slot_refund {
+            None => StorageAccessRefund::Cold,
+            Some(val) => {
+                let (refund_type, val) = *val.lock().unwrap();
+
+                match refund_type {
+                    StorageRefund::Cold => dbg!(StorageAccessRefund::Cold),
+                    StorageRefund::Warm => dbg!(StorageAccessRefund::Warm { ergs: val }),
+                }
+            }
+        }
+    }
+
+    #[track_caller]
+    fn execute_partial_query(
+        &mut self,
+        monotonic_cycle_counter: u32,
+        query: LogQuery,
+    ) -> (LogQuery, PubdataCost) {
+        self.storage
+            .execute_partial_query(monotonic_cycle_counter, query)
+    }
+
+    #[track_caller]
+    fn start_frame(&mut self, timestamp: Timestamp) {
+        self.storage.start_frame(timestamp)
+    }
+
+    #[track_caller]
+    fn finish_frame(&mut self, timestamp: Timestamp, panicked: bool) {
+        self.storage.finish_frame(timestamp, panicked)
+    }
+
+    #[track_caller]
+    fn start_new_tx(&mut self, timestamp: Timestamp) {
+        self.storage.start_new_tx(timestamp)
+    }
+}

src/tests/utils/mod.rs

@@ -1,10 +1,2 @@
 pub mod preprocess_asm;
-pub mod storage;
 pub mod testing_tracer;
-
-/// Enum holding the types of storage refunds
-#[derive(Debug, Copy, Clone)]
-pub(crate) enum StorageRefund {
-    Cold,
-    Warm,
-}

src/tests/utils/testing_tracer.rs

@@ -20,6 +20,7 @@ use zkevm_assembly::zkevm_opcode_defs::RetOpcode;
 use zkevm_assembly::zkevm_opcode_defs::REGISTERS_COUNT;
 
 use crate::ethereum_types::U256;
+use crate::tests::storage::StorageRefund;
 use crate::zk_evm::reference_impls::memory::SimpleMemory;
 use crate::zk_evm::tracing::*;
 
@@ -30,7 +31,6 @@ use crate::tests::utils::preprocess_asm::PRINT_REG_PREFIX;
 
 use super::preprocess_asm::STORAGE_REFUND_COLD_PREFIX;
 use super::preprocess_asm::STORAGE_REFUND_WARM_PREFIX;
-use super::StorageRefund;
 
 #[derive(Debug, Clone, PartialEq, Default)]
 enum TracerState {