extend mock runtime to allow executing constructors and cross contract calls

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

Cargo.lock

@@ -68,7 +68,7 @@ dependencies = [
  "proc-macro-error",
  "proc-macro2",
  "quote",
- "syn 2.0.61",
+ "syn 2.0.63",
  "syn-solidity",
  "tiny-keccak",
 ]
@@ -284,7 +284,7 @@ checksum = "3c87f3f15e7794432337fc718554eaa4dc8f04c9677a950ffe366f20a162ae42"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn 2.0.61",
+ "syn 2.0.63",
 ]
 
 [[package]]
@@ -473,7 +473,7 @@ dependencies = [
  "heck 0.5.0",
  "proc-macro2",
  "quote",
- "syn 2.0.61",
+ "syn 2.0.63",
 ]
 
 [[package]]
@@ -1047,7 +1047,7 @@ source = "git+https://github.com/TheDan64/inkwell.git#6c0fb56b3554e939f9ca61b465
 dependencies = [
  "proc-macro2",
  "quote",
- "syn 2.0.61",
+ "syn 2.0.63",
 ]
 
 [[package]]
@@ -1712,6 +1712,7 @@ dependencies = [
 name = "revive-differential"
 version = "0.1.0"
 dependencies = [
+ "alloy-primitives",
  "evm-interpreter",
  "primitive-types",
 ]
@@ -2025,7 +2026,7 @@ checksum = "c5e405930b9796f1c00bee880d03fc7e0bb4b9a11afc776885ffe84320da2865"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn 2.0.61",
+ "syn 2.0.63",
 ]
 
 [[package]]
@@ -2197,9 +2198,9 @@ dependencies = [
 
 [[package]]
 name = "syn"
-version = "2.0.61"
+version = "2.0.63"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c993ed8ccba56ae856363b1845da7266a7cb78e1d146c8a32d54b45a8b831fc9"
+checksum = "bf5be731623ca1a1fb7d8be6f261a3be6d3e2337b8a1f97be944d020c8fcb704"
 dependencies = [
  "proc-macro2",
  "quote",
@@ -2215,7 +2216,7 @@ dependencies = [
  "paste",
  "proc-macro2",
  "quote",
- "syn 2.0.61",
+ "syn 2.0.63",
 ]
 
 [[package]]
@@ -2262,7 +2263,7 @@ checksum = "e2470041c06ec3ac1ab38d0356a6119054dedaea53e12fbefc0de730a1c08524"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn 2.0.61",
+ "syn 2.0.63",
 ]
 
 [[package]]
@@ -2413,7 +2414,7 @@ dependencies = [
  "once_cell",
  "proc-macro2",
  "quote",
- "syn 2.0.61",
+ "syn 2.0.63",
  "wasm-bindgen-shared",
 ]
 
@@ -2435,7 +2436,7 @@ checksum = "e94f17b526d0a461a191c78ea52bbce64071ed5c04c9ffe424dcb38f74171bb7"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn 2.0.61",
+ "syn 2.0.63",
  "wasm-bindgen-backend",
  "wasm-bindgen-shared",
 ]
@@ -2674,7 +2675,7 @@ checksum = "15e934569e47891f7d9411f1a451d947a60e000ab3bd24fbb970f000387d1b3b"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn 2.0.61",
+ "syn 2.0.63",
 ]
 
 [[package]]
@@ -2694,7 +2695,7 @@ checksum = "ce36e65b0d2999d2aafac989fb249189a141aee1f53c612c1f37d72631959f69"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn 2.0.61",
+ "syn 2.0.63",
 ]
 
 [[package]]

crates/benchmarks/benches/execute.rs

