feat(vm): Allow caching signature verification (#3505)

## What ❔ Allows caching signature verification (more precisely, `ecrecover` output for L2 transactions using the default AA) in the fast VM. ## Why ❔ Signature verification takes ~50% of the execution time for "simple" transactions (e.g., base token and ERC-20 token transfers), so caching it may improve performance. ## Checklist - [x] PR title corresponds to the body of PR (we generate changelog entries from PRs). - [x] Tests for the changes have been added / updated. - [x] Documentation comments have been added / updated. - [x] Code has been formatted via `zkstack dev fmt` and `zkstack dev lint`.
matter-labs · Feb 3, 2025 · 7bb5ed3 · 7bb5ed3
1 parent 8ebb9f9
commit 7bb5ed3
Show file tree

Hide file tree

Showing 21 changed files with 435 additions and 209 deletions.
diff --git a/core/Cargo.lock b/core/Cargo.lock
diff --git a/core/Cargo.toml b/core/Cargo.toml
@@ -244,7 +244,7 @@ fflonk = "=0.30.13"
 bellman = { package = "zksync_bellman", version = "=0.30.13" }
 
 # New VM; pinned to a specific commit because of instability
-zksync_vm2 = { git = "https://github.com/matter-labs/vm2.git", rev = "457d8a7eea9093af9440662e33e598c13ba41633" }
+zksync_vm2 = { git = "https://github.com/matter-labs/vm2.git", rev = "3841f5a430288a63c8207853eca11560bf7a5712" }
 
 # Consensus dependencies.
 zksync_concurrency = "=0.8.0"

diff --git a/core/lib/multivm/src/versions/vm_fast/mod.rs b/core/lib/multivm/src/versions/vm_fast/mod.rs
@@ -15,3 +15,4 @@ mod tracers;
 mod utils;
 mod version;
 mod vm;
+mod world;
diff --git a/core/lib/multivm/src/versions/vm_fast/tests/precompiles.rs b/core/lib/multivm/src/versions/vm_fast/tests/precompiles.rs
@@ -1,5 +1,12 @@
+use circuit_sequencer_api::geometry_config::ProtocolGeometry;
+use zksync_types::Execute;
+
 use crate::{
-    versions::testonly::precompiles::{test_ecrecover, test_keccak, test_sha256},
+    interface::{InspectExecutionMode, TxExecutionMode, VmInterface, VmInterfaceExt},
+    versions::testonly::{
+        precompiles::{test_ecrecover, test_keccak, test_sha256},
+        VmTesterBuilder,
+    },
     vm_fast::Vm,
 };
 
@@ -17,3 +24,39 @@ fn sha256() {
 fn ecrecover() {
     test_ecrecover::<Vm<_>>();
 }
+
+#[test]
+fn caching_ecrecover_result() {
+    let mut vm = VmTesterBuilder::new()
+        .with_empty_in_memory_storage()
+        .with_rich_accounts(1)
+        .with_execution_mode(TxExecutionMode::VerifyExecute)
+        .build::<Vm<_>>();
+    vm.vm.skip_signature_verification();
+
+    let account = &mut vm.rich_accounts[0];
+    let tx = account.get_l2_tx_for_execute(
+        Execute {
+            contract_address: Some(account.address),
+            calldata: vec![],
+            value: 0.into(),
+            factory_deps: vec![],
+        },
+        None,
+    );
+    vm.vm.push_transaction(tx);
+
+    assert!(vm.vm.world.precompiles.expected_ecrecover_call.is_some());
+    assert_eq!(vm.vm.world.precompiles.expected_calls.get(), 0);
+
+    let exec_result = vm.vm.execute(InspectExecutionMode::OneTx);
+    assert!(!exec_result.result.is_failed(), "{exec_result:#?}");
+    assert_eq!(vm.vm.world.precompiles.expected_calls.get(), 1);
+
+    // Cycle stats should still be produced for the cached call
+    let ecrecover_count = exec_result.statistics.circuit_statistic.ecrecover
+        * ProtocolGeometry::V1_5_0
+            .config()
+            .cycles_per_ecrecover_circuit as f32;
+    assert!((ecrecover_count - 1.0).abs() < 1e-4, "{ecrecover_count}");
+}
diff --git a/core/lib/multivm/src/versions/vm_fast/vm.rs b/core/lib/multivm/src/versions/vm_fast/vm.rs
@@ -3,30 +3,21 @@ use std::{collections::HashMap, fmt, mem, rc::Rc};
 use zk_evm_1_5_0::{
     aux_structures::LogQuery, zkevm_opcode_defs::system_params::INITIAL_FRAME_FORMAL_EH_LOCATION,
 };
-use zksync_contracts::SystemContractCode;
 use zksync_types::{
-    bytecode::BytecodeHash,
-    h256_to_u256,
-    l1::is_l1_tx_type,
-    l2_to_l1_log::UserL2ToL1Log,
-    u256_to_h256,
-    utils::key_for_eth_balance,
-    writes::{
-        compression::compress_with_best_strategy, StateDiffRecord, BYTES_PER_DERIVED_KEY,
-        BYTES_PER_ENUMERATION_INDEX,
-    },
-    AccountTreeId, StorageKey, StorageLog, StorageLogKind, StorageLogWithPreviousValue,
-    Transaction, BOOTLOADER_ADDRESS, H160, H256, KNOWN_CODES_STORAGE_ADDRESS, L1_MESSENGER_ADDRESS,
-    L2_BASE_TOKEN_ADDRESS, U256,
+    bytecode::BytecodeHash, h256_to_u256, l1::is_l1_tx_type, l2_to_l1_log::UserL2ToL1Log,
+    u256_to_h256, writes::StateDiffRecord, AccountTreeId, StorageKey, StorageLog, StorageLogKind,
+    StorageLogWithPreviousValue, Transaction, BOOTLOADER_ADDRESS, H160, H256,
+    KNOWN_CODES_STORAGE_ADDRESS, L1_MESSENGER_ADDRESS, U256,
 };
 use zksync_vm2::{
     interface::{CallframeInterface, HeapId, StateInterface, Tracer},
-    ExecutionEnd, FatPointer, Program, Settings, StorageSlot, VirtualMachine,
+    ExecutionEnd, FatPointer, Settings, VirtualMachine,
 };
 
 use super::{
     bytecode::compress_bytecodes,
-    tracers::{DynamicBytecodes, ValidationTracer, WithBuiltinTracers},
+    tracers::{ValidationTracer, WithBuiltinTracers},
+    world::World,
 };
 use crate::{
     glue::GlueInto,
@@ -99,6 +90,7 @@ pub struct Vm<S, Tr = (), Val = ()> {
     pub(super) system_env: SystemEnv,
     snapshot: Option<VmSnapshot>,
     vm_version: FastVmVersion,
+    skip_signature_verification: bool,
     #[cfg(test)]
     enforced_state_diffs: Option<Vec<StateDiffRecord>>,
 }
@@ -171,13 +163,18 @@ impl<S: ReadStorage, Tr: Tracer, Val: ValidationTracer> Vm<S, Tr, Val> {
             batch_env,
             snapshot: None,
             vm_version,
+            skip_signature_verification: false,
             #[cfg(test)]
             enforced_state_diffs: None,
         };
         this.write_to_bootloader_heap(bootloader_memory);
         this
     }
 
+    pub fn skip_signature_verification(&mut self) {
+        self.skip_signature_verification = true;
+    }
+
     fn get_hook_params(&self) -> [U256; 3] {
         (get_vm_hook_params_start_position(self.vm_version.into())
             ..get_vm_hook_params_start_position(self.vm_version.into()) + VM_HOOK_PARAMS_COUNT)
@@ -281,17 +278,28 @@ impl<S: ReadStorage, Tr: Tracer, Val: ValidationTracer> Vm<S, Tr, Val> {
         };
 
         let trusted_ergs_limit = tx.trusted_ergs_limit();
-
-        let memory = self.bootloader_state.push_tx(
+        let tx_origin = tx.from;
+        let (memory, ecrecover_call) = self.bootloader_state.push_tx(
             tx,
             overhead,
             refund,
             compressed_bytecodes,
             trusted_ergs_limit,
             self.system_env.chain_id,
         );
-
         self.write_to_bootloader_heap(memory);
+
+        // The expected `ecrecover` call params *must* be reset on each transaction.
+        // We only set call params for transactions using the default AA. Other AAs may expect another
+        // address returned from `ecrecover`, or may not invoke `ecrecover` at all during validation
+        // (both of which would make caching a security hazard).
+        let needs_default_aa_check = self.skip_signature_verification && ecrecover_call.is_some();
+        self.world.precompiles.expected_ecrecover_call =
+            if needs_default_aa_check && self.world.has_default_aa(&tx_origin) {
+                ecrecover_call
+            } else {
+                None
+            };
     }
 
     #[cfg(test)]
@@ -889,163 +897,10 @@ impl<S: fmt::Debug, Tr: fmt::Debug, Val: fmt::Debug> fmt::Debug for Vm<S, Tr, Va
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.debug_struct("Vm")
             .field("bootloader_state", &self.bootloader_state)
-            .field("storage", &self.world.storage)
-            .field("program_cache", &self.world.program_cache)
+            .field("world", &self.world)
             .field("batch_env", &self.batch_env)
             .field("system_env", &self.system_env)
             .field("snapshot", &self.snapshot.as_ref().map(|_| ()))
             .finish()
     }
 }
-
-#[derive(Debug)]
-pub(crate) struct World<S, T> {
-    pub(crate) storage: S,
-    dynamic_bytecodes: DynamicBytecodes,
-    program_cache: HashMap<U256, Program<T, Self>>,
-    pub(crate) bytecode_cache: HashMap<U256, Vec<u8>>,
-}
-
-impl<S: ReadStorage, T: Tracer> World<S, T> {
-    fn new(storage: S, program_cache: HashMap<U256, Program<T, Self>>) -> Self {
-        Self {
-            storage,
-            dynamic_bytecodes: DynamicBytecodes::default(),
-            program_cache,
-            bytecode_cache: HashMap::default(),
-        }
-    }
-
-    fn convert_system_contract_code(
-        code: &SystemContractCode,
-        is_bootloader: bool,
-    ) -> (U256, Program<T, Self>) {
-        (
-            h256_to_u256(code.hash),
-            Program::new(&code.code, is_bootloader),
-        )
-    }
-
-    fn decommit_dynamic_bytecodes(
-        &self,
-        candidate_hashes: impl Iterator<Item = H256>,
-    ) -> HashMap<H256, Vec<u8>> {
-        let bytecodes = candidate_hashes.filter_map(|hash| {
-            let bytecode = self
-                .dynamic_bytecodes
-                .map(h256_to_u256(hash), <[u8]>::to_vec)?;
-            Some((hash, bytecode))
-        });
-        bytecodes.collect()
-    }
-}
-
-impl<S: ReadStorage, T: Tracer> zksync_vm2::StorageInterface for World<S, T> {
-    fn read_storage(&mut self, contract: H160, key: U256) -> StorageSlot {
-        let key = &StorageKey::new(AccountTreeId::new(contract), u256_to_h256(key));
-        let value = U256::from_big_endian(self.storage.read_value(key).as_bytes());
-        // `is_write_initial` value can be true even if the slot has previously been written to / has non-zero value!
-        // This can happen during oneshot execution (i.e., executing a single transaction) since it emulates
-        // execution starting in the middle of a batch in the general case. Hence, a slot that was first written to in the batch
-        // must still be considered an initial write by the refund logic.
-        let is_write_initial = self.storage.is_write_initial(key);
-        StorageSlot {
-            value,
-            is_write_initial,
-        }
-    }
-
-    fn read_storage_value(&mut self, contract: H160, key: U256) -> U256 {
-        let key = &StorageKey::new(AccountTreeId::new(contract), u256_to_h256(key));
-        U256::from_big_endian(self.storage.read_value(key).as_bytes())
-    }
-
-    fn cost_of_writing_storage(&mut self, slot: StorageSlot, new_value: U256) -> u32 {
-        if slot.value == new_value {
-            return 0;
-        }
-
-        // Since we need to publish the state diffs onchain, for each of the updated storage slot
-        // we basically need to publish the following pair: `(<storage_key, compressed_new_value>)`.
-        // For key we use the following optimization:
-        //   - The first time we publish it, we use 32 bytes.
-        //         Then, we remember a 8-byte id for this slot and assign it to it. We call this initial write.
-        //   - The second time we publish it, we will use the 4/5 byte representation of this 8-byte instead of the 32
-        //     bytes of the entire key.
-        // For value compression, we use a metadata byte which holds the length of the value and the operation from the
-        // previous state to the new state, and the compressed value. The maximum for this is 33 bytes.
-        // Total bytes for initial writes then becomes 65 bytes and repeated writes becomes 38 bytes.
-        let compressed_value_size = compress_with_best_strategy(slot.value, new_value).len() as u32;
-
-        if slot.is_write_initial {
-            (BYTES_PER_DERIVED_KEY as u32) + compressed_value_size
-        } else {
-            (BYTES_PER_ENUMERATION_INDEX as u32) + compressed_value_size
-        }
-    }
-
-    fn is_free_storage_slot(&self, contract: &H160, key: &U256) -> bool {
-        contract == &zksync_system_constants::SYSTEM_CONTEXT_ADDRESS
-            || contract == &L2_BASE_TOKEN_ADDRESS
-                && u256_to_h256(*key) == key_for_eth_balance(&BOOTLOADER_ADDRESS)
-    }
-}
-
-/// It may look like that an append-only cache for EVM bytecodes / `Program`s can lead to the following scenario:
-///
-/// 1. A transaction deploys an EVM bytecode with hash `H`, then reverts.
-/// 2. A following transaction in the same VM run queries a bytecode with hash `H` and gets it.
-///
-/// This would be incorrect behavior because bytecode deployments must be reverted along with transactions.
-///
-/// In reality, this cannot happen because both `decommit()` and `decommit_code()` calls perform storage-based checks
-/// before a decommit:
-///
-/// - `decommit_code()` is called from the `CodeOracle` system contract, which checks that the decommitted bytecode is known.
-/// - `decommit()` is called during far calls, which obtains address -> bytecode hash mapping beforehand.
-///
-/// Thus, if storage is reverted correctly, additional EVM bytecodes occupy the cache, but are unreachable.
-impl<S: ReadStorage, T: Tracer> zksync_vm2::World<T> for World<S, T> {
-    fn decommit(&mut self, hash: U256) -> Program<T, Self> {
-        self.program_cache
-            .entry(hash)
-            .or_insert_with(|| {
-                let cached = self
-                    .bytecode_cache
-                    .get(&hash)
-                    .map(|code| Program::new(code, false))
-                    .or_else(|| {
-                        self.dynamic_bytecodes
-                            .map(hash, |code| Program::new(code, false))
-                    });
-
-                if let Some(cached) = cached {
-                    cached
-                } else {
-                    let code = self
-                        .storage
-                        .load_factory_dep(u256_to_h256(hash))
-                        .unwrap_or_else(|| {
-                            panic!("VM tried to decommit nonexistent bytecode: {hash:?}");
-                        });
-                    let program = Program::new(&code, false);
-                    self.bytecode_cache.insert(hash, code);
-                    program
-                }
-            })
-            .clone()
-    }
-
-    fn decommit_code(&mut self, hash: U256) -> Vec<u8> {
-        self.decommit(hash)
-            .code_page()
-            .as_ref()
-            .iter()
-            .flat_map(|u| {
-                let mut buffer = [0u8; 32];
-                u.to_big_endian(&mut buffer);
-                buffer
-            })
-            .collect()
-    }
-}