VM storage cache

arbitrator/arbutil/src/evm/api.rs

@@ -3,28 +3,24 @@
 
 use crate::{evm::user::UserOutcomeKind, Bytes20, Bytes32};
 use eyre::Result;
+use num_enum::IntoPrimitive;
 use std::sync::Arc;
 
-#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+#[derive(Clone, Copy, Debug, PartialEq, Eq, IntoPrimitive)]
 #[repr(u8)]
 pub enum EvmApiStatus {
     Success,
     Failure,
-}
-
-impl From<EvmApiStatus> for UserOutcomeKind {
-    fn from(value: EvmApiStatus) -> Self {
-        match value {
-            EvmApiStatus::Success => UserOutcomeKind::Success,
-            EvmApiStatus::Failure => UserOutcomeKind::Revert,
-        }
-    }
+    OutOfGas,
+    WriteProtection,
 }
 
 impl From<u8> for EvmApiStatus {
     fn from(value: u8) -> Self {
         match value {
             0 => Self::Success,
+            2 => Self::OutOfGas,
+            3 => Self::WriteProtection,
             _ => Self::Failure,
         }
     }
@@ -34,7 +30,7 @@ impl From<u8> for EvmApiStatus {
 #[repr(u32)]
 pub enum EvmApiMethod {
     GetBytes32,
-    SetBytes32,
+    SetTrieSlots,
     ContractCall,
     DelegateCall,
     StaticCall,
@@ -81,10 +77,13 @@ pub trait EvmApi<D: DataReader>: Send + 'static {
     /// Analogous to `vm.SLOAD`.
     fn get_bytes32(&mut self, key: Bytes32) -> (Bytes32, u64);
 
-    /// Stores the given value at the given key in the EVM state trie.
-    /// Returns the access cost on success.
-    /// Analogous to `vm.SSTORE`.
-    fn set_bytes32(&mut self, key: Bytes32, value: Bytes32) -> Result<u64>;
+    /// Stores the given value at the given key in Stylus VM's cache of the EVM state trie.
+    /// Note that the actual values only get written after calls to `set_trie_slots`.
+    fn cache_bytes32(&mut self, key: Bytes32, value: Bytes32) -> u64;
+
+    /// Persists any dirty values in the storage cache to the EVM state trie, dropping the cache entirely if requested.
+    /// Analogous to repeated invocations of `vm.SSTORE`.
+    fn flush_storage_cache(&mut self, clear: bool, gas_left: u64) -> Result<u64>;
 
     /// Calls the contract at the given address.
     /// Returns the EVM return data's length, the gas cost, and whether the call succeeded.
@@ -141,7 +140,7 @@ pub trait EvmApi<D: DataReader>: Send + 'static {
     ) -> (eyre::Result<Bytes20>, u32, u64);
 
     /// Returns the EVM return data.
-    /// Analogous to `vm.RETURNDATASIZE`.
+    /// Analogous to `vm.RETURNDATA`.
     fn get_return_data(&self) -> D;
 
     /// Emits an EVM log with the given number of topics and data, the first bytes of which should be the topic data.

arbitrator/arbutil/src/evm/mod.rs

@@ -1,10 +1,11 @@
-// Copyright 2023, Offchain Labs, Inc.
+// Copyright 2023-2024, Offchain Labs, Inc.
 // For license information, see https://github.com/OffchainLabs/nitro/blob/master/LICENSE
 
 use crate::{Bytes20, Bytes32};
 
 pub mod api;
 pub mod req;
+pub mod storage;
 pub mod user;
 
 // params.SstoreSentryGasEIP2200
@@ -13,9 +14,12 @@ pub const SSTORE_SENTRY_GAS: u64 = 2300;
 // params.ColdAccountAccessCostEIP2929
 pub const COLD_ACCOUNT_GAS: u64 = 2600;
 
-// params.ColdSloadCostEIP2929
+// params.WarmStorageReadCostEIP2929
 pub const COLD_SLOAD_GAS: u64 = 2100;
 
+// params.WarmSloadCostEIP2929;
+pub const WARM_SLOAD_GAS: u64 = 100;
+
 // params.LogGas and params.LogDataGas
 pub const LOG_TOPIC_GAS: u64 = 375;
 pub const LOG_DATA_GAS: u64 = 8;

arbitrator/arbutil/src/evm/req.rs

@@ -3,12 +3,16 @@
 
 use crate::{
     evm::{
-        api::{DataReader, EvmApi, EvmApiMethod},
+        api::{DataReader, EvmApi, EvmApiMethod, EvmApiStatus},
+        storage::{StorageCache, StorageWord},
         user::UserOutcomeKind,
     },
+    format::Utf8OrHex,
+    pricing::EVM_API_INK,
     Bytes20, Bytes32,
 };
 use eyre::{bail, eyre, Result};
+use std::collections::hash_map::Entry;
 
 pub trait RequestHandler<D: DataReader>: Send + 'static {
     fn handle_request(&mut self, req_type: EvmApiMethod, req_data: &[u8]) -> (Vec<u8>, D, u64);
@@ -18,6 +22,7 @@ pub struct EvmApiRequestor<D: DataReader, H: RequestHandler<D>> {
     handler: H,
     last_code: Option<(Bytes20, D)>,
     last_return_data: Option<D>,
+    storage_cache: StorageCache,
 }
 
 impl<D: DataReader, H: RequestHandler<D>> EvmApiRequestor<D, H> {
@@ -26,6 +31,7 @@ impl<D: DataReader, H: RequestHandler<D>> EvmApiRequestor<D, H> {
             handler,
             last_code: None,
             last_return_data: None,
+            storage_cache: StorageCache::default(),
         }
     }
 
@@ -93,20 +99,45 @@ impl<D: DataReader, H: RequestHandler<D>> EvmApiRequestor<D, H> {
 
 impl<D: DataReader, H: RequestHandler<D>> EvmApi<D> for EvmApiRequestor<D, H> {
     fn get_bytes32(&mut self, key: Bytes32) -> (Bytes32, u64) {
-        let (res, _, cost) = self.handle_request(EvmApiMethod::GetBytes32, key.as_slice());
-        (res.try_into().unwrap(), cost)
+        let cache = &mut self.storage_cache;
+        let mut cost = cache.read_gas();
+
+        let value = cache.entry(key).or_insert_with(|| {
+            let (res, _, gas) = self
+                .handler
+                .handle_request(EvmApiMethod::GetBytes32, key.as_slice());
+            cost = cost.saturating_add(gas).saturating_add(EVM_API_INK);
+            StorageWord::known(res.try_into().unwrap())
+        });
+        (value.value, cost)
     }
 
-    fn set_bytes32(&mut self, key: Bytes32, value: Bytes32) -> Result<u64> {
-        let mut request = Vec::with_capacity(64);
-        request.extend(key);
-        request.extend(value);
-        let (res, _, cost) = self.handle_request(EvmApiMethod::SetBytes32, &request);
-        if res.len() != 1 {
-            bail!("bad response from set_bytes32")
+    fn cache_bytes32(&mut self, key: Bytes32, value: Bytes32) -> u64 {
+        match self.storage_cache.entry(key) {
+            Entry::Occupied(mut key) => key.get_mut().value = value,
+            Entry::Vacant(slot) => drop(slot.insert(StorageWord::unknown(value))),
+        };
+        self.storage_cache.write_gas()
+    }
+
+    fn flush_storage_cache(&mut self, clear: bool, gas_left: u64) -> Result<u64> {
+        let mut data = Vec::with_capacity(64 * self.storage_cache.len() + 8);
+        data.extend(gas_left.to_be_bytes());
+
+        for (key, value) in &mut self.storage_cache.slots {
+            if value.dirty() {
+                data.extend(*key);
+                data.extend(*value.value);
+                value.known = Some(value.value);
+            }
         }
-        if res[0] != 1 {
-            bail!("write protected")
+        if clear {
+            self.storage_cache.clear();
+        }
+
+        let (res, _, cost) = self.handle_request(EvmApiMethod::SetTrieSlots, &data);
+        if res[0] != EvmApiStatus::Success.into() {
+            bail!("{}", String::from_utf8_or_hex(res));
         }
         Ok(cost)
     }
@@ -175,6 +206,7 @@ impl<D: DataReader, H: RequestHandler<D>> EvmApi<D> for EvmApiRequestor<D, H> {
     }
 
     fn emit_log(&mut self, data: Vec<u8>, topics: u32) -> Result<()> {
+        // TODO: remove copy
         let mut request = Vec::with_capacity(4 + data.len());
         request.extend(topics.to_be_bytes());
         request.extend(data);

arbitrator/arbutil/src/evm/storage.rs

@@ -0,0 +1,64 @@
+// Copyright 2022-2024, Offchain Labs, Inc.
+// For license information, see https://github.com/nitro/blob/master/LICENSE
+
+use crate::Bytes32;
+use std::{
+    collections::HashMap,
+    ops::{Deref, DerefMut},
+};
+
+/// Represents the EVM word at a given key.
+pub struct StorageWord {
+    /// The current value of the slot.
+    pub value: Bytes32,
+    /// The value in Geth, if known.
+    pub known: Option<Bytes32>,
+}
+
+impl StorageWord {
+    pub fn known(value: Bytes32) -> Self {
+        let known = Some(value);
+        Self { value, known }
+    }
+
+    pub fn unknown(value: Bytes32) -> Self {
+        Self { value, known: None }
+    }
+
+    pub fn dirty(&self) -> bool {
+        Some(self.value) != self.known
+    }
+}
+
+#[derive(Default)]
+pub struct StorageCache {
+    pub(crate) slots: HashMap<Bytes32, StorageWord>,
+}
+
+impl StorageCache {
+    pub const REQUIRED_ACCESS_GAS: u64 = crate::evm::COLD_SLOAD_GAS;
+
+    pub fn read_gas(&self) -> u64 {
+        //self.slots.len().ilog2() as u64
+        self.slots.len() as u64
+    }
+
+    pub fn write_gas(&self) -> u64 {
+        //self.slots.len().ilog2() as u64
+        self.slots.len() as u64
+    }
+}
+
+impl Deref for StorageCache {
+    type Target = HashMap<Bytes32, StorageWord>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.slots
+    }
+}
+
+impl DerefMut for StorageCache {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.slots
+    }
+}

arbitrator/arbutil/src/format.rs

@@ -50,3 +50,16 @@ impl<T: Debug> DebugBytes for T {
         format!("{:?}", self).as_bytes().to_vec()
     }
 }
+
+pub trait Utf8OrHex {
+    fn from_utf8_or_hex(data: impl Into<Vec<u8>>) -> String;
+}
+
+impl Utf8OrHex for String {
+    fn from_utf8_or_hex(data: impl Into<Vec<u8>>) -> String {
+        match String::from_utf8(data.into()) {
+            Ok(string) => string,
+            Err(error) => hex::encode(error.as_bytes()),
+        }
+    }
+}

arbitrator/stylus/cbindgen.toml

@@ -10,4 +10,4 @@ extra_bindings = ["arbutil", "prover"]
 prefix_with_name = true
 
 [export]
-include = ["EvmApiMethod"]
+include = ["EvmApiMethod", "EvmApiStatus"]

arbitrator/stylus/src/env.rs

@@ -1,4 +1,4 @@
-// Copyright 2022-2023, Offchain Labs, Inc.
+// Copyright 2022-2024, Offchain Labs, Inc.
 // For license information, see https://github.com/nitro/blob/master/LICENSE
 
 use arbutil::{

arbitrator/stylus/src/evm_api.rs

@@ -3,7 +3,7 @@
 
 use crate::{GoSliceData, RustSlice};
 use arbutil::evm::{
-    api::{EvmApiMethod, EvmApiStatus, EVM_API_METHOD_REQ_OFFSET},
+    api::{EvmApiMethod, EVM_API_METHOD_REQ_OFFSET},
     req::RequestHandler,
 };
 
@@ -16,7 +16,7 @@ pub struct NativeRequestHandler {
         gas_cost: *mut u64,
         result: *mut GoSliceData,
         raw_data: *mut GoSliceData,
-    ) -> EvmApiStatus,
+    ),
     pub id: usize,
 }
 
@@ -35,7 +35,7 @@ impl RequestHandler<GoSliceData> for NativeRequestHandler {
         let mut result = GoSliceData::null();
         let mut raw_data = GoSliceData::null();
         let mut cost = 0;
-        let status = unsafe {
+        unsafe {
             (self.handle_request_fptr)(
                 self.id,
                 req_type as u32 + EVM_API_METHOD_REQ_OFFSET,
@@ -45,7 +45,6 @@ impl RequestHandler<GoSliceData> for NativeRequestHandler {
                 ptr!(raw_data),
             )
         };
-        assert_eq!(status, EvmApiStatus::Success);
         (result.slice().to_vec(), raw_data, cost)
     }
 }

arbitrator/stylus/src/host.rs

@@ -126,12 +126,19 @@ pub(crate) fn storage_load_bytes32<D: DataReader, E: EvmApi<D>>(
     hostio!(env, storage_load_bytes32(key, dest))
 }
 
-pub(crate) fn storage_store_bytes32<D: DataReader, E: EvmApi<D>>(
+pub(crate) fn storage_cache_bytes32<D: DataReader, E: EvmApi<D>>(
     mut env: WasmEnvMut<D, E>,
     key: GuestPtr,
     value: GuestPtr,
 ) -> MaybeEscape {
-    hostio!(env, storage_store_bytes32(key, value))
+    hostio!(env, storage_cache_bytes32(key, value))
+}
+
+pub(crate) fn storage_flush_cache<D: DataReader, E: EvmApi<D>>(
+    mut env: WasmEnvMut<D, E>,
+    clear: u32,
+) -> MaybeEscape {
+    hostio!(env, storage_flush_cache(clear != 0))
 }
 
 pub(crate) fn call_contract<D: DataReader, E: EvmApi<D>>(

arbitrator/stylus/src/native.rs

@@ -130,7 +130,8 @@ impl<D: DataReader, E: EvmApi<D>> NativeInstance<D, E> {
                 "write_result" => func!(host::write_result),
                 "exit_early" => func!(host::exit_early),
                 "storage_load_bytes32" => func!(host::storage_load_bytes32),
-                "storage_store_bytes32" => func!(host::storage_store_bytes32),
+                "storage_cache_bytes32" => func!(host::storage_cache_bytes32),
+                "storage_flush_cache" => func!(host::storage_flush_cache),
                 "call_contract" => func!(host::call_contract),
                 "delegate_call_contract" => func!(host::delegate_call_contract),
                 "static_call_contract" => func!(host::static_call_contract),
@@ -339,7 +340,8 @@ pub fn module(wasm: &[u8], compile: CompileConfig) -> Result<Vec<u8>> {
             "write_result" => stub!(|_: u32, _: u32|),
             "exit_early" => stub!(|_: u32|),
             "storage_load_bytes32" => stub!(|_: u32, _: u32|),
-            "storage_store_bytes32" => stub!(|_: u32, _: u32|),
+            "storage_cache_bytes32" => stub!(|_: u32, _: u32|),
+            "storage_flush_cache" => stub!(|_: u32|),
             "call_contract" => stub!(u8 <- |_: u32, _: u32, _: u32, _: u32, _: u64, _: u32|),
             "delegate_call_contract" => stub!(u8 <- |_: u32, _: u32, _: u32, _: u64, _: u32|),
             "static_call_contract" => stub!(u8 <- |_: u32, _: u32, _: u32, _: u64, _: u32|),

arbitrator/stylus/src/test/api.rs

@@ -74,11 +74,17 @@ impl EvmApi<VecReader> for TestEvmApi {
         (value, 2100) // pretend worst case
     }
 
-    fn set_bytes32(&mut self, key: Bytes32, value: Bytes32) -> Result<u64> {
+    fn cache_bytes32(&mut self, key: Bytes32, value: Bytes32) -> u64 {
         let storage = &mut self.storage.lock();
         let storage = storage.get_mut(&self.program).unwrap();
         storage.insert(key, value);
-        Ok(22100) // pretend worst case
+        0
+    }
+
+    fn flush_storage_cache(&mut self, _clear: bool, _gas_left: u64) -> Result<u64> {
+        let storage = &mut self.storage.lock();
+        let storage = storage.get_mut(&self.program).unwrap();
+        Ok(22100 * storage.len() as u64) // pretend worst case
     }
 
     /// Simulates a contract call.