✨ Add step fn to MipsEVM

crates/mipsevm/Cargo.toml

@@ -14,13 +14,13 @@ serde = { version = "1.0.188", features = ["derive"] }
 serde_json = "1.0.106"
 fnv = "1.0.7"
 revm = { version = "3.3.0", features = ["no_gas_measuring"] }
+preimage-oracle = { path = "../preimage" }
 
 # misc
 anyhow = "1.0.75"
 tracing = "0.1.37"
 
 [dev-dependencies]
-preimage-oracle = { path = "../preimage" }
 rand = "0.8.5"
 criterion = { version = "0.5.1", features = ["html_reports"] }
 

crates/mipsevm/src/evm.rs

@@ -1,6 +1,7 @@
 //! This module contains a wrapper around a [revm] inspector with an in-memory backend
 //! that has the MIPS & PreimageOracle smart contracts deployed at deterministic addresses.
 
+use crate::{StateWitness, StateWitnessHasher, StepWitness};
 use alloy_primitives::{hex, Address, U256};
 use anyhow::Result;
 use revm::{
@@ -94,6 +95,89 @@ impl MipsEVM<CacheDB<EmptyDB>> {
         }
     }
 
+    /// Perform a single instruction step on the MIPS smart contract from the VM state encoded
+    /// in the [StepWitness] passed.
+    ///
+    /// ### Takes
+    /// - `witness`: The [StepWitness] containing the VM state to step.
+    ///
+    /// ### Returns
+    /// - A [Result] containing the post-state hash of the MIPS VM or an error returned during
+    /// execution.
+    pub fn step(&mut self, witness: StepWitness) -> Result<B256> {
+        if witness.has_preimage() {
+            #[cfg(feature = "tracing")]
+            tracing::info!(
+                target: "mipsevm::evm",
+                "Reading preimage key {:x} at offset {}",
+                witness.preimage_key,
+                witness.preimage_offset
+            );
+
+            let preimage_oracle_input =
+                witness
+                    .encode_preimage_oracle_input()
+                    .ok_or(anyhow::anyhow!(
+                        "Failed to ABI encode preimage oracle input."
+                    ))?;
+            self.fill_tx_env(
+                TransactTo::Call(PREIMAGE_ORACLE_ADDR.into()),
+                preimage_oracle_input.0,
+            );
+            self.inner.transact_commit().map_err(|_| {
+                anyhow::anyhow!("Failed to commit preimage to PreimageOracle contract")
+            })?;
+        }
+
+        #[cfg(feature = "tracing")]
+        tracing::debug!(
+            target: "mipsevm::evm",
+            "Performing EVM step",
+        );
+
+        let step_input = witness.encode_step_input();
+        self.fill_tx_env(TransactTo::Call(MIPS_ADDR.into()), step_input.0);
+        if let Ok(ResultAndState {
+            result:
+                revm::primitives::ExecutionResult::Success {
+                    reason: _,
+                    gas_used: _,
+                    gas_refunded: _,
+                    logs,
+                    output: Output::Call(output),
+                },
+            state: _,
+        }) = self.inner.transact_ref()
+        {
+            let output = B256::from_slice(&output);
+
+            #[cfg(feature = "tracing")]
+            tracing::debug!(
+                target: "mipsevm::evm",
+                "EVM step successful with resulting post-state hash: {:x}",
+                output,
+            );
+
+            if logs.len() != 1 {
+                anyhow::bail!("Expected 1 log, got {}", logs.len());
+            }
+
+            let post_state: StateWitness = logs[0].data.to_vec().as_slice().try_into()?;
+
+            if post_state.state_hash().as_slice() != output.as_slice() {
+                anyhow::bail!(
+                    "Post-state hash does not match state hash in log: {:x} != {:x}",
+                    output,
+                    post_state.state_hash()
+                );
+            }
+
+            Ok(output)
+        } else {
+            anyhow::bail!("Failed to step MIPS contract");
+        }
+    }
+
     /// Deploys a contract with the given code at the given address.
     ///
     /// ### Takes

crates/mipsevm/src/witness.rs

@@ -1,7 +1,9 @@
 //! This module contains the various witness types.
 
 use crate::{State, StateWitness, StateWitnessHasher};
