extend mock runtime to allow executing constructors and cross contract calls

Signed-off-by: xermicus <cyrill@parity.io>

crates/integration/src/cases.rs

@@ -1,6 +1,8 @@
 use alloy_primitives::{I256, U256};
 use alloy_sol_types::{sol, SolCall};
 
+use crate::mock_runtime::{CallOutput, State};
+
 #[derive(Clone)]
 pub struct Contract {
     pub name: &'static str,
@@ -103,6 +105,24 @@ sol!(
 );
 
 impl Contract {
+    /// Execute the contract.
+    ///
+    /// Useful helper if the contract state can be ignored,
+    /// as it spares the deploy transaciton.
+    ///
+    /// - Inserts an account with given `code` into a new state.
+    /// - Callee and caller account will be `Transaction::default_address()`.
+    /// - Sets the calldata.
+    /// - Doesn't execute the constructor or deploy code.
+    /// - Calls the "call" export on a default backend config.
+    pub fn execute(&self) -> (State, CallOutput) {
+        State::default()
+            .transaction()
+            .with_default_account(&self.pvm_runtime)
+            .calldata(self.calldata.clone())
+            .call()
+    }
+
     pub fn baseline() -> Self {
         let code = include_str!("../contracts/Baseline.sol");
         let name = "Baseline";

crates/integration/src/lib.rs

@@ -1,5 +1,7 @@
 use cases::Contract;
-use mock_runtime::State;
+use mock_runtime::{CallOutput, State};
+
+use crate::mock_runtime::ReturnFlags;
 
 pub mod cases;
 pub mod mock_runtime;
@@ -73,15 +75,15 @@ pub fn compile_blob_with_options(
     hex::decode(bytecode).expect("hex encoding should always be valid")
 }
 
-pub fn assert_success(contract: Contract, differential: bool) -> State {
-    let (mut instance, export) = mock_runtime::prepare(&contract.pvm_runtime, None);
-    let state = mock_runtime::call(State::new(contract.calldata.clone()), &mut instance, export);
-    assert_eq!(state.output.flags, 0);
+pub fn assert_success(contract: &Contract, differential: bool) -> (State, CallOutput) {
+    let (state, output) = contract.execute();
+    assert_eq!(output.flags, ReturnFlags::Success);
 
     if differential {
-        let evm = revive_differential::prepare(contract.evm_runtime, contract.calldata);
-        assert_eq!(state.output.data.clone(), revive_differential::execute(evm));
+        let evm =
+            revive_differential::prepare(contract.evm_runtime.clone(), contract.calldata.clone());
+        assert_eq!(output.data.clone(), revive_differential::execute(evm));
     }
 
-    state
+    (state, output)
 }

crates/integration/src/mock_runtime.rs

@@ -1,72 +1,338 @@
 //! Mock environment used for integration tests.
-//! TODO: Switch to drink! once RISCV is ready in polkadot-sdk
 use std::collections::HashMap;
 
-use alloy_primitives::{Keccak256, U256};
+use alloy_primitives::{keccak256, Address, Keccak256, B256, U256};
 use polkavm::{
     Caller, Config, Engine, ExportIndex, GasMeteringKind, Instance, Linker, Module, ModuleConfig,
     ProgramBlob, Trap,
 };
 use revive_llvm_context::polkavm_const::runtime_api;
 
-#[derive(Default, Clone, Debug)]
-pub struct State {
-    pub input: Vec<u8>,
-    pub output: CallOutput,
-    pub value: u128,
-    pub storage: HashMap<U256, U256>,
+/// The mocked blockchain account.
+#[derive(Debug, Default, Clone)]
+struct Account {
+    value: U256,
+    contract: Option<U256>,
+    storage: HashMap<U256, U256>,
 }
 
-#[derive(Clone, Debug)]
-pub struct CallOutput {
-    pub flags: u32,
+/// Emitted event data.
+#[derive(Debug, Default, Clone)]
+pub struct Event {
     pub data: Vec<u8>,
+    pub topics: Vec<U256>,
 }
 
-impl Default for CallOutput {
-    fn default() -> Self {
-        Self {
-            flags: u32::MAX,
-            data: Vec::new(),
+/// The result of the contract call.
+#[derive(Debug, Default, Clone)]
+pub struct CallOutput {
+    /// The return flags.
+    pub flags: ReturnFlags,
+    /// The contract call output.
+    pub data: Vec<u8>,
+    /// The emitted events.
+    pub events: Event,
+}
+
+/// The contract blob export to be called.
+#[derive(Clone, Debug, Default)]
+enum Export {
+    #[default]
+    Call,
+    Deploy(U256),
+}
+
+/// Possible contract call return flags.
+#[derive(Debug, Default, Clone, PartialEq)]
+#[repr(u32)]
+pub enum ReturnFlags {
+    /// The contract execution returned normally.
+    Success = 0,
+    /// The contract execution returned normally but state changes should be reverted.
+    Revert = 1,
+    /// The contract trapped unexpectedly during execution.
+    #[default]
+    Trap = u32::MAX,
+}
+
+impl From<u32> for ReturnFlags {
+    fn from(value: u32) -> Self {
+        match value {
+            0 => Self::Success,
+            1 => Self::Revert,
+            u32::MAX => Self::Trap,
+            _ => panic!("invalid return flag: {value}"),
         }
     }
 }
 
+/// The local context inside the call stack.
+#[derive(Debug, Clone)]
+struct Frame {
+    /// The account that is being executed.
+    callee: Address,
+    /// The caller account.
+    caller: Address,
+    /// The value transferred with this transaction.
+    callvalue: U256,
+    /// The calldata for the contract execution.
+    input: Vec<u8>,
+    // The contract call output.
+    output: CallOutput,
+    /// The export to call.
+    export: Export,
+}
+
+impl Default for Frame {
+    fn default() -> Self {
+        Self {
+            callee: Transaction::default_address(),
+            caller: Transaction::default_address(),
+            callvalue: Default::default(),
+            input: Default::default(),
+            output: Default::default(),
+            export: Default::default(),
+        }
+    }
+}
+
+/// The transaction can modify the state by calling contracts.
+///
+/// Use the [TransactionBuilder] to create new transactions.
+#[derive(Default, Clone, Debug)]
+pub struct Transaction {
+    state: State,
+    stack: Vec<Frame>,
+}
+
+impl Transaction {
+    pub fn default_address() -> Address {
+        Address::default().create2(B256::default(), keccak256([]).0)
+    }
+
+    fn top_frame(&self) -> &Frame {
+        self.stack.last().expect("transactions should have a frame")
+    }
+
+    fn top_frame_mut(&mut self) -> &mut Frame {
+        self.stack
+            .last_mut()
+            .expect("transactions should have a frame")
+    }
+
+    fn top_account_mut(&mut self) -> &mut Account {
+        let account = self.top_frame_mut().callee;
+        self.state
+            .accounts
+            .get_mut(&account)
+            .unwrap_or_else(|| panic!("callee has no associated account: {account}"))
+    }
+
+    fn create2(&self, salt: U256, blob_hash: U256) -> Address {
+        self.top_frame()
+            .callee
+            .create2(salt.to_be_bytes(), blob_hash.to_be_bytes())
+    }
+}
+
+/// Helper to create valid transactions.
+#[derive(Default, Clone, Debug)]
+pub struct TransactionBuilder {
+    context: Transaction,
+    state_before: State,
+}
+
+impl TransactionBuilder {
+    /// Set the caller account.
+    pub fn caller(mut self, account: Address) -> Self {
+        self.context.top_frame_mut().caller = account;
+        self
+    }
+
+    /// Set the callee account.
+    pub fn callee(mut self, account: Address) -> Self {
+        self.context.top_frame_mut().callee = account;
+        self
+    }
+
+    /// Set the transferred callvalue.
+    pub fn callvalue(mut self, amount: U256) -> Self {
+        self.context.top_frame_mut().callvalue = amount;
+        self
+    }
+
+    /// Set the calldata.
+    pub fn calldata(mut self, data: Vec<u8>) -> Self {
+        self.context.top_frame_mut().input = data;
+        self
+    }
+
+    /// Set the transaction to deploy code.
+    pub fn deploy(mut self, code: &[u8]) -> Self {
+        let blob_hash = self.context.state.upload_code(code);
+        self.context.top_frame_mut().export = Export::Deploy(blob_hash);
+        self
+    }
+
+    /// Set the account at [Transaction::default_address] to the given `code`.
+    ///
+    /// Useful helper to spare the deploy transaction.
+    pub fn with_default_account(mut self, code: &[u8]) -> Self {
+        self.context.state.upload_code(code);
+        self.context.state.create_account(
+            Transaction::default_address(),
+            Default::default(),
+            keccak256(code).into(),
+        );
+        self
+    }
+
+    /// Execute the transaction with a default config backend.
+    ///
+    /// Reverts any state changes if the contract reverts or the exuection traps.
+    pub fn call(mut self) -> (State, CallOutput) {
+        let blob_hash = match self.context.top_frame().export {
+            Export::Call => self.context.top_account_mut().contract.unwrap_or_else(|| {
+                panic!(
+                    "not a contract account: {}",
+                    self.context.top_frame().callee
+                )
+            }),
+            Export::Deploy(blob_hash) => blob_hash,
+        };
+        let code = self
+            .context
+            .state
+            .blobs
+            .get(&blob_hash)
+            .unwrap_or_else(|| panic!("contract code not found: {blob_hash}"));
+        let (mut instance, export) = prepare(code, None);
+        self.call_on(&mut instance, export)
+    }
+
+    /// Execute the transaction on a given instance and export.
+    /// The `instance` and `export` are expected to match that of the `Transaction`.
+
+    /// Reverts any state changes if the contract reverts or the exuection traps.
+    pub fn call_on(
+        mut self,
+        instance: &mut Instance<Transaction>,
+        export: ExportIndex,
+    ) -> (State, CallOutput) {
+        let mut state_args = polkavm::StateArgs::default();
+        state_args.set_gas(polkavm::Gas::MAX);
+
+        if let Export::Deploy(blob_hash) = self.context.top_frame().export {
+            self.context.state.create_account(
+                self.context.create2(Default::default(), blob_hash),
+                self.context.top_frame().callvalue,
+                blob_hash,
+            );
+        }
+
+        let callvalue = self.context.top_frame().callvalue;
+        self.context.top_account_mut().value += callvalue;
+
+        let call_args = polkavm::CallArgs::new(&mut self.context, export);
+
+        init_logs();
+
+        match instance.call(state_args, call_args) {
+            Err(polkavm::ExecutionError::Trap(_)) => self.finalize(),
+            Err(other) => panic!("unexpected error: {other}"),
+            Ok(_) => panic!("unexpected return"),
+        }
+    }
+
+    /// Commits or reverts the state changes based on the call flags.
+    fn finalize(mut self) -> (State, CallOutput) {
+        let state = match self.context.top_frame().output.flags {
+            ReturnFlags::Success => self.context.state,
+            _ => self.state_before,
+        };
+        let output = self.context.stack.pop().unwrap().output;
+        (state, output)
+    }
+}
+
+impl From<State> for TransactionBuilder {
+    fn from(state: State) -> Self {
+        TransactionBuilder {
+            state_before: state.clone(),
+            context: Transaction {
+                state,
+                stack: Default::default(),
+            },
+        }
+    }
+}
+
+/// The mocked blockchain state.
+#[derive(Default, Clone, Debug)]
+pub struct State {
+    blobs: HashMap<U256, Vec<u8>>,
+    accounts: HashMap<Address, Account>,
+}
+
 impl State {
-    pub const ADDRESS: [u8; 20] = [1; 20];
     pub const BLOCK_NUMBER: u64 = 123;
     pub const BLOCK_TIMESTAMP: u64 = 456;
-    pub const CALLER: [u8; 20] = [2; 20];
 
-    pub fn new(input: Vec<u8>) -> Self {
-        Self {
-            input,
-            ..Default::default()
+    pub fn transaction(self) -> TransactionBuilder {
+        TransactionBuilder {
+            state_before: self.clone(),
+            context: Transaction {
+                state: self,
+                stack: vec![Default::default()],
+            },
         }
     }
 
-    pub fn reset_output(&mut self) {
-        self.output = Default::default();
+    pub fn upload_code(&mut self, code: &[u8]) -> U256 {
+        let blob_hash = keccak256(code).into();
+        self.blobs.insert(blob_hash, code.to_vec());
+        blob_hash
     }
 
-    pub fn assert_storage_key(&self, at: U256, expect: U256) {
-        assert_eq!(self.storage[&at], expect);
+    pub fn assert_storage_key(&self, account: Address, key: U256, expected: U256) {
+        assert_eq!(
+            self.accounts
+                .get(&account)
+                .unwrap_or_else(|| panic!("unknown account: {account}"))
+                .storage
+                .get(&key)
+                .copied()
+                .unwrap_or_default(),
+            expected
+        );
+    }
+
+    pub fn create_account(&mut self, address: Address, value: U256, blob_hash: U256) {
+        self.accounts.insert(
+            address,
+            Account {
+                value,
+                contract: Some(blob_hash),
+                storage: HashMap::new(),
+            },
+        );
     }
 }
 
-fn link_host_functions(engine: &Engine) -> Linker<State> {
+fn link_host_functions(engine: &Engine) -> Linker<Transaction> {
     let mut linker = Linker::new(engine);
 
     linker
         .func_wrap(
             runtime_api::INPUT,
-            |caller: Caller<State>, out_ptr: u32, out_len_ptr: u32| -> Result<(), Trap> {
-                let (mut caller, state) = caller.split();
+            |caller: Caller<Transaction>, out_ptr: u32, out_len_ptr: u32| -> Result<(), Trap> {
+                let (mut caller, transaction) = caller.split();
 
-                assert!(state.input.len() <= caller.read_u32(out_len_ptr).unwrap() as usize);
+                let input = &transaction.top_frame().input;
+                assert!(input.len() <= caller.read_u32(out_len_ptr).unwrap() as usize);
 
-                caller.write_memory(out_ptr, &state.input)?;
-                caller.write_memory(out_len_ptr, &(state.input.len() as u32).to_le_bytes())?;
+                caller.write_memory(out_ptr, input)?;
+                caller.write_memory(out_len_ptr, &(input.len() as u32).to_le_bytes())?;
 
                 Ok(())
             },
@@ -76,11 +342,16 @@ fn link_host_functions(engine: &Engine) -> Linker<State> {
     linker
         .func_wrap(
             runtime_api::RETURN,
-            |caller: Caller<State>, flags: u32, data_ptr: u32, data_len: u32| -> Result<(), Trap> {
-                let (caller, state) = caller.split();
+            |caller: Caller<Transaction>,
+             flags: u32,
+             data_ptr: u32,
+             data_len: u32|
+             -> Result<(), Trap> {
+                let (caller, transaction) = caller.split();
 
-                state.output.flags = flags;
-                state.output.data = caller.read_memory_into_vec(data_ptr, data_len)?;
+                let frame = transaction.top_frame_mut();
+                frame.output.flags = flags.into();
+                frame.output.data = caller.read_memory_into_vec(data_ptr, data_len)?;
 
                 Err(Default::default())
             },
@@ -90,8 +361,8 @@ fn link_host_functions(engine: &Engine) -> Linker<State> {
     linker
         .func_wrap(
             runtime_api::VALUE_TRANSFERRED,
-            |caller: Caller<State>, out_ptr: u32, out_len_ptr: u32| -> Result<(), Trap> {
-                let (mut caller, state) = caller.split();
+            |caller: Caller<Transaction>, out_ptr: u32, out_len_ptr: u32| -> Result<(), Trap> {
+                let (mut caller, transaction) = caller.split();
 
                 let out_len = caller.read_u32(out_len_ptr)? as usize;
                 assert_eq!(
@@ -100,8 +371,7 @@ fn link_host_functions(engine: &Engine) -> Linker<State> {
                     "spurious output buffer size: {out_len}"
                 );
 
-                let value = state.value.to_le_bytes();
-
+                let value = transaction.top_frame().callvalue.as_le_bytes();
                 caller.write_memory(out_ptr, &value)?;
                 caller.write_memory(out_len_ptr, &(value.len() as u32).to_le_bytes())?;
 
@@ -113,7 +383,7 @@ fn link_host_functions(engine: &Engine) -> Linker<State> {
     linker
         .func_wrap(
             "debug_message",
-            |caller: Caller<State>, str_ptr: u32, str_len: u32| -> Result<u32, Trap> {
+            |caller: Caller<Transaction>, str_ptr: u32, str_len: u32| -> Result<u32, Trap> {
                 let (caller, _) = caller.split();
 
                 let data = caller.read_memory_into_vec(str_ptr, str_len)?;
@@ -127,13 +397,13 @@ fn link_host_functions(engine: &Engine) -> Linker<State> {
     linker
         .func_wrap(
             runtime_api::SET_STORAGE,
-            |caller: Caller<State>,
+            |caller: Caller<Transaction>,
              key_ptr: u32,
              key_len: u32,
              value_ptr: u32,
              value_len: u32|
              -> Result<u32, Trap> {
-                let (caller, state) = caller.split();
+                let (caller, transaction) = caller.split();
 
                 assert_eq!(
                     key_len as usize,
@@ -149,7 +419,7 @@ fn link_host_functions(engine: &Engine) -> Linker<State> {
                 let key = caller.read_memory_into_vec(key_ptr, key_len)?;
                 let value = caller.read_memory_into_vec(value_ptr, value_len)?;
 
-                state.storage.insert(
+                transaction.top_account_mut().storage.insert(
                     U256::from_be_bytes::<32>(key.try_into().unwrap()),
                     U256::from_be_bytes::<32>(value.try_into().unwrap()),
                 );
@@ -162,13 +432,13 @@ fn link_host_functions(engine: &Engine) -> Linker<State> {
     linker
         .func_wrap(
             runtime_api::GET_STORAGE,
-            |caller: Caller<State>,
+            |caller: Caller<Transaction>,
              key_ptr: u32,
              key_len: u32,
              out_ptr: u32,
              out_len_ptr: u32|
              -> Result<u32, Trap> {
-                let (mut caller, state) = caller.split();
+                let (mut caller, transaction) = caller.split();
 
                 let key = caller.read_memory_into_vec(key_ptr, key_len)?;
                 let out_len = caller.read_u32(out_len_ptr)? as usize;
@@ -178,7 +448,8 @@ fn link_host_functions(engine: &Engine) -> Linker<State> {
                     "spurious output buffer size: {out_len}"
                 );
 
-                let value = state
+                let value = transaction
+                    .top_account_mut()
                     .storage
                     .get(&U256::from_be_bytes::<32>(key.try_into().unwrap()))
                     .map(U256::to_be_bytes::<32>)
@@ -195,7 +466,7 @@ fn link_host_functions(engine: &Engine) -> Linker<State> {
     linker
         .func_wrap(
             runtime_api::HASH_KECCAK_256,
-            |caller: Caller<State>,
+            |caller: Caller<Transaction>,
              input_ptr: u32,
              input_len: u32,
              out_ptr: u32|
@@ -216,7 +487,7 @@ fn link_host_functions(engine: &Engine) -> Linker<State> {
     linker
         .func_wrap(
             runtime_api::NOW,
-            |caller: Caller<State>, out_ptr: u32, out_len_ptr: u32| {
+            |caller: Caller<Transaction>, out_ptr: u32, out_len_ptr: u32| {
                 let (mut caller, _) = caller.split();
 
                 let out_len = caller.read_u32(out_len_ptr)? as usize;
@@ -237,7 +508,7 @@ fn link_host_functions(engine: &Engine) -> Linker<State> {
     linker
         .func_wrap(
             runtime_api::BLOCK_NUMBER,
-            |caller: Caller<State>, out_ptr: u32, out_len_ptr: u32| {
+            |caller: Caller<Transaction>, out_ptr: u32, out_len_ptr: u32| {
                 let (mut caller, _) = caller.split();
 
                 let out_len = caller.read_u32(out_len_ptr)? as usize;
@@ -258,18 +529,19 @@ fn link_host_functions(engine: &Engine) -> Linker<State> {
     linker
         .func_wrap(
             runtime_api::ADDRESS,
-            |caller: Caller<State>, out_ptr: u32, out_len_ptr: u32| {
-                let (mut caller, _) = caller.split();
+            |caller: Caller<Transaction>, out_ptr: u32, out_len_ptr: u32| {
+                let (mut caller, transaction) = caller.split();
 
                 let out_len = caller.read_u32(out_len_ptr)? as usize;
                 assert_eq!(
                     out_len,
-                    revive_common::BYTE_LENGTH_WORD,
+                    revive_common::BYTE_LENGTH_ETH_ADDRESS,
                     "spurious output buffer size: {out_len}"
                 );
 
-                caller.write_memory(out_ptr, &State::ADDRESS)?;
-                caller.write_memory(out_len_ptr, &(State::ADDRESS.len() as u32).to_le_bytes())?;
+                let address = transaction.top_frame().callee.as_slice();
+                caller.write_memory(out_ptr, address)?;
+                caller.write_memory(out_len_ptr, &(address.len() as u32).to_le_bytes())?;
 
                 Ok(())
             },
@@ -279,18 +551,19 @@ fn link_host_functions(engine: &Engine) -> Linker<State> {
     linker
         .func_wrap(
             runtime_api::CALLER,
-            |caller: Caller<State>, out_ptr: u32, out_len_ptr: u32| {
-                let (mut caller, _) = caller.split();
+            |caller: Caller<Transaction>, out_ptr: u32, out_len_ptr: u32| {
+                let (mut caller, transaction) = caller.split();
 
                 let out_len = caller.read_u32(out_len_ptr)? as usize;
                 assert_eq!(
                     out_len,
-                    revive_common::BYTE_LENGTH_WORD,
+                    revive_common::BYTE_LENGTH_ETH_ADDRESS,
                     "spurious output buffer size: {out_len}"
                 );
 
-                caller.write_memory(out_ptr, &State::CALLER)?;
-                caller.write_memory(out_len_ptr, &(State::CALLER.len() as u32).to_le_bytes())?;
+                let address = transaction.top_frame().caller.as_slice();
+                caller.write_memory(out_ptr, address)?;
+                caller.write_memory(out_len_ptr, &(address.len() as u32).to_le_bytes())?;
 
                 Ok(())
             },
@@ -311,7 +584,10 @@ pub fn recompile_code(code: &[u8], engine: &Engine) -> Module {
     Module::new(engine, &module_config, code).unwrap()
 }
 
-pub fn instantiate_module(module: &Module, engine: &Engine) -> (Instance<State>, ExportIndex) {
+pub fn instantiate_module(
+    module: &Module,
+    engine: &Engine,
+) -> (Instance<Transaction>, ExportIndex) {
     let export = module.lookup_export(runtime_api::CALL).unwrap();
     let func = link_host_functions(engine).instantiate_pre(module).unwrap();
     let instance = func.instantiate().unwrap();
@@ -319,7 +595,7 @@ pub fn instantiate_module(module: &Module, engine: &Engine) -> (Instance<State>,
     (instance, export)
 }
 
-pub fn prepare(code: &[u8], config: Option<Config>) -> (Instance<State>, ExportIndex) {
+pub fn prepare(code: &[u8], config: Option<Config>) -> (Instance<Transaction>, ExportIndex) {
     let blob = ProgramBlob::parse(code).unwrap();
 
     let engine = Engine::new(&config.unwrap_or_default()).unwrap();
@@ -337,23 +613,6 @@ pub fn prepare(code: &[u8], config: Option<Config>) -> (Instance<State>, ExportI
     (instance, export)
 }
 
-pub fn call(mut state: State, on: &mut Instance<State>, export: ExportIndex) -> State {
-    state.reset_output();
-
-    let mut state_args = polkavm::StateArgs::default();
-    state_args.set_gas(polkavm::Gas::MAX);
-
-    let call_args = polkavm::CallArgs::new(&mut state, export);
-
-    init_logs();
-
-    match on.call(state_args, call_args) {
-        Err(polkavm::ExecutionError::Trap(_)) => state,
-        Err(other) => panic!("unexpected error: {other}"),
-        Ok(_) => panic!("unexpected return"),
-    }
-}
-
 fn init_logs() {
     if std::env::var("RUST_LOG").is_ok() {
         #[cfg(test)]

crates/integration/src/tests.rs

@@ -1,4 +1,4 @@
-use alloy_primitives::{Address, FixedBytes, Keccak256, I256, U256};
+use alloy_primitives::{keccak256, Address, FixedBytes, I256, U256};
 use alloy_sol_types::{sol, SolCall};
 use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
 use sha1::Digest;
@@ -6,7 +6,7 @@ use sha1::Digest;
 use crate::{
     assert_success,
     cases::Contract,
-    mock_runtime::{self, State},
+    mock_runtime::{self, ReturnFlags, State, Transaction},
 };
 
 #[test]
@@ -18,8 +18,8 @@ fn fibonacci() {
         Contract::fib_iterative(parameter),
         Contract::fib_binet(parameter),
     ] {
-        let state = assert_success(contract, true);
-        let received = U256::from_be_bytes::<32>(state.output.data.try_into().unwrap());
+        let (_, output) = assert_success(&contract, true);
+        let received = U256::from_be_bytes::<32>(output.data.try_into().unwrap());
         let expected = U256::from(8);
         assert_eq!(received, expected);
     }
@@ -28,15 +28,15 @@ fn fibonacci() {
 #[test]
 fn flipper() {
     let contract = Contract::flipper();
-    let (mut instance, export) = mock_runtime::prepare(&contract.pvm_runtime, None);
+    let address = Transaction::default_address();
 
-    let state = crate::mock_runtime::call(State::new(contract.calldata), &mut instance, export);
-    assert_eq!(state.output.flags, 0);
-    assert_eq!(state.storage[&U256::ZERO], U256::try_from(1).unwrap());
+    let (state, output) = contract.execute();
+    assert_eq!(output.flags, ReturnFlags::Success);
+    state.assert_storage_key(address, U256::ZERO, U256::from(1));
 
-    let state = crate::mock_runtime::call(state, &mut instance, export);
-    assert_eq!(state.output.flags, 0);
-    assert_eq!(state.storage[&U256::ZERO], U256::ZERO);
+    let (state, output) = state.transaction().calldata(contract.calldata).call();
+    assert_eq!(output.flags, ReturnFlags::Success);
+    state.assert_storage_key(address, U256::ZERO, U256::ZERO);
 }
 
 #[test]
@@ -57,16 +57,16 @@ fn hash_keccak_256() {
     let param = "hello";
     let input = TestSha3::testCall::new((param.to_string(),)).abi_encode();
 
-    let state = State::new(input);
-    let (mut instance, export) = mock_runtime::prepare(&code, None);
-    let state = crate::mock_runtime::call(state, &mut instance, export);
+    let (_, output) = State::default()
+        .transaction()
+        .with_default_account(&code)
+        .calldata(input)
+        .call();
 
-    assert_eq!(state.output.flags, 0);
+    assert_eq!(output.flags, ReturnFlags::Success);
 
-    let mut hasher = Keccak256::new();
-    hasher.update(param);
-    let expected = hasher.finalize();
-    let received = FixedBytes::<32>::from_slice(&state.output.data);
+    let expected = keccak256(param.as_bytes());
+    let received = FixedBytes::<32>::from_slice(&output.data);
     assert_eq!(received, expected);
 }
 
@@ -77,16 +77,16 @@ fn erc20() {
 
 #[test]
 fn triangle_number() {
-    let state = assert_success(Contract::triangle_number(13), true);
-    let received = U256::from_be_bytes::<32>(state.output.data.try_into().unwrap());
+    let (_, output) = assert_success(&Contract::triangle_number(13), true);
+    let received = U256::from_be_bytes::<32>(output.data.try_into().unwrap());
     let expected = U256::try_from(91).unwrap();
     assert_eq!(received, expected);
 }
 
 #[test]
 fn odd_product() {
-    let state = assert_success(Contract::odd_product(5), true);
-    let received = I256::from_be_bytes::<32>(state.output.data.try_into().unwrap());
+    let (_, output) = assert_success(&Contract::odd_product(5), true);
+    let received = I256::from_be_bytes::<32>(output.data.try_into().unwrap());
     let expected = I256::try_from(945i64).unwrap();
     assert_eq!(received, expected);
 }
@@ -105,16 +105,19 @@ fn msize_plain() {
         false,
         revive_solidity::SolcPipeline::EVMLA,
     );
-    let (mut instance, export) = mock_runtime::prepare(&code, None);
 
     let input = MSize::mSizeCall::new(()).abi_encode();
-    let state = crate::mock_runtime::call(State::new(input), &mut instance, export);
+    let (_, output) = State::default()
+        .transaction()
+        .calldata(input)
+        .with_default_account(&code)
+        .call();
 
-    assert_eq!(state.output.flags, 0);
+    assert_eq!(output.flags, ReturnFlags::Success);
 
     // Solidity always stores the "free memory pointer" (32 byte int) at offset 64.
     let expected = U256::try_from(64 + 32).unwrap();
-    let received = U256::from_be_bytes::<32>(state.output.data.try_into().unwrap());
+    let received = U256::from_be_bytes::<32>(output.data.try_into().unwrap());
     assert_eq!(received, expected);
 }
 
@@ -126,16 +129,18 @@ fn transferred_value() {
         }
     );
     let code = crate::compile_blob("Value", include_str!("../contracts/Value.sol"));
-    let mut state = State::new(Value::valueCall::SELECTOR.to_vec());
-    state.value = 123;
 
-    let (mut instance, export) = mock_runtime::prepare(&code, None);
-    let state = crate::mock_runtime::call(state, &mut instance, export);
+    let (_, output) = State::default()
+        .transaction()
+        .calldata(Value::valueCall::SELECTOR.to_vec())
+        .callvalue(U256::from(123))
+        .with_default_account(&code)
+        .call();
 
-    assert_eq!(state.output.flags, 0);
+    assert_eq!(output.flags, ReturnFlags::Success);
 
-    let expected = I256::try_from(state.value).unwrap();
-    let received = I256::from_be_bytes::<32>(state.output.data.try_into().unwrap());
+    let expected = I256::try_from(123).unwrap();
+    let received = I256::from_be_bytes::<32>(output.data.try_into().unwrap());
     assert_eq!(received, expected);
 }
 
@@ -153,18 +158,21 @@ fn msize_non_word_sized_access() {
         false,
         revive_solidity::SolcPipeline::Yul,
     );
-    let (mut instance, export) = mock_runtime::prepare(&code, None);
 
     let input = MSize::mStore100Call::new(()).abi_encode();
-    let state = crate::mock_runtime::call(State::new(input), &mut instance, export);
+    let (_, output) = State::default()
+        .transaction()
+        .with_default_account(&code)
+        .calldata(input)
+        .call();
 
-    assert_eq!(state.output.flags, 0);
+    assert_eq!(output.flags, ReturnFlags::Success);
 
     // https://docs.zksync.io/build/developer-reference/differences-with-ethereum.html#mstore-mload
     // "Unlike EVM, where the memory growth is in words, on zkEVM the memory growth is counted in bytes."
     // "For example, if you write mstore(100, 0) the msize on zkEVM will be 132, but on the EVM it will be 160."
     let expected = U256::try_from(132).unwrap();
-    let received = U256::from_be_bytes::<32>(state.output.data.try_into().unwrap());
+    let received = U256::from_be_bytes::<32>(output.data.try_into().unwrap());
     assert_eq!(received, expected);
 }
 
@@ -232,8 +240,8 @@ fn mstore8() {
     ]
     .par_iter()
     .map(|(parameter, expected)| {
-        let state = assert_success(Contract::mstore8(*parameter), true);
-        let received = U256::from_be_bytes::<32>(state.output.data.try_into().unwrap());
+        let (_, output) = assert_success(&Contract::mstore8(*parameter), true);
+        let received = U256::from_be_bytes::<32>(output.data.try_into().unwrap());
         (received, *expected)
     })
     .collect::<Vec<_>>()
@@ -249,41 +257,42 @@ fn sha1() {
     hasher.update(&pre);
     let hash = hasher.finalize();
 
-    let state = assert_success(Contract::sha1(pre), true);
+    let (_, output) = assert_success(&Contract::sha1(pre), true);
     let expected = FixedBytes::<20>::from_slice(&hash[..]);
-    let received = FixedBytes::<20>::from_slice(&state.output.data[..20]);
+    let received = FixedBytes::<20>::from_slice(&output.data[..20]);
     assert_eq!(received, expected);
 }
 
 #[test]
 fn block_number() {
-    let state = assert_success(Contract::block_number(), true);
-    let received = U256::from_be_bytes::<32>(state.output.data.try_into().unwrap());
+    let (_, output) = assert_success(&Contract::block_number(), true);
+    let received = U256::from_be_bytes::<32>(output.data.try_into().unwrap());
     let expected = U256::from(mock_runtime::State::BLOCK_NUMBER);
     assert_eq!(received, expected);
 }
 
 #[test]
 fn block_timestamp() {
-    let state = assert_success(Contract::block_timestamp(), true);
-    let received = U256::from_be_bytes::<32>(state.output.data.try_into().unwrap());
+    let (_, output) = assert_success(&Contract::block_timestamp(), true);
+    let received = U256::from_be_bytes::<32>(output.data.try_into().unwrap());
     let expected = U256::from(mock_runtime::State::BLOCK_TIMESTAMP);
     assert_eq!(received, expected);
 }
 
 #[test]
 fn address() {
-    let state = assert_success(Contract::context_address(), true);
-    let received = Address::from_slice(&state.output.data[12..]);
-    let expected = Address::from(&mock_runtime::State::ADDRESS);
+    let contract = Contract::context_address();
+    let (_, output) = assert_success(&contract, true);
+    let received = Address::from_slice(&output.data[12..]);
+    let expected = Transaction::default_address();
     assert_eq!(received, expected);
 }
 
 #[test]
 fn caller() {
-    let state = assert_success(Contract::context_caller(), true);
-    let received = Address::from_slice(&state.output.data[12..]);
-    let expected = Address::from(&mock_runtime::State::CALLER);
+    let (_, output) = assert_success(&Contract::context_caller(), true);
+    let received = Address::from_slice(&output.data[12..]);
+    let expected = Transaction::default_address();
     assert_eq!(received, expected);
 }
 
@@ -301,8 +310,8 @@ fn unsigned_division() {
     ]
     .par_iter()
     .map(|(n, d, q)| {
-        let state = assert_success(Contract::division_arithmetics_div(*n, *d), true);
-        let received = U256::from_be_bytes::<32>(state.output.data.try_into().unwrap());
+        let (_, output) = assert_success(&Contract::division_arithmetics_div(*n, *d), true);
+        let received = U256::from_be_bytes::<32>(output.data.try_into().unwrap());
         (received, *q)
     })
     .collect::<Vec<_>>()
@@ -333,8 +342,8 @@ fn signed_division() {
     ]
     .par_iter()
     .map(|(n, d, q)| {
-        let state = assert_success(Contract::division_arithmetics_sdiv(*n, *d), true);
-        let received = I256::from_be_bytes::<32>(state.output.data.try_into().unwrap());
+        let (_, output) = assert_success(&Contract::division_arithmetics_sdiv(*n, *d), true);
+        let received = I256::from_be_bytes::<32>(output.data.try_into().unwrap());
         (received, *q)
     })
     .collect::<Vec<_>>()
@@ -359,8 +368,8 @@ fn unsigned_remainder() {
     ]
     .par_iter()
     .map(|(n, d, q)| {
-        let state = assert_success(Contract::division_arithmetics_mod(*n, *d), true);
-        let received = U256::from_be_bytes::<32>(state.output.data.try_into().unwrap());
+        let (_, output) = assert_success(&Contract::division_arithmetics_mod(*n, *d), true);
+        let received = U256::from_be_bytes::<32>(output.data.try_into().unwrap());
         (received, *q)
     })
     .collect::<Vec<_>>()
@@ -397,8 +406,8 @@ fn signed_remainder() {
     ]
     .par_iter()
     .map(|(n, d, q)| {
-        let state = assert_success(Contract::division_arithmetics_smod(*n, *d), true);
-        let received = I256::from_be_bytes::<32>(state.output.data.try_into().unwrap());
+        let (_, output) = assert_success(&Contract::division_arithmetics_smod(*n, *d), true);
+        let received = I256::from_be_bytes::<32>(output.data.try_into().unwrap());
         (received, *q)
     })
     .collect::<Vec<_>>()