VM Storage Cache

.dockerignore

@@ -30,6 +30,7 @@ arbitrator/wasm-testsuite/target/
 arbitrator/wasm-libraries/target/
 arbitrator/tools/wasmer/target/
 arbitrator/tools/wasm-tools/
+arbitrator/tools/pricers/
 arbitrator/tools/module_roots/
 arbitrator/langs/rust/target/
 arbitrator/langs/bf/target/

arbitrator/Cargo.toml

@@ -1,6 +1,7 @@
 [workspace]
 members = [
         "arbutil",
+        "caller-env",
         "prover",
         "stylus",
         "jit",

arbitrator/arbutil/Cargo.toml

@@ -6,6 +6,7 @@ edition = "2021"
 [dependencies]
 digest = "0.9.0"
 eyre = "0.6.5"
+fnv = "1.0.7"
 hex = "0.4.3"
 num-traits = "0.2.17"
 siphasher = "0.3.10"

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
@@ -16,6 +17,9 @@ pub const COLD_ACCOUNT_GAS: u64 = 2600;
 // params.ColdSloadCostEIP2929
 pub const COLD_SLOAD_GAS: u64 = 2100;
 
+// params.WarmStorageReadCostEIP2929
+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,21 +3,26 @@
 
 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);
+    fn request(&mut self, req_type: EvmApiMethod, req_data: impl AsRef<[u8]>) -> (Vec<u8>, D, u64);
 }
 
 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,11 +31,12 @@ impl<D: DataReader, H: RequestHandler<D>> EvmApiRequestor<D, H> {
             handler,
             last_code: None,
             last_return_data: None,
+            storage_cache: StorageCache::default(),
         }
     }
 
-    fn handle_request(&mut self, req_type: EvmApiMethod, req_data: &[u8]) -> (Vec<u8>, D, u64) {
-        self.handler.handle_request(req_type, req_data)
+    fn request(&mut self, req_type: EvmApiMethod, req_data: impl AsRef<[u8]>) -> (Vec<u8>, D, u64) {
+        self.handler.request(req_type, req_data)
     }
 
     /// Call out to a contract.
@@ -48,7 +54,7 @@ impl<D: DataReader, H: RequestHandler<D>> EvmApiRequestor<D, H> {
         request.extend(gas.to_be_bytes());
         request.extend(input);
 
-        let (res, data, cost) = self.handle_request(call_type, &request);
+        let (res, data, cost) = self.request(call_type, &request);
         let status: UserOutcomeKind = res[0].try_into().expect("unknown outcome");
         let data_len = data.slice().len() as u32;
         self.last_return_data = Some(data);
@@ -75,7 +81,7 @@ impl<D: DataReader, H: RequestHandler<D>> EvmApiRequestor<D, H> {
         }
         request.extend(code);
 
-        let (mut res, data, cost) = self.handle_request(create_type, &request);
+        let (mut res, data, cost) = self.request(create_type, request);
         if res.len() != 21 || res[0] == 0 {
             if !res.is_empty() {
                 res.remove(0);
@@ -93,20 +99,47 @@ 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.request(EvmApiMethod::GetBytes32, key);
+            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 {
+        let cost = self.storage_cache.write_gas();
+        match self.storage_cache.entry(key) {
+            Entry::Occupied(mut key) => key.get_mut().value = value,
+            Entry::Vacant(slot) => drop(slot.insert(StorageWord::unknown(value))),
+        };
+        cost
+    }
+
+    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 clear {
+            self.storage_cache.clear();
         }
-        if res[0] != 1 {
-            bail!("write protected")
+        if data.len() == 8 {
+            return Ok(0); // no need to make request
+        }
+
+        let (res, _, cost) = self.request(EvmApiMethod::SetTrieSlots, data);
+        if res[0] != EvmApiStatus::Success.into() {
+            bail!("{}", String::from_utf8_or_hex(res));
         }
         Ok(cost)
     }
@@ -175,19 +208,20 @@ 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);
 
-        let (res, _, _) = self.handle_request(EvmApiMethod::EmitLog, &request);
+        let (res, _, _) = self.request(EvmApiMethod::EmitLog, request);
         if !res.is_empty() {
             bail!(String::from_utf8(res).unwrap_or("malformed emit-log response".into()))
         }
         Ok(())
     }
 
     fn account_balance(&mut self, address: Bytes20) -> (Bytes32, u64) {
-        let (res, _, cost) = self.handle_request(EvmApiMethod::AccountBalance, address.as_slice());
+        let (res, _, cost) = self.request(EvmApiMethod::AccountBalance, address);
         (res.try_into().unwrap(), cost)
     }
 
@@ -201,19 +235,18 @@ impl<D: DataReader, H: RequestHandler<D>> EvmApi<D> for EvmApiRequestor<D, H> {
         req.extend(address);
         req.extend(gas_left.to_be_bytes());
 
-        let (_, data, cost) = self.handle_request(EvmApiMethod::AccountCode, &req);
+        let (_, data, cost) = self.request(EvmApiMethod::AccountCode, req);
         self.last_code = Some((address, data.clone()));
         (data, cost)
     }
 
     fn account_codehash(&mut self, address: Bytes20) -> (Bytes32, u64) {
-        let (res, _, cost) = self.handle_request(EvmApiMethod::AccountCodeHash, address.as_slice());
+        let (res, _, cost) = self.request(EvmApiMethod::AccountCodeHash, address);
         (res.try_into().unwrap(), cost)
     }
 
     fn add_pages(&mut self, pages: u16) -> u64 {
-        self.handle_request(EvmApiMethod::AddPages, &pages.to_be_bytes())
-            .2
+        self.request(EvmApiMethod::AddPages, pages.to_be_bytes()).2
     }
 
     fn capture_hostio(
@@ -233,6 +266,6 @@ impl<D: DataReader, H: RequestHandler<D>> EvmApi<D> for EvmApiRequestor<D, H> {
         request.extend(name.as_bytes());
         request.extend(args);
         request.extend(outs);
-        self.handle_request(EvmApiMethod::CaptureHostIO, &request);
+        self.request(EvmApiMethod::CaptureHostIO, request);
     }
 }

arbitrator/arbutil/src/evm/storage.rs

@@ -0,0 +1,73 @@
+// Copyright 2022-2024, Offchain Labs, Inc.
+// For license information, see https://github.com/nitro/blob/master/LICENSE
+
+use crate::Bytes32;
+use fnv::FnvHashMap as HashMap;
+use std::ops::{Deref, DerefMut};
+
+/// Represents the EVM word at a given key.
+#[derive(Debug)]
+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>,
+    reads: usize,
+    writes: usize,
+}
+
+impl StorageCache {
+    pub const REQUIRED_ACCESS_GAS: u64 = 10;
+
+    pub fn read_gas(&mut self) -> u64 {
+        self.reads += 1;
+        match self.reads {
+            0..=32 => 0,
+            33..=128 => 2,
+            _ => 10,
+        }
+    }
+
+    pub fn write_gas(&mut self) -> u64 {
+        self.writes += 1;
+        match self.writes {
+            0..=8 => 0,
+            9..=64 => 7,
+            _ => 10,
+        }
+    }
+}
+
+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()),
+        }
+    }
+}