feat: Allow customizing precompiles logic

Cargo.toml

@@ -23,11 +23,13 @@ categories = ["cryptography"]
 afl = "0.15"
 anyhow = "1"
 arbitrary = "1"
+assert_matches = "1.5.0"
 divan = "0.1"
 enum_dispatch = "0.3"
 pretty_assertions = "1.4.0"
 primitive-types = "0.12.1"
 proptest = "1.4"
+rand = "0.8.5"
 
 # "Internal" dependencies
 zkevm_opcode_defs = "0.150.0"

crates/vm2-interface/src/state_interface.rs

@@ -216,15 +216,15 @@ pub struct L2ToL1Log {
 }
 
 #[cfg(test)]
-pub mod tests {
+pub(crate) mod testonly {
     use primitive_types::{H160, U256};
 
     use super::{
         CallframeInterface, Event, Flags, GlobalStateInterface, HeapId, L2ToL1Log, StateInterface,
     };
 
     #[derive(Debug)]
-    pub struct DummyState;
+    pub(crate) struct DummyState;
 
     impl StateInterface for DummyState {
         fn read_register(&self, _: u8) -> (U256, bool) {

crates/vm2-interface/src/tracer_interface.rs

@@ -303,7 +303,7 @@ impl ShouldStop {
 }
 
 /// Cycle statistics emitted by the VM and supplied to [`Tracer::on_extra_prover_cycles()`].
-#[derive(Debug, Clone, Copy)]
+#[derive(Debug, Clone, Copy, PartialEq)]
 pub enum CycleStats {
     /// Call to the `keccak256` precompile with the specified number of hash cycles.
     Keccak256(u32),
@@ -349,7 +349,7 @@ impl<A: Tracer, B: Tracer> Tracer for (A, B) {
 #[cfg(test)]
 mod tests {
     use super::{CallingMode, OpcodeType};
-    use crate::{opcodes, tests::DummyState, GlobalStateInterface, Tracer};
+    use crate::{opcodes, testonly::DummyState, GlobalStateInterface, Tracer};
 
     struct FarCallCounter(usize);
 

crates/vm2/Cargo.toml

@@ -24,8 +24,10 @@ zk_evm = { workspace = true, optional = true }
 anyhow = { workspace = true, optional = true }
 
 [dev-dependencies]
+assert_matches.workspace = true
 divan.workspace = true
 proptest.workspace = true
+rand.workspace = true
 
 [lints]
 workspace = true

crates/vm2/src/instruction_handlers/precompiles.rs

@@ -1,141 +1,130 @@
-use primitive_types::{H160, U256};
-use zk_evm_abstractions::{
-    aux::Timestamp,
-    precompiles::{
-        ecrecover::ecrecover_function, keccak256::keccak256_rounds_function,
-        secp256r1_verify::secp256r1_verify_function, sha256::sha256_rounds_function,
-    },
-    queries::LogQuery,
-    vm::Memory,
-};
-use zkevm_opcode_defs::{
-    system_params::{
-        ECRECOVER_INNER_FUNCTION_PRECOMPILE_ADDRESS, KECCAK256_ROUND_FUNCTION_PRECOMPILE_ADDRESS,
-        SECP256R1_VERIFY_PRECOMPILE_ADDRESS, SHA256_ROUND_FUNCTION_PRECOMPILE_ADDRESS,
-    },
-    PrecompileAuxData, PrecompileCallABI,
-};
-use zksync_vm2_interface::{opcodes, CycleStats, HeapId, Tracer};
+use primitive_types::U256;
+use zksync_vm2_interface::{opcodes, HeapId, Tracer};
 
 use super::{common::boilerplate_ext, ret::spontaneous_panic};
 use crate::{
     addressing_modes::{Arguments, Destination, Register1, Register2, Source},
-    heap::Heaps,
     instruction::ExecutionStatus,
+    precompiles::{PrecompileMemoryReader, Precompiles},
     Instruction, VirtualMachine, World,
 };
 
-fn precompile_call<T: Tracer, W: World<T>>(
-    vm: &mut VirtualMachine<T, W>,
-    world: &mut W,
-    tracer: &mut T,
-) -> ExecutionStatus {
-    boilerplate_ext::<opcodes::PrecompileCall, _, _>(vm, world, tracer, |vm, args, _, tracer| {
-        // The user gets to decide how much gas to burn
-        // This is safe because system contracts are trusted
-        let aux_data = PrecompileAuxData::from_u256(Register2::get(args, &mut vm.state));
-        let Ok(()) = vm.state.use_gas(aux_data.extra_ergs_cost) else {
-            vm.state.current_frame.pc = spontaneous_panic();
-            return;
-        };
+#[derive(Debug)]
+struct PrecompileAuxData {
+    extra_ergs_cost: u32,
+    extra_pubdata_cost: u32,
+}
 
-        #[allow(clippy::cast_possible_wrap)]
-        {
-            vm.world_diff.pubdata.0 += aux_data.extra_pubdata_cost as i32;
-        }
+impl PrecompileAuxData {
+    #[allow(clippy::cast_possible_truncation)]
+    fn from_u256(raw_value: U256) -> Self {
+        let raw = raw_value.0;
+        let extra_ergs_cost = raw[0] as u32;
+        let extra_pubdata_cost = (raw[0] >> 32) as u32;
 
-        let mut abi = PrecompileCallABI::from_u256(Register1::get(args, &mut vm.state));
-        if abi.memory_page_to_read == 0 {
-            abi.memory_page_to_read = vm.state.current_frame.heap.as_u32();
-        }
-        if abi.memory_page_to_write == 0 {
-            abi.memory_page_to_write = vm.state.current_frame.heap.as_u32();
+        Self {
+            extra_ergs_cost,
+            extra_pubdata_cost,
         }
+    }
+}
 
-        let query = LogQuery {
-            timestamp: Timestamp(0),
-            key: abi.to_u256(),
-            // only two first fields are read by the precompile
-            tx_number_in_block: Default::default(),
-            aux_byte: Default::default(),
-            shard_id: Default::default(),
-            address: H160::default(),
-            read_value: U256::default(),
-            written_value: U256::default(),
-            rw_flag: Default::default(),
-            rollback: Default::default(),
-            is_service: Default::default(),
-        };
+#[derive(Debug)]
+struct PrecompileCallAbi {
+    input_memory_offset: u32,
+    input_memory_length: u32,
+    output_memory_offset: u32,
+    output_memory_length: u32,
+    memory_page_to_read: HeapId,
+    memory_page_to_write: HeapId,
+    precompile_interpreted_data: u64,
+}
 
-        let address_bytes = vm.state.current_frame.address.0;
-        let address_low = u16::from_le_bytes([address_bytes[19], address_bytes[18]]);
-        let heaps = &mut vm.state.heaps;
+impl PrecompileCallAbi {
+    #[allow(clippy::cast_possible_truncation)]
+    fn from_u256(raw_value: U256) -> Self {
+        let raw = raw_value.0;
+        let input_memory_offset = raw[0] as u32;
+        let input_memory_length = (raw[0] >> 32) as u32;
+        let output_memory_offset = raw[1] as u32;
+        let output_memory_length = (raw[1] >> 32) as u32;
+        let memory_page_to_read = HeapId::from_u32_unchecked(raw[2] as u32);
+        let memory_page_to_write = HeapId::from_u32_unchecked((raw[2] >> 32) as u32);
+        let precompile_interpreted_data = raw[3];
 
-        #[allow(clippy::cast_possible_truncation)]
-        // if we're having `> u32::MAX` cycles, we've got larger issues
-        match address_low {
-            KECCAK256_ROUND_FUNCTION_PRECOMPILE_ADDRESS => {
-                tracer.on_extra_prover_cycles(CycleStats::Keccak256(
-                    keccak256_rounds_function::<_, false>(0, query, heaps).0 as u32,
-                ));
-            }
-            SHA256_ROUND_FUNCTION_PRECOMPILE_ADDRESS => {
-                tracer.on_extra_prover_cycles(CycleStats::Sha256(
-                    sha256_rounds_function::<_, false>(0, query, heaps).0 as u32,
-                ));
-            }
-            ECRECOVER_INNER_FUNCTION_PRECOMPILE_ADDRESS => {
-                tracer.on_extra_prover_cycles(CycleStats::EcRecover(
-                    ecrecover_function::<_, false>(0, query, heaps).0 as u32,
-                ));
-            }
-            SECP256R1_VERIFY_PRECOMPILE_ADDRESS => {
-                tracer.on_extra_prover_cycles(CycleStats::Secp256r1Verify(
-                    secp256r1_verify_function::<_, false>(0, query, heaps).0 as u32,
-                ));
-            }
-            _ => {
-                // A precompile call may be used just to burn gas
-            }
+        Self {
+            input_memory_offset,
+            input_memory_length,
+            output_memory_offset,
+            output_memory_length,
+            memory_page_to_read,
+            memory_page_to_write,
+            precompile_interpreted_data,
         }
-
-        Register1::set(args, &mut vm.state, 1.into());
-    })
+    }
 }
 
-impl Memory for Heaps {
-    fn execute_partial_query(
-        &mut self,
-        _monotonic_cycle_counter: u32,
-        mut query: zk_evm_abstractions::queries::MemoryQuery,
-    ) -> zk_evm_abstractions::queries::MemoryQuery {
-        let page = HeapId::from_u32_unchecked(query.location.page.0);
+fn precompile_call<T: Tracer, W: World<T>>(
+    vm: &mut VirtualMachine<T, W>,
+    world: &mut W,
+    tracer: &mut T,
+) -> ExecutionStatus {
+    boilerplate_ext::<opcodes::PrecompileCall, _, _>(
+        vm,
+        world,
+        tracer,
+        |vm, args, world, tracer| {
+            // The user gets to decide how much gas to burn
+            // This is safe because system contracts are trusted
+            let aux_data = PrecompileAuxData::from_u256(Register2::get(args, &mut vm.state));
+            let Ok(()) = vm.state.use_gas(aux_data.extra_ergs_cost) else {
+                vm.state.current_frame.pc = spontaneous_panic();
+                return;
+            };
 
-        let start = query.location.index.0 * 32;
-        if query.rw_flag {
-            self.write_u256(page, start, query.value);
-        } else {
-            query.value = self[page].read_u256(start);
-            query.value_is_pointer = false;
-        }
-        query
-    }
+            #[allow(clippy::cast_possible_wrap)]
+            {
+                vm.world_diff.pubdata.0 += aux_data.extra_pubdata_cost as i32;
+            }
 
-    fn specialized_code_query(
-        &mut self,
-        _monotonic_cycle_counter: u32,
-        _query: zk_evm_abstractions::queries::MemoryQuery,
-    ) -> zk_evm_abstractions::queries::MemoryQuery {
-        todo!()
-    }
+            let mut abi = PrecompileCallAbi::from_u256(Register1::get(args, &mut vm.state));
+            if abi.memory_page_to_read.as_u32() == 0 {
+                abi.memory_page_to_read = vm.state.current_frame.heap;
+            }
+            if abi.memory_page_to_write.as_u32() == 0 {
+                abi.memory_page_to_write = vm.state.current_frame.heap;
+            }
 
-    fn read_code_query(
-        &self,
-        _monotonic_cycle_counter: u32,
-        _query: zk_evm_abstractions::queries::MemoryQuery,
-    ) -> zk_evm_abstractions::queries::MemoryQuery {
-        todo!()
-    }
+            let address_bytes = vm.state.current_frame.address.0;
+            let address_low = u16::from_le_bytes([address_bytes[19], address_bytes[18]]);
+            let heap_to_read = &vm.state.heaps[abi.memory_page_to_read];
+            let memory = PrecompileMemoryReader::new(
+                heap_to_read,
+                abi.input_memory_offset,
+                abi.input_memory_length,
+            );
+            let output = world.precompiles().call_precompile(
+                address_low,
+                memory,
+                abi.precompile_interpreted_data,
+            );
+
+            if let Some(cycle_stats) = output.cycle_stats {
+                tracer.on_extra_prover_cycles(cycle_stats);
+            }
+
+            let mut write_offset = abi.output_memory_offset * 32;
+            for i in 0..output.len.min(abi.output_memory_length) {
+                vm.state.heaps.write_u256(
+                    abi.memory_page_to_write,
+                    write_offset,
+                    output.buffer[i as usize],
+                );
+                write_offset += 32;
+            }
+            Register1::set(args, &mut vm.state, 1.into());
+        },
+    )
 }
 
 impl<T: Tracer, W: World<T>> Instruction<T, W> {

crates/vm2/src/lib.rs

@@ -20,6 +20,7 @@ pub use self::{
     vm::{Settings, VirtualMachine},
     world_diff::{Snapshot, StorageChange, WorldDiff},
 };
+use crate::precompiles::{LegacyPrecompiles, Precompiles};
 
 pub mod addressing_modes;
 #[cfg(not(feature = "single_instruction_test"))]
@@ -33,6 +34,7 @@ mod heap;
 mod instruction;
 mod instruction_handlers;
 mod mode_requirements;
+pub mod precompiles;
 mod predication;
 #[cfg(not(feature = "single_instruction_test"))]
 mod program;
@@ -98,6 +100,11 @@ pub trait World<T: Tracer>: StorageInterface + Sized {
 
     /// Loads bytecode bytes for the `decommit` opcode.
     fn decommit_code(&mut self, hash: U256) -> Vec<u8>;
+
+    /// Returns precompiles to be used.
+    fn precompiles(&self) -> &impl Precompiles {
+        &LegacyPrecompiles
+    }
 }
 
 /// Deterministic (across program runs and machines) hash that can be used for `Debug` implementations