fix(vm): Fix used bytecodes divergence

Cargo.toml

@@ -218,7 +218,7 @@ zk_evm_1_4_1 = { package = "zk_evm", version = "0.141.0" }
 zk_evm_1_5_0 = { package = "zk_evm", version = "=0.150.4" }
 
 # New VM; pinned to a specific commit because of instability
-vm2 = { git = "https://github.com/matter-labs/vm2.git", rev = "9a38900d7af9b1d72b47ce3be980e77c1239a61d" }
+vm2 = { git = "https://github.com/matter-labs/vm2.git", rev = "2276b7b5af520fca0477bdafe43781b51896d235" }
 
 # Consensus dependencies.
 zksync_concurrency = "=0.1.0-rc.11"

core/lib/multivm/Cargo.toml

@@ -40,6 +40,7 @@ tracing.workspace = true
 vise.workspace = true
 
 [dev-dependencies]
+assert_matches.workspace = true
 tokio = { workspace = true, features = ["time"] }
 zksync_test_account.workspace = true
 ethabi.workspace = true

core/lib/multivm/src/versions/vm_fast/tests/get_used_contracts.rs

@@ -1,17 +1,23 @@
-use std::collections::HashSet;
+use std::{collections::HashSet, iter};
 
+use assert_matches::assert_matches;
+use ethabi::Token;
 use itertools::Itertools;
+use zk_evm_1_3_1::zkevm_opcode_defs::decoding::{EncodingModeProduction, VmEncodingMode};
 use zksync_system_constants::CONTRACT_DEPLOYER_ADDRESS;
 use zksync_test_account::Account;
-use zksync_types::{Execute, U256};
+use zksync_types::{Address, Execute, U256};
 use zksync_utils::{bytecode::hash_bytecode, h256_to_u256};
 
 use crate::{
-    interface::{storage::ReadStorage, TxExecutionMode, VmExecutionMode, VmInterface},
+    interface::{
+        storage::ReadStorage, ExecutionResult, TxExecutionMode, VmExecutionMode,
+        VmExecutionResultAndLogs, VmInterface,
+    },
     vm_fast::{
         tests::{
-            tester::{TxType, VmTesterBuilder},
-            utils::{read_test_contract, BASE_SYSTEM_CONTRACTS},
+            tester::{TxType, VmTester, VmTesterBuilder},
+            utils::{read_proxy_counter_contract, read_test_contract, BASE_SYSTEM_CONTRACTS},
         },
         vm::Vm,
     },
@@ -88,8 +94,90 @@ fn known_bytecodes_without_aa_code<S: ReadStorage>(vm: &Vm<S>) -> HashSet<U256>
         .keys()
         .cloned()
         .collect::<HashSet<_>>();
-
     known_bytecodes_without_aa_code.remove(&h256_to_u256(BASE_SYSTEM_CONTRACTS.default_aa.hash));
-
     known_bytecodes_without_aa_code
 }
+
+/// Counter test contract bytecode inflated by appending lots of `NOP` opcodes at the end. This leads to non-trivial
+/// decommitment cost (>10,000 gas).
+fn inflated_counter_bytecode() -> Vec<u8> {
+    let mut counter_bytecode = read_test_contract();
+    counter_bytecode.extend(
+        iter::repeat(EncodingModeProduction::nop_encoding().to_be_bytes())
+            .take(10_000)
+            .flatten(),
+    );
+    counter_bytecode
+}
+
+fn execute_proxy_counter(gas: u32) -> (VmTester, U256, VmExecutionResultAndLogs) {
+    let counter_bytecode = inflated_counter_bytecode();
+    let counter_bytecode_hash = h256_to_u256(hash_bytecode(&counter_bytecode));
+    let counter_address = Address::repeat_byte(0x23);
+
+    let mut vm = VmTesterBuilder::new()
+        .with_empty_in_memory_storage()
+        .with_custom_contracts(vec![(counter_bytecode, counter_address, false)])
+        .with_execution_mode(TxExecutionMode::VerifyExecute)
+        .with_random_rich_accounts(1)
+        .build();
+
+    let (proxy_counter_bytecode, proxy_counter_abi) = read_proxy_counter_contract();
+    let account = &mut vm.rich_accounts[0];
+    let deploy_tx = account.get_deploy_tx(
+        &proxy_counter_bytecode,
+        Some(&[Token::Address(counter_address)]),
+        TxType::L2,
+    );
+    let (compression_result, exec_result) = vm
+        .vm
+        .execute_transaction_with_bytecode_compression(deploy_tx.tx, true);
+    compression_result.unwrap();
+    assert!(!exec_result.result.is_failed(), "{exec_result:#?}");
+
+    let decommitted_hashes = vm.vm.decommitted_hashes().collect::<HashSet<_>>();
+    assert!(
+        !decommitted_hashes.contains(&counter_bytecode_hash),
+        "{decommitted_hashes:?}"
+    );
+
+    let increment = proxy_counter_abi.function("increment").unwrap();
+    let increment_tx = account.get_l2_tx_for_execute(
+        Execute {
+            contract_address: deploy_tx.address,
+            calldata: increment
+                .encode_input(&[Token::Uint(1.into()), Token::Uint(gas.into())])
+                .unwrap(),
+            value: 0.into(),
+            factory_deps: vec![],
+        },
+        None,
+    );
+    let (compression_result, exec_result) = vm
+        .vm
+        .execute_transaction_with_bytecode_compression(increment_tx, true);
+    compression_result.unwrap();
+    (vm, counter_bytecode_hash, exec_result)
+}
+
+#[test]
+fn get_used_contracts_with_far_call() {
+    let (vm, counter_bytecode_hash, exec_result) = execute_proxy_counter(100_000);
+    assert!(!exec_result.result.is_failed(), "{exec_result:#?}");
+    let decommitted_hashes = vm.vm.decommitted_hashes().collect::<HashSet<_>>();
+    assert!(
+        decommitted_hashes.contains(&counter_bytecode_hash),
+        "{decommitted_hashes:?}"
+    );
+}
+
+#[test]
+fn get_used_contracts_with_out_of_gas_far_call() {
+    let (vm, counter_bytecode_hash, exec_result) = execute_proxy_counter(10_000);
+    assert_matches!(exec_result.result, ExecutionResult::Revert { .. });
+    let decommitted_hashes = vm.vm.decommitted_hashes().collect::<HashSet<_>>();
+    assert!(
+        decommitted_hashes.contains(&counter_bytecode_hash),
+        "{decommitted_hashes:?}"
+    );
+}