-use alloy_primitives::{keccak256, B256};
+use alloy_primitives::{keccak256, Bytes, B256, U256};
+use alloy_sol_types::{sol, SolCall};
+use preimage_oracle::KeyType;
 
 /// The size of an encoded [StateWitness] in bytes.
 pub(crate) const STATE_WITNESS_SIZE: usize = 226;
@@ -44,3 +46,96 @@ impl Default for StepWitness {
         }
     }
 }
+
+sol! {
+    // `PreimageOracle` loadLocalData function.
+    function loadLocalData(uint256,bytes32,uint256,uint256) external returns (bytes32);
+
+    // `PreimageOracle` loadKeccak256PreimagePart function.
+    function loadKeccak256PreimagePart(uint256,bytes) external;
+
+    // `MIPS` step function.
+    function step(bytes,bytes) external returns (bytes32);
+}
+
+impl StepWitness {
+    /// Returns `true` if the step witness has a preimage.
+    pub fn has_preimage(&self) -> bool {
+        self.preimage_key != B256::ZERO
+    }
+
+    /// ABI encodes the input to the preimage oracle, if the [StepWitness] has a preimage request.
+    ///
+    /// ### Returns
+    /// - `Some(input)` if the [StepWitness] has a preimage request.
+    /// - `None` if the [StepWitness] does not have a preimage request.
+    pub fn encode_preimage_oracle_input(&self) -> Option<Bytes> {
+        if self.preimage_key == B256::ZERO {
+            #[cfg(feature = "tracing")]
+            tracing::warn!("Cannot encode preimage oracle input without preimage");
+            return None;
+        }
+
+        match KeyType::from(self.preimage_key[0]) {
+            KeyType::_Illegal => {
+                #[cfg(feature = "tracing")]
+                tracing::error!("Illegal key type");
+                None
+            }
+            KeyType::Local => {
+                if self.preimage_value.len() > 32 + 8 {
+                    #[cfg(feature = "tracing")]
+                    tracing::error!(
+                        target: "mipsevm::step_witness",
+                        "Local preimage value exceeds maximum size of 32 bytes with key 0x{:x}",
+                        self.preimage_key
+                    );
+                    return None;
+                }
+
+                let preimage_part = &self.preimage_value[8..];
+                let mut tmp = [0u8; 32];
+                tmp[0..preimage_part.len()].copy_from_slice(preimage_part);
+
+                let call = loadLocalDataCall {
+                    _0: self.preimage_key.into(),
+                    _1: tmp,
+                    _2: U256::from(self.preimage_value.len() - 8),
+                    _3: U256::from(self.preimage_offset),
+                };
+
+                Some(call.encode().into())
+            }
+            KeyType::GlobalKeccak => {
+                let call = loadKeccak256PreimagePartCall {
+                    _0: U256::from(self.preimage_offset),
+                    _1: self.preimage_value[8..].into(),
+                };
+
+                Some(call.encode().into())
+            }
+        }
+    }
+
+    /// ABI encodes the input to the MIPS step function.
+    ///
+    /// ### Returns
+    /// - The ABI encoded input to the MIPS step function.
+    pub fn encode_step_input(&self) -> Bytes {
+        let mut abi_state_len = self.state.len();
+        if abi_state_len % 32 != 0 {
+            abi_state_len += 32 - (abi_state_len % 32);
+        }
+
+        // Pad state to 32 byte multiple per ABI
+        let mut abi_state = vec![0u8; abi_state_len];
+        abi_state[..self.state.len()].copy_from_slice(&self.state);
+
+        let call = stepCall {
+            _0: abi_state,
+            _1: self.mem_proof.clone(),
+        };
+
+        call.encode().into()
+    }
+}

crates/preimage/src/types.rs

@@ -26,6 +26,16 @@ pub enum KeyType {
     GlobalKeccak = 2,
 }
 
+impl From<u8> for KeyType {
+    fn from(n: u8) -> Self {
+        match n {
+            1 => KeyType::Local,
+            2 => KeyType::GlobalKeccak,
+            _ => KeyType::_Illegal,
+        }
+    }
+}
+
 impl Key for LocalIndexKey {
     fn preimage_key(self) -> B256 {
         let mut out = B256::ZERO;