Refactor rollbacks

.gitmodules

@@ -1,3 +1,4 @@
 [submodule "era-compiler-tester"]
 	path = era-compiler-tester
 	url = https://github.com/lambdaclass/era-compiler-tester.git
+	branch = refactor-rollbacks

src/call_frame.rs

@@ -2,16 +2,15 @@ use std::num::Saturating;
 use u256::U256;
 use zkevm_opcode_defs::ethereum_types::Address;
 
-use crate::{state::Stack, store::SnapShot, utils::is_kernel};
+use crate::{state::Stack, utils::is_kernel, world::WorldSnapshot};
 
 #[derive(Debug, Clone)]
 pub struct CallFrame {
     pub pc: u64,
-    pub transient_storage_snapshot: SnapShot,
     pub gas_left: Saturating<u32>,
     pub exception_handler: u64,
     pub sp: u32,
-    pub storage_snapshot: SnapShot,
+    pub snapshot: WorldSnapshot,
 }
 
 #[derive(Debug, Clone)]
@@ -64,17 +63,11 @@ impl Context {
         calldata_heap_id: u32,
         exception_handler: u64,
         context_u128: u128,
-        transient_storage_snapshot: SnapShot,
-        storage_snapshot: SnapShot,
+        snapshot: WorldSnapshot,
         is_static: bool,
     ) -> Self {
         Self {
-            frame: CallFrame::new_far_call_frame(
-                gas_stipend,
-                exception_handler,
-                storage_snapshot,
-                transient_storage_snapshot,
-            ),
+            frame: CallFrame::new_far_call_frame(gas_stipend, exception_handler, snapshot),
             near_call_frames: vec![],
             contract_address,
             caller,
@@ -98,16 +91,14 @@ impl CallFrame {
     pub fn new_far_call_frame(
         gas_stipend: u32,
         exception_handler: u64,
-        storage_snapshot: SnapShot,
-        transient_storage_snapshot: SnapShot,
+        snapshot: WorldSnapshot,
     ) -> Self {
         Self {
             pc: 0,
             gas_left: Saturating(gas_stipend),
-            transient_storage_snapshot,
             exception_handler,
             sp: 0,
-            storage_snapshot,
+            snapshot,
         }
     }
 
@@ -116,17 +107,15 @@ impl CallFrame {
         sp: u32,
         pc: u64,
         gas_stipend: u32,
-        transient_storage_snapshot: SnapShot,
         exception_handler: u64,
-        storage_snapshot: SnapShot,
+        snapshot: WorldSnapshot,
     ) -> Self {
         Self {
             pc,
             gas_left: Saturating(gas_stipend),
-            transient_storage_snapshot,
             exception_handler,
             sp,
-            storage_snapshot,
+            snapshot,
         }
     }
 }

src/lib.rs

@@ -16,4 +16,6 @@ pub mod vm;
 pub use opcode::Opcode;
 pub use state::VMState;
 pub use vm::EraVM;
+mod rollbacks;
+pub mod world;
 use zkevm_opcode_defs::Opcode as Variant;

src/op_handlers/event.rs

@@ -3,11 +3,12 @@ use zkevm_opcode_defs::ADDRESS_EVENT_WRITER;
 
 use crate::{
     eravm_error::EraVmError,
-    state::{Event, VMState},
+    state::VMState,
+    world::{Event, World},
     Opcode,
 };
 
-pub fn event(vm: &mut VMState, opcode: &Opcode) -> Result<(), EraVmError> {
+pub fn event(vm: &mut VMState, opcode: &Opcode, world: &mut World) -> Result<(), EraVmError> {
     if vm.current_context()?.contract_address == H160::from_low_u64_be(ADDRESS_EVENT_WRITER as u64)
     {
         let key = vm.get_register(opcode.src0_index).value;
@@ -20,7 +21,7 @@ pub fn event(vm: &mut VMState, opcode: &Opcode) -> Result<(), EraVmError> {
             tx_number: vm.tx_number as u16,
         };
 
-        vm.events.push(event);
+        world.record_event(event);
     }
     Ok(())
 }

src/op_handlers/far_call.rs

@@ -8,10 +8,12 @@ use zkevm_opcode_defs::{
 use crate::{
     address_operands::address_operands_read,
     eravm_error::{EraVmError, HeapError},
+    rollbacks::Rollbackable,
     state::VMState,
-    store::{ContractStorage, StateStorage, StorageError, StorageKey},
+    store::{StorageError, StorageKey},
     utils::{address_into_u256, is_kernel},
     value::{FatPointer, TaggedValue},
+    world::World,
     Opcode,
 };
 #[allow(dead_code)]
@@ -122,7 +124,7 @@ fn address_from_u256(register_value: &U256) -> H160 {
 }
 
 fn decommit_code_hash(
-    state_storage: &mut StateStorage,
+    world: &mut World,
     address: Address,
     default_aa_code_hash: [u8; 32],
     evm_interpreter_code_hash: [u8; 32],
@@ -133,10 +135,10 @@ fn decommit_code_hash(
         Address::from_low_u64_be(DEPLOYER_SYSTEM_CONTRACT_ADDRESS_LOW as u64);
     let storage_key = StorageKey::new(deployer_system_contract_address, address_into_u256(address));
 
-    let code_info = match state_storage.storage_read(storage_key)? {
+    let code_info = match world.storage_read(storage_key)? {
         Some(code_info) => code_info,
         None => {
-            state_storage.storage_write(storage_key, U256::zero())?;
+            world.storage_write(storage_key, U256::zero());
             U256::zero()
         }
     };
@@ -203,23 +205,20 @@ pub fn far_call(
     vm: &mut VMState,
     opcode: &Opcode,
     far_call: &FarCallOpcode,
-    state_storage: &mut StateStorage,
-    contract_storage: &mut dyn ContractStorage,
-    transient_storage: &StateStorage,
+    world: &mut World,
 ) -> Result<(), EraVmError> {
     let (src0, src1) = address_operands_read(vm, opcode)?;
     let contract_address = address_from_u256(&src1.value);
 
     let exception_handler = opcode.imm0 as u64;
-    let storage_snapshot = state_storage.create_snapshot();
-    let transient_storage_snapshot = transient_storage.create_snapshot();
+    let snapshot = world.snapshot();
 
     let mut abi = get_far_call_arguments(src0.value);
     abi.is_constructor_call = abi.is_constructor_call && vm.current_context()?.is_kernel();
     abi.is_system_call = abi.is_system_call && is_kernel(&contract_address);
 
     let (code_key, is_evm) = decommit_code_hash(
-        state_storage,
+        world,
         contract_address,
         vm.default_aa_code_hash,
         vm.evm_interpreter_code_hash,
@@ -240,7 +239,7 @@ pub fn far_call(
         .checked_add(stipend)
         .expect("stipend must not cause overflow");
 
-    let program_code = contract_storage
+    let program_code = world
         .decommit(code_key)?
         .ok_or(StorageError::KeyNotPresent)?;
     let new_heap = vm.heaps.allocate();
@@ -260,8 +259,7 @@ pub fn far_call(
                 forward_memory.page,
                 exception_handler,
                 vm.register_context_u128,
-                transient_storage_snapshot,
-                storage_snapshot,
+                snapshot,
                 is_new_frame_static && !is_evm,
             )?;
         }
@@ -286,8 +284,7 @@ pub fn far_call(
                 forward_memory.page,
                 exception_handler,
                 vm.register_context_u128,
-                transient_storage_snapshot,
-                storage_snapshot,
+                snapshot,
                 is_new_frame_static && !is_evm,
             )?;
         }
@@ -306,8 +303,7 @@ pub fn far_call(
                 forward_memory.page,
                 exception_handler,
                 this_context.context_u128,
-                transient_storage_snapshot,
-                storage_snapshot,
+                snapshot,
                 is_new_frame_static && !is_evm,
             )?;
         }

src/op_handlers/log.rs

@@ -3,78 +3,75 @@ use u256::U256;
 use crate::{
     eravm_error::EraVmError,
     state::VMState,
-    store::{L2ToL1Log, StateStorage, StorageKey},
+    store::StorageKey,
     value::TaggedValue,
+    world::{L2ToL1Log, World},
     Opcode,
 };
 
 pub fn storage_write(
     vm: &mut VMState,
     opcode: &Opcode,
-    state_storage: &mut StateStorage,
+    world: &mut World,
 ) -> Result<(), EraVmError> {
     let key_for_contract_storage = vm.get_register(opcode.src0_index).value;
     let address = vm.current_context()?.contract_address;
     let key = StorageKey::new(address, key_for_contract_storage);
     let value = vm.get_register(opcode.src1_index).value;
-    state_storage.storage_write(key, value)?;
+    world.storage_write(key, value);
     Ok(())
 }
 
-pub fn storage_read(
-    vm: &mut VMState,
-    opcode: &Opcode,
-    state_storage: &StateStorage,
-) -> Result<(), EraVmError> {
+pub fn storage_read(vm: &mut VMState, opcode: &Opcode, world: &World) -> Result<(), EraVmError> {
     let key_for_contract_storage = vm.get_register(opcode.src0_index).value;
     let address = vm.current_context()?.contract_address;
     let key = StorageKey::new(address, key_for_contract_storage);
-    let value = state_storage.storage_read(key)?.unwrap_or(U256::zero());
+    let value = world.storage_read(key)?.unwrap_or(U256::zero());
     vm.set_register(opcode.dst0_index, TaggedValue::new_raw_integer(value));
     Ok(())
 }
 
 pub fn transient_storage_write(
     vm: &mut VMState,
     opcode: &Opcode,
-    transient_storage: &mut StateStorage,
+    world: &mut World,
 ) -> Result<(), EraVmError> {
     let key_for_contract_storage = vm.get_register(opcode.src0_index).value;
     let address = vm.current_context()?.contract_address;
     let key = StorageKey::new(address, key_for_contract_storage);
     let value = vm.get_register(opcode.src1_index).value;
-    transient_storage.storage_write(key, value)?;
+    world.transient_storage_write(key, value);
     Ok(())
 }
 
 pub fn transient_storage_read(
     vm: &mut VMState,
     opcode: &Opcode,
-    transient_storage: &StateStorage,
+    world: &World,
 ) -> Result<(), EraVmError> {
     let key_for_contract_storage = vm.get_register(opcode.src0_index).value;
     let address = vm.current_context()?.contract_address;
     let key = StorageKey::new(address, key_for_contract_storage);
-    let value = transient_storage.storage_read(key)?.unwrap_or(U256::zero());
+    let value = world.transient_storage_read(key)?.unwrap_or(U256::zero());
     vm.set_register(opcode.dst0_index, TaggedValue::new_raw_integer(value));
     Ok(())
 }
 
 pub fn add_l2_to_l1_message(
     vm_state: &mut VMState,
     opcode: &Opcode,
-    state_storage: &mut StateStorage,
+    world: &mut World,
 ) -> Result<(), EraVmError> {
     let key = vm_state.get_register(opcode.src0_index).value;
     let value = vm_state.get_register(opcode.src1_index).value;
     let is_service = opcode.imm0 == 1;
-    state_storage.record_l2_to_l1_log(L2ToL1Log {
+    world.record_l2_to_l1_log(L2ToL1Log {
         key,
         value,
         is_service,
         address: vm_state.current_context()?.contract_address,
         shard_id: 0,
         tx_number: vm_state.tx_number as u16,
-    })?;
+    });
     Ok(())
 }

src/op_handlers/near_call.rs

@@ -2,15 +2,11 @@ use u256::U256;
 
 use crate::call_frame::CallFrame;
 use crate::eravm_error::EraVmError;
-use crate::store::StateStorage;
+use crate::rollbacks::Rollbackable;
+use crate::world::World;
 use crate::{opcode::Opcode, state::VMState};
 
-pub fn near_call(
-    vm: &mut VMState,
-    opcode: &Opcode,
-    state_storage: &StateStorage,
-    transient_storage: &StateStorage,
-) -> Result<(), EraVmError> {
+pub fn near_call(vm: &mut VMState, opcode: &Opcode, world: &World) -> Result<(), EraVmError> {
     let abi_reg = vm.get_register(opcode.src0_index);
     let call_pc = opcode.imm0 as u64;
     let exception_handler = opcode.imm1;
@@ -28,16 +24,14 @@ pub fn near_call(
     let current_frame = vm.current_frame_mut()?;
 
     let new_sp = current_frame.sp;
-    let transient_storage_snapshot = transient_storage.create_snapshot();
 
     // Create new frame
     let new_frame = CallFrame::new_near_call_frame(
         new_sp,
         call_pc,
         callee_ergs,
-        transient_storage_snapshot,
         exception_handler as u64,
-        state_storage.create_snapshot(),
+        world.snapshot(),
     );
 
     vm.push_near_call_frame(new_frame)

src/op_handlers/opcode_decommit.rs

@@ -2,15 +2,15 @@ use crate::{
     address_operands::{address_operands_read, address_operands_store},
     eravm_error::{EraVmError, HeapError},
     state::VMState,
-    store::ContractStorage,
     value::{FatPointer, TaggedValue},
+    world::World,
     Opcode,
 };
 
 pub fn opcode_decommit(
     vm: &mut VMState,
     opcode: &Opcode,
-    storage: &mut dyn ContractStorage,
+    world: &mut World,
 ) -> Result<(), EraVmError> {
     let (src0, src1) = address_operands_read(vm, opcode)?;
 
@@ -20,7 +20,7 @@ pub fn opcode_decommit(
 
     vm.decrease_gas(extra_cost)?;
 
-    let code = storage
+    let code = world
         .decommit(code_hash)?
         .ok_or(EraVmError::DecommitFailed)?;