core/lib/multivm/src/versions/vm_fast/tests/utils.rs

@@ -127,3 +127,8 @@ pub(crate) fn read_expensive_contract() -> (Vec<u8>, Contract) {
         "etc/contracts-test-data/artifacts-zk/contracts/expensive/expensive.sol/Expensive.json";
     (read_bytecode(PATH), load_contract(PATH))
 }
+
+pub(crate) fn read_proxy_counter_contract() -> (Vec<u8>, Contract) {
+    const PATH: &str = "etc/contracts-test-data/artifacts-zk/contracts/counter/proxy_counter.sol/ProxyCounter.json";
+    (read_bytecode(PATH), load_contract(PATH))
+}

core/lib/multivm/src/versions/vm_latest/tests/get_used_contracts.rs

@@ -1,25 +1,32 @@
 use std::{
     collections::{HashMap, HashSet},
+    iter,
     str::FromStr,
 };
 
+use assert_matches::assert_matches;
+use ethabi::Token;
 use itertools::Itertools;
+use zk_evm_1_3_1::zkevm_opcode_defs::decoding::{EncodingModeProduction, VmEncodingMode};
 use zk_evm_1_5_0::{
     abstractions::DecommittmentProcessor,
     aux_structures::{DecommittmentQuery, MemoryPage, Timestamp},
     zkevm_opcode_defs::{VersionedHashHeader, VersionedHashNormalizedPreimage},
 };
 use zksync_system_constants::CONTRACT_DEPLOYER_ADDRESS;
 use zksync_test_account::Account;
-use zksync_types::{Execute, U256};
+use zksync_types::{Address, Execute, U256};
 use zksync_utils::{bytecode::hash_bytecode, h256_to_u256};
