From 4121a00d4821a69fefa1dcb873688fb717649be4 Mon Sep 17 00:00:00 2001
From: Joonatan Saarhelo <joon.saar@gmail.com>
Date: Thu, 2 May 2024 17:51:31 +0200
Subject: [PATCH] VM2 no longer rolls back the root frame

---
 Cargo.lock                                  |  2 +-
 compiler_tester/src/vm/eravm/vm2_adapter.rs | 84 ++++++++++++---------
 2 files changed, 49 insertions(+), 37 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 15e8f451..58421f86 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2582,7 +2582,7 @@ checksum = "2266fcb904c50fb17fda4c9a751a1715629ecf8b21f4c9d78b4890fb71525d71"
 [[package]]
 name = "vm2"
 version = "0.1.0"
-source = "git+https://github.com/matter-labs/vm2#f745d472bdae30cffeb47a0422f5cde8e0af2a97"
+source = "git+https://github.com/matter-labs/vm2#33d7d7a6775a20a7527d4822814237d6e61c761c"
 dependencies = [
  "enum_dispatch",
  "primitive-types",
diff --git a/compiler_tester/src/vm/eravm/vm2_adapter.rs b/compiler_tester/src/vm/eravm/vm2_adapter.rs
index 97c6ef22..d3305e52 100644
--- a/compiler_tester/src/vm/eravm/vm2_adapter.rs
+++ b/compiler_tester/src/vm/eravm/vm2_adapter.rs
@@ -98,12 +98,49 @@ pub fn run_vm(
     vm.state.registers[4] = abi_params.r4_value.unwrap_or_default();
     vm.state.registers[5] = abi_params.r5_value.unwrap_or_default();
 
+    let mut storage_changes = HashMap::new();
+    let mut deployed_contracts = HashMap::new();
+
     let output = match vm.run() {
-        ExecutionEnd::ProgramFinished(return_value) => Output {
-            return_data: chunk_return_data(&return_value),
-            exception: false,
-            events: merge_events(vm.world.events()),
-        },
+        ExecutionEnd::ProgramFinished(return_value) => {
+            // Only successful transactions can have side effects
+            // The VM doesn't undo side effects done in the initial frame
+            // because that would mess with the bootloader.
+
+            storage_changes = vm
+                .world
+                .get_storage_changes()
+                .iter()
+                .map(|(&(address, key), value)| {
+                    (StorageKey { address, key }, H256::from_uint(value))
+                })
+                .collect::<HashMap<_, _>>();
+            deployed_contracts = vm
+            .world
+            .get_storage_changes()
+            .iter()
+            .filter_map(|((address, key), value)| {
+                if *address == *zkevm_assembly::zkevm_opcode_defs::system_params::DEPLOYER_SYSTEM_CONTRACT_ADDRESS {
+                    let mut buffer = [0u8; 32];
+                    key.to_big_endian(&mut buffer);
+                    let deployed_address = web3::ethabi::Address::from_slice(&buffer[12..]);
+                    if let Some(code) = known_contracts.get(&value) {
+                        Some((deployed_address, code.clone()))
+                    } else {
+                        None
+                    }
+                } else {
+                    None
+                }
+            })
+            .collect::<HashMap<_, _>>();
+
+            Output {
+                return_data: chunk_return_data(&return_value),
+                exception: false,
+                events: merge_events(vm.world.events()),
+            }
+        }
         ExecutionEnd::Reverted(return_value) => Output {
             return_data: chunk_return_data(&return_value),
             exception: true,
@@ -117,32 +154,6 @@ pub fn run_vm(
         ExecutionEnd::SuspendedOnHook { .. } => unreachable!(),
     };
 
-    let storage_changes = vm
-        .world
-        .get_storage_changes()
-        .iter()
-        .map(|(&(address, key), value)| (StorageKey { address, key }, H256::from_uint(value)))
-        .collect::<HashMap<_, _>>();
-    let deployed_contracts = vm
-        .world
-        .get_storage_changes()
-        .iter()
-        .filter_map(|((address, key), value)| {
-            if *address == *zkevm_assembly::zkevm_opcode_defs::system_params::DEPLOYER_SYSTEM_CONTRACT_ADDRESS {
-                let mut buffer = [0u8; 32];
-                key.to_big_endian(&mut buffer);
-                let deployed_address = web3::ethabi::Address::from_slice(&buffer[12..]);
-                if let Some(code) = known_contracts.get(&value) {
-                    Some((deployed_address, code.clone()))
-                } else {
-                    None
-                }
-            } else {
-                None
-            }
-        })
-        .collect::<HashMap<_, _>>();
-
     Ok((
         ExecutionResult {
             output,
@@ -162,13 +173,14 @@ struct TestWorld {
 
 impl World for TestWorld {
     fn decommit(&mut self, hash: U256) -> Program {
-        let bytecode = self
+        let Some(bytecode) = self
             .contracts
             .get(&hash)
-            .unwrap()
-            .clone()
-            .compile_to_bytecode()
-            .unwrap();
+            .map(|assembly| assembly.clone().compile_to_bytecode().unwrap())
+        else {
+            // This case only happens when a contract fails to deploy but the assumes it did deploy
+            return Program::new(vec![vm2::Instruction::from_invalid()], vec![]);
+        };
         let instructions = bytecode
             .iter()
             .flat_map(|x| {