@@ -31,14 +31,14 @@ where
         #[cfg(feature = "bench-pvm-interpreter")]
         {
             let contract = contract(p.clone());
-            let (state, mut instance, export) = revive_benchmarks::prepare_pvm(
+            let (transaction, mut instance, export) = revive_benchmarks::prepare_pvm(
                 &contract.pvm_runtime,
-                &contract.calldata,
+                contract.calldata,
                 polkavm::BackendKind::Interpreter,
             );
             group.bench_with_input(BenchmarkId::new("PVMInterpreter", l), p, |b, _| {
                 b.iter(|| {
-                    revive_integration::mock_runtime::call(state.clone(), &mut instance, export);
+                    transaction.clone().call_on(&mut instance, export);
                 });
             });
         }
@@ -46,14 +46,14 @@ where
         #[cfg(feature = "bench-pvm")]
         {
             let contract = contract(p.clone());
-            let (state, mut instance, export) = revive_benchmarks::prepare_pvm(
+            let (transaction, mut instance, export) = revive_benchmarks::prepare_pvm(
                 &contract.pvm_runtime,
-                &contract.calldata,
+                contract.calldata,
                 polkavm::BackendKind::Compiler,
             );
             group.bench_with_input(BenchmarkId::new("PVM", l), p, |b, _| {
                 b.iter(|| {
-                    revive_integration::mock_runtime::call(state.clone(), &mut instance, export);
+                    transaction.clone().call_on(&mut instance, export);
                 });
             });
         }

crates/benchmarks/src/lib.rs

@@ -1,19 +1,20 @@
 use polkavm::{BackendKind, Config, Engine, ExportIndex, Instance, SandboxKind};
-use revive_integration::mock_runtime;
+use revive_integration::mock_runtime::{self, TransactionBuilder};
-use revive_integration::mock_runtime::State;
+use revive_integration::mock_runtime::{State, Transaction};
 
 pub fn prepare_pvm(
     code: &[u8],
-    input: &[u8],
+    input: Vec<u8>,
     backend: BackendKind,
-) -> (State, Instance<State>, ExportIndex) {
+) -> (TransactionBuilder, Instance<Transaction>, ExportIndex) {
     let mut config = Config::new();
     config.set_backend(Some(backend));
     config.set_sandbox(Some(SandboxKind::Linux));
 
     let (instance, export_index) = mock_runtime::prepare(code, Some(config));
+    let transaction = State::default().transaction().calldata(input);
 
-    (State::new(input.to_vec()), instance, export_index)
+    (transaction, instance, export_index)
 }
 
 pub fn instantiate_engine(backend: BackendKind) -> Engine {

crates/common/src/byte_length.rs

@@ -19,7 +19,7 @@ pub const BYTE_LENGTH_ETH_ADDRESS: usize = 20;
 pub const BYTE_LENGTH_WORD: usize = 32;
 
 /// Byte length of the runtime value type.
-pub const BYTE_LENGTH_VALUE: usize = 16;
+pub const BYTE_LENGTH_VALUE: usize = 32;
 
 /// Byte length of the runtime block number type.
 pub const BYTE_LENGTH_BLOCK_NUMBER: usize = 8;

crates/differential/Cargo.toml

@@ -7,4 +7,5 @@ edition = "2021"
 
 [dependencies]
 evm-interpreter = { workspace = true }
-primitive-types = { workspace = true }
+primitive-types = { workspace = true }
+alloy-primitives = { workspace = true }

crates/differential/src/lib.rs

@@ -1,3 +1,4 @@
+use alloy_primitives::{keccak256, Address, B256};
 use evm_interpreter::{
     interpreter::{EtableInterpreter, RunInterpreter},
     trap::CallCreateTrap,
@@ -123,10 +124,13 @@ pub struct PreparedEvm {
 }
 
 pub fn prepare(code: Vec<u8>, data: Vec<u8>) -> PreparedEvm {
+    let address = Address::default().create2(B256::default(), keccak256([]).0);
+    let caller = Address::default().create2(B256::default(), keccak256([]).0);
+
     let state = RuntimeState {
         context: Context {
-            address: H160::from_slice(&[1u8; 20]),
+            address: H160::from(address.0 .0),
-            caller: H160::from_slice(&[2u8; 20]),
+            caller: H160::from(caller.0 .0),
             apparent_value: U256::default(),
         },
         transaction_context: TransactionContext {

crates/integration/codesize.json

@@ -1,9 +1,9 @@
 {
-  "Baseline": 3917,
+  "Baseline": 3944,
-  "Computation": 7363,
+  "Computation": 7444,
-  "DivisionArithmetics": 42649,
+  "DivisionArithmetics": 42784,
-  "ERC20": 53193,
+  "ERC20": 53524,
-  "FibonacciIterative": 5965,
+  "FibonacciIterative": 6019,
-  "Flipper": 4336,
+  "Flipper": 4392,
-  "SHA1": 36053
+  "SHA1": 36107
 }

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 {
+pub fn assert_success(contract: &Contract, differential: bool) -> (State, CallOutput) {
-    let (mut instance, export) = mock_runtime::prepare(&contract.pvm_runtime, None);
+    let (state, output) = contract.execute();
-    let state = mock_runtime::call(State::new(contract.calldata.clone()), &mut instance, export);
+    assert_eq!(output.flags, ReturnFlags::Success);
-    assert_eq!(state.output.flags, 0);
 
     if differential {
-        let evm = revive_differential::prepare(contract.evm_runtime, contract.calldata);
+        let evm =
-        assert_eq!(state.output.data.clone(), revive_differential::execute(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)]
+/// The mocked blockchain account.
-pub struct State {
+#[derive(Debug, Default, Clone)]
-    pub input: Vec<u8>,
+struct Account {
-    pub output: CallOutput,
+    value: U256,
-    pub value: u128,
+    contract: Option<U256>,
-    pub storage: HashMap<U256, U256>,
+    storage: HashMap<U256, U256>,
 }
 
-#[derive(Clone, Debug)]
+/// Emitted event data.
-pub struct CallOutput {
+#[derive(Debug, Default, Clone)]
-    pub flags: u32,
+pub struct Event {
     pub data: Vec<u8>,
+    pub topics: Vec<U256>,
 }
 
-impl Default for CallOutput {
+/// The result of the contract call.
-    fn default() -> Self {
+#[derive(Debug, Default, Clone)]
-        Self {
+pub struct CallOutput {
-            flags: u32::MAX,
+    /// The return flags.
-            data: Vec::new(),
+    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 {
+    pub fn transaction(self) -> TransactionBuilder {
-        Self {
+        TransactionBuilder {
-            input,
+            state_before: self.clone(),
-            ..Default::default()
+            context: Transaction {
+                state: self,
+                stack: vec![Default::default()],
+            },
         }
     }
 
-    pub fn reset_output(&mut self) {
+    pub fn upload_code(&mut self, code: &[u8]) -> U256 {
-        self.output = Default::default();
+        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) {
+    pub fn assert_storage_key(&self, account: Address, key: U256, expected: U256) {
-        assert_eq!(self.storage[&at], expect);
+        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> {
+            |caller: Caller<Transaction>, out_ptr: u32, out_len_ptr: u32| -> Result<(), Trap> {
-                let (mut caller, state) = caller.split();
+                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_ptr, input)?;
-                caller.write_memory(out_len_ptr, &(state.input.len() as u32).to_le_bytes())?;
+                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> {
+            |caller: Caller<Transaction>,
-                let (caller, state) = caller.split();
+             flags: u32,
+             data_ptr: u32,
+             data_len: u32|
+             -> Result<(), Trap> {
+                let (caller, transaction) = caller.split();
 
-                state.output.flags = flags;
+                let frame = transaction.top_frame_mut();
-                state.output.data = caller.read_memory_into_vec(data_ptr, data_len)?;
+                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> {
+            |caller: Caller<Transaction>, out_ptr: u32, out_len_ptr: u32| -> Result<(), Trap> {
-                let (mut caller, state) = caller.split();
+                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| {
+            |caller: Caller<Transaction>, out_ptr: u32, out_len_ptr: u32| {
-                let (mut caller, _) = caller.split();
+                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)?;
+                let address = transaction.top_frame().callee.as_slice();
-                caller.write_memory(out_len_ptr, &(State::ADDRESS.len() as u32).to_le_bytes())?;
+                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| {
+            |caller: Caller<Transaction>, out_ptr: u32, out_len_ptr: u32| {
-                let (mut caller, _) = caller.split();
+                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)?;
+                let address = transaction.top_frame().caller.as_slice();
-                caller.write_memory(out_len_ptr, &(State::CALLER.len() as u32).to_le_bytes())?;
+                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 (_, output) = assert_success(&contract, true);
-        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());
         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);
+    let (state, output) = contract.execute();
-    assert_eq!(state.output.flags, 0);
+    assert_eq!(output.flags, ReturnFlags::Success);
-    assert_eq!(state.storage[&U256::ZERO], U256::try_from(1).unwrap());
+    state.assert_storage_key(address, U256::ZERO, U256::from(1));
 
-    let state = crate::mock_runtime::call(state, &mut instance, export);
+    let (state, output) = state.transaction().calldata(contract.calldata).call();
-    assert_eq!(state.output.flags, 0);
+    assert_eq!(output.flags, ReturnFlags::Success);
-    assert_eq!(state.storage[&U256::ZERO], U256::ZERO);
+    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 (_, output) = State::default()
-    let (mut instance, export) = mock_runtime::prepare(&code, None);
+        .transaction()
-    let state = crate::mock_runtime::call(state, &mut instance, export);
+        .with_default_account(&code)
+        .calldata(input)
+        .call();
 
-    assert_eq!(state.output.flags, 0);
+    assert_eq!(output.flags, ReturnFlags::Success);
 
-    let mut hasher = Keccak256::new();
+    let expected = keccak256(param.as_bytes());
-    hasher.update(param);
+    let received = FixedBytes::<32>::from_slice(&output.data);
-    let expected = hasher.finalize();
-    let received = FixedBytes::<32>::from_slice(&state.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 (_, output) = assert_success(&Contract::triangle_number(13), true);
-    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());
     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 (_, output) = assert_success(&Contract::odd_product(5), true);
-    let received = I256::from_be_bytes::<32>(state.output.data.try_into().unwrap());
+    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 (_, output) = State::default()
-    let state = crate::mock_runtime::call(state, &mut instance, export);
+        .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 expected = I256::try_from(123).unwrap();
-    let received = I256::from_be_bytes::<32>(state.output.data.try_into().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 (_, output) = assert_success(&Contract::mstore8(*parameter), true);
-        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());
         (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 (_, output) = assert_success(&Contract::block_number(), true);
-    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());
     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 (_, output) = assert_success(&Contract::block_timestamp(), true);
-    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());
     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 contract = Contract::context_address();
-    let received = Address::from_slice(&state.output.data[12..]);
+    let (_, output) = assert_success(&contract, true);
-    let expected = Address::from(&mock_runtime::State::ADDRESS);
+    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 (_, output) = assert_success(&Contract::context_caller(), true);
-    let received = Address::from_slice(&state.output.data[12..]);
+    let received = Address::from_slice(&output.data[12..]);
-    let expected = Address::from(&mock_runtime::State::CALLER);
+    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 (_, output) = assert_success(&Contract::division_arithmetics_div(*n, *d), true);
-        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());
         (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 (_, output) = assert_success(&Contract::division_arithmetics_sdiv(*n, *d), true);
-        let received = I256::from_be_bytes::<32>(state.output.data.try_into().unwrap());
+        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 (_, output) = assert_success(&Contract::division_arithmetics_mod(*n, *d), true);
-        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());
         (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 (_, output) = assert_success(&Contract::division_arithmetics_smod(*n, *d), true);
-        let received = I256::from_be_bytes::<32>(state.output.data.try_into().unwrap());
+        let received = I256::from_be_bytes::<32>(output.data.try_into().unwrap());
         (received, *q)
     })
     .collect::<Vec<_>>()

crates/llvm-context/src/polkavm/const/mod.rs

@@ -46,7 +46,7 @@ pub static GLOBAL_VERBATIM_GETTER_PREFIX: &str = "get_global::";
 pub static GLOBAL_I256_SIZE: &str = "i256_size";
 
 /// The static value size.
-pub static GLOBAL_I128_SIZE: &str = "i128_size";
+pub static GLOBAL_I160_SIZE: &str = "i160_size";
 
 /// The static i64 size.
 pub static GLOBAL_I64_SIZE: &str = "i64_size";

crates/llvm-context/src/polkavm/context/function/intrinsics.rs

@@ -16,7 +16,9 @@ pub struct Intrinsics<'ctx> {
     /// The memory copy from a generic page.
     pub memory_copy_from_generic: FunctionDeclaration<'ctx>,
     /// Performs endianness swaps on i256 values
-    pub byte_swap: FunctionDeclaration<'ctx>,
+    pub byte_swap_word: FunctionDeclaration<'ctx>,
+    /// Performs endianness swaps on i160 values
+    pub byte_swap_eth_address: FunctionDeclaration<'ctx>,
 }
 
 impl<'ctx> Intrinsics<'ctx> {
@@ -30,7 +32,10 @@ impl<'ctx> Intrinsics<'ctx> {
     pub const FUNCTION_MEMORY_COPY_FROM_GENERIC: &'static str = "llvm.memcpy.p3.p1.i256";
 
     /// The corresponding intrinsic function name.
-    pub const FUNCTION_BYTE_SWAP: &'static str = "llvm.bswap.i256";
+    pub const FUNCTION_BYTE_SWAP_WORD: &'static str = "llvm.bswap.i256";
+
+    /// The corresponding intrinsic function name.
+    pub const FUNCTION_BYTE_SWAP_ETH_ADDRESS: &'static str = "llvm.bswap.i160";
 
     /// A shortcut constructor.
     pub fn new(
@@ -40,6 +45,7 @@ impl<'ctx> Intrinsics<'ctx> {
         let void_type = llvm.void_type();
         let bool_type = llvm.bool_type();
         let word_type = llvm.custom_width_int_type(revive_common::BIT_LENGTH_WORD as u32);
+        let address_type = llvm.custom_width_int_type(revive_common::BIT_LENGTH_ETH_ADDRESS as u32);
         let _stack_field_pointer_type = llvm.ptr_type(AddressSpace::Stack.into());
         let heap_field_pointer_type = llvm.ptr_type(AddressSpace::Heap.into());
         let generic_byte_pointer_type = llvm.ptr_type(AddressSpace::Generic.into());
@@ -78,18 +84,25 @@ impl<'ctx> Intrinsics<'ctx> {
                 false,
             ),
         );
-        let byte_swap = Self::declare(
+        let byte_swap_word = Self::declare(
             llvm,
             module,
-            Self::FUNCTION_BYTE_SWAP,
+            Self::FUNCTION_BYTE_SWAP_WORD,
             word_type.fn_type(&[word_type.as_basic_type_enum().into()], false),
         );
+        let byte_swap_eth_address = Self::declare(
+            llvm,
+            module,
+            Self::FUNCTION_BYTE_SWAP_ETH_ADDRESS,
+            address_type.fn_type(&[address_type.as_basic_type_enum().into()], false),
+        );
 
         Self {
             trap,
             memory_copy,
             memory_copy_from_generic,
-            byte_swap,
+            byte_swap_word,
+            byte_swap_eth_address,
         }
     }
 
@@ -131,7 +144,12 @@ impl<'ctx> Intrinsics<'ctx> {
                     .as_basic_type_enum(),
                 word_type.as_basic_type_enum(),
             ],
-            name if name == Self::FUNCTION_BYTE_SWAP => vec![word_type.as_basic_type_enum()],
+            name if name == Self::FUNCTION_BYTE_SWAP_WORD => vec![word_type.as_basic_type_enum()],
+            name if name == Self::FUNCTION_BYTE_SWAP_ETH_ADDRESS => {
+                vec![llvm
+                    .custom_width_int_type(revive_common::BIT_LENGTH_ETH_ADDRESS as u32)
+                    .as_basic_type_enum()]
+            }
             _ => vec![],
         }
     }

crates/llvm-context/src/polkavm/context/function/runtime/entry.rs

@@ -95,12 +95,12 @@ impl Entry {
         );
 
         context.set_global(
-            crate::polkavm::GLOBAL_I128_SIZE,
+            crate::polkavm::GLOBAL_I160_SIZE,
             context.xlen_type(),
             AddressSpace::Stack,
             context.integer_const(
                 crate::polkavm::XLEN,
-                revive_common::BYTE_LENGTH_X64 as u64 * 2,
+                revive_common::BYTE_LENGTH_X64 as u64 * 2 + revive_common::BYTE_LENGTH_X32 as u64,
             ),
         );
 

crates/llvm-context/src/polkavm/context/mod.rs

@@ -622,9 +622,9 @@ where
         let buffer_pointer = self.build_alloca(self.integer_type(bit_length), name);
         let symbol = match bit_length {
             revive_common::BIT_LENGTH_WORD => GLOBAL_I256_SIZE,
-            revive_common::BIT_LENGTH_VALUE => GLOBAL_I128_SIZE,
+            revive_common::BIT_LENGTH_ETH_ADDRESS => GLOBAL_I160_SIZE,
             revive_common::BIT_LENGTH_BLOCK_NUMBER => GLOBAL_I64_SIZE,
-            _ => unreachable!("invalid stack parameter bit width: {bit_length}"),
+            _ => panic!("invalid stack parameter bit width: {bit_length}"),
         };
         let length_pointer = self.get_global(symbol).expect("should be declared");
         (buffer_pointer, length_pointer.into())
@@ -837,13 +837,18 @@ where
         &self,
         value: inkwell::values::BasicValueEnum<'ctx>,
     ) -> anyhow::Result<inkwell::values::BasicValueEnum<'ctx>> {
+        let intrinsic = match value.get_type().into_int_type().get_bit_width() as usize {
+            revive_common::BIT_LENGTH_WORD => self.intrinsics().byte_swap_word.value,
+            revive_common::BIT_LENGTH_ETH_ADDRESS => self.intrinsics().byte_swap_eth_address.value,
+            _ => panic!(
+                "invalid byte swap parameter: {:?} {}",
+                value.get_name(),
+                value.get_type()
+            ),
+        };
         Ok(self
             .builder()
-            .build_call(
+            .build_call(intrinsic, &[value.into()], "call_byte_swap")?
-                self.intrinsics().byte_swap.value,
-                &[value.into()],
-                "call_byte_swap",
-            )?
             .try_as_basic_value()
             .left()
             .unwrap())
@@ -1177,7 +1182,9 @@ where
     /// Truncate a memory offset to register size, trapping if it doesn't fit.
     /// Pointers are represented as opaque 256 bit integer values in EVM.
     /// In practice, they should never exceed a register sized bit value.
-    /// However, we still protect against this possibility here.
+    /// However, we still protect against this possibility here. Heap index
+    /// offsets are generally untrusted and potentially represent valid
+    /// (but wrong) pointers when truncated.
     pub fn safe_truncate_int_to_xlen(
         &self,
         value: inkwell::values::IntValue<'ctx>,
@@ -1191,16 +1198,12 @@ where
             "expected XLEN or WORD sized int type for memory offset",
         );
 
-        let truncated = self.builder().build_int_truncate_or_bit_cast(
+        let truncated =
-            value,
+            self.builder()
-            self.xlen_type(),
+                .build_int_truncate(value, self.xlen_type(), "offset_truncated")?;
-            "offset_truncated",
+        let extended =
-        )?;
+            self.builder()
-        let extended = self.builder().build_int_z_extend_or_bit_cast(
+                .build_int_z_extend(truncated, self.word_type(), "offset_extended")?;
-            truncated,
-            self.word_type(),
-            "offset_extended",
-        )?;
         let is_overflow = self.builder().build_int_compare(
             inkwell::IntPredicate::NE,
             value,
@@ -1208,15 +1211,15 @@ where
             "compare_truncated_extended",
         )?;
 
-        let continue_block = self.append_basic_block("offset_pointer_ok");
+        let block_continue = self.append_basic_block("offset_pointer_ok");
-        let trap_block = self.append_basic_block("offset_pointer_overflow");
+        let block_trap = self.append_basic_block("offset_pointer_overflow");
-        self.build_conditional_branch(is_overflow, trap_block, continue_block)?;
+        self.build_conditional_branch(is_overflow, block_trap, block_continue)?;
 
-        self.set_basic_block(trap_block);
+        self.set_basic_block(block_trap);
         self.build_call(self.intrinsics().trap, &[], "invalid_trap");
         self.build_unreachable();
 
-        self.set_basic_block(continue_block);
+        self.set_basic_block(block_continue);
         Ok(truncated)
     }
 

crates/llvm-context/src/polkavm/evm/context.rs

@@ -172,7 +172,7 @@ where
     D: Dependency + Clone,
 {
     let (output_pointer, output_length_pointer) =
-        context.build_stack_parameter(revive_common::BIT_LENGTH_WORD, "address_output");
+        context.build_stack_parameter(revive_common::BIT_LENGTH_ETH_ADDRESS, "address_output");
     context.build_runtime_call(
         runtime_api::ADDRESS,
         &[
@@ -180,7 +180,11 @@ where
             output_length_pointer.to_int(context).into(),
         ],
     );
-    context.build_load(output_pointer, "address")
+    let value = context.build_byte_swap(context.build_load(output_pointer, "address")?)?;
+    Ok(context
+        .builder()
+        .build_int_z_extend(value.into_int_value(), context.word_type(), "address_zext")?
+        .into())
 }
 
 /// Translates the `caller` instruction.
@@ -191,7 +195,7 @@ where
     D: Dependency + Clone,
 {
     let (output_pointer, output_length_pointer) =
-        context.build_stack_parameter(revive_common::BIT_LENGTH_WORD, "caller_output");
+        context.build_stack_parameter(revive_common::BIT_LENGTH_ETH_ADDRESS, "caller_output");
     context.build_runtime_call(
         runtime_api::CALLER,
         &[
@@ -199,5 +203,9 @@ where
             output_length_pointer.to_int(context).into(),
         ],
     );
-    context.build_load(output_pointer, "caller")
+    let value = context.build_byte_swap(context.build_load(output_pointer, "caller")?)?;
+    Ok(context
+        .builder()
+        .build_int_z_extend(value.into_int_value(), context.word_type(), "caller_zext")?
+        .into())
 }