+use zksync_vm_interface::VmExecutionResultAndLogs;
 
 use crate::{
-    interface::{storage::WriteStorage, TxExecutionMode, VmExecutionMode, VmInterface},
+    interface::{
+        storage::WriteStorage, ExecutionResult, TxExecutionMode, VmExecutionMode, VmInterface,
+    },
     vm_latest::{
         tests::{
-            tester::{TxType, VmTesterBuilder},
-            utils::{read_test_contract, BASE_SYSTEM_CONTRACTS},
+            tester::{TxType, VmTester, VmTesterBuilder},
+            utils::{read_proxy_counter_contract, read_test_contract, BASE_SYSTEM_CONTRACTS},
         },
         HistoryDisabled, Vm,
     },
@@ -148,10 +155,92 @@ fn known_bytecodes_without_aa_code<S: WriteStorage, H: HistoryMode>(
         .known_bytecodes
         .inner()
         .clone();
-
     known_bytecodes_without_aa_code
         .remove(&h256_to_u256(BASE_SYSTEM_CONTRACTS.default_aa.hash))
         .unwrap();
-
     known_bytecodes_without_aa_code
 }
+
+/// Counter test contract bytecode inflated by appending lots of `NOP` opcodes at the end. This leads to non-trivial
+/// decommitment cost (>10,000 gas).
+fn inflated_counter_bytecode() -> Vec<u8> {
+    let mut counter_bytecode = read_test_contract();
+    counter_bytecode.extend(
+        iter::repeat(EncodingModeProduction::nop_encoding().to_be_bytes())
+            .take(10_000)
+            .flatten(),
+    );
+    counter_bytecode
+}
+
+fn execute_proxy_counter(gas: u32) -> (VmTester<HistoryDisabled>, U256, VmExecutionResultAndLogs) {
+    let counter_bytecode = inflated_counter_bytecode();
+    let counter_bytecode_hash = h256_to_u256(hash_bytecode(&counter_bytecode));
+    let counter_address = Address::repeat_byte(0x23);
+
+    let mut vm = VmTesterBuilder::new(HistoryDisabled)
+        .with_empty_in_memory_storage()
+        .with_custom_contracts(vec![(counter_bytecode, counter_address, false)])
+        .with_execution_mode(TxExecutionMode::VerifyExecute)
+        .with_random_rich_accounts(1)
+        .build();
+
+    let (proxy_counter_bytecode, proxy_counter_abi) = read_proxy_counter_contract();
+    let account = &mut vm.rich_accounts[0];
+    let deploy_tx = account.get_deploy_tx(
+        &proxy_counter_bytecode,
+        Some(&[Token::Address(counter_address)]),
+        TxType::L2,
+    );
+    let (compression_result, exec_result) = vm
+        .vm
+        .execute_transaction_with_bytecode_compression(deploy_tx.tx, true);
+    compression_result.unwrap();
+    assert!(!exec_result.result.is_failed(), "{exec_result:#?}");
+
+    let decommitted_hashes = vm.vm.get_used_contracts();
+    assert!(
+        !decommitted_hashes.contains(&counter_bytecode_hash),
+        "{decommitted_hashes:?}"
+    );
+
+    let increment = proxy_counter_abi.function("increment").unwrap();
+    let increment_tx = account.get_l2_tx_for_execute(
+        Execute {
+            contract_address: deploy_tx.address,
+            calldata: increment
+                .encode_input(&[Token::Uint(1.into()), Token::Uint(gas.into())])
+                .unwrap(),
+            value: 0.into(),
+            factory_deps: vec![],
+        },
+        None,
+    );
+    let (compression_result, exec_result) = vm
+        .vm
+        .execute_transaction_with_bytecode_compression(increment_tx, true);
+    compression_result.unwrap();
+    (vm, counter_bytecode_hash, exec_result)
+}
+
+#[test]
+fn get_used_contracts_with_far_call() {
+    let (vm, counter_bytecode_hash, exec_result) = execute_proxy_counter(100_000);
+    assert!(!exec_result.result.is_failed(), "{exec_result:#?}");
+    let decommitted_hashes = vm.vm.get_used_contracts();
+    assert!(
+        decommitted_hashes.contains(&counter_bytecode_hash),
+        "{decommitted_hashes:?}"
+    );
+}
+
+#[test]
+fn get_used_contracts_with_out_of_gas_far_call() {
+    let (vm, counter_bytecode_hash, exec_result) = execute_proxy_counter(10_000);
+    assert_matches!(exec_result.result, ExecutionResult::Revert { .. });
+    let decommitted_hashes = vm.vm.get_used_contracts();
+    assert!(
+        decommitted_hashes.contains(&counter_bytecode_hash),
+        "{decommitted_hashes:?}"
+    );
+}

core/lib/multivm/src/versions/vm_latest/tests/utils.rs

@@ -137,3 +137,8 @@ pub(crate) fn read_expensive_contract() -> (Vec<u8>, Contract) {
         "etc/contracts-test-data/artifacts-zk/contracts/expensive/expensive.sol/Expensive.json";
     (read_bytecode(PATH), load_contract(PATH))
 }
+
+pub(crate) fn read_proxy_counter_contract() -> (Vec<u8>, Contract) {
+    const PATH: &str = "etc/contracts-test-data/artifacts-zk/contracts/counter/proxy_counter.sol/ProxyCounter.json";
+    (read_bytecode(PATH), load_contract(PATH))
+}

etc/contracts-test-data/contracts/counter/counter.sol

@@ -5,7 +5,7 @@ pragma solidity ^0.8.0;
 contract Counter {
     uint256 value;
 
-    function increment(uint256 x) public {
+    function increment(uint256 x) external {
         value += x;
     }
 

etc/contracts-test-data/contracts/counter/proxy_counter.sol

@@ -0,0 +1,22 @@
+// SPDX-License-Identifier: MIT OR Apache-2.0
+
+pragma solidity ^0.8.0;
+
+interface ICounter {
+    function increment(uint256 x) external;
+}
+
+contract ProxyCounter {
+    ICounter counter;
+
+    constructor(ICounter _counter) {
+        counter = _counter;
+    }
+
+    function increment(uint256 x, uint gasToPass) public {
+        while (gasleft() > gasToPass) {
+            // Burn gas so that there's about `gasToPass` left before the external call.
+        }
+        counter.increment(x);
+    }
+}