feat(arbos-revm): track stylus pages

Cargo.toml

@@ -277,7 +277,7 @@ stylus = { git = "https://github.com/bernard-wagner/nitro", rev = "e1d47755c5fe4
 wasmer-types = { git = "https://github.com/bernard-wagner/nitro", rev = "e1d47755c5fe476cf197c9304faf2ec61ea0e850", default-features = false }
 
 ## alloy-evm
-alloy-evm = { git = "https://github.com/iosiro/alloy-evm", rev = "c18cc373d3de47b2dfcd57d0149972aecaa1ccfc", default-features = false }
+alloy-evm = { git = "https://github.com/iosiro/alloy-evm", rev = "60dc4ec58100af8938ad090e7231bc064380dd34", default-features = false }
 
 ## cli
 anstream = "0.6"

crates/arbos-revm/src/constants.rs

@@ -6,7 +6,7 @@ pub const STYLUS_DISCRIMINANT: &[u8] =
     &[STYLUS_EOF_MAGIC, STYLUS_EOF_MAGIC_SUFFIX, STYLUS_EOF_VERSION];
 
 pub const INITIAL_ARBOS_VERSION: u16 = 42;
-pub const INITIAL_STYLUS_VERSION: u16 = 1;
+pub const INITIAL_STYLUS_VERSION: u16 = 2;
 pub const INITIAL_MAX_WASM_SIZE: u32 = 128 * 1024; // max decompressed wasm size (programs are also bounded by compressed size)
 pub const INITIAL_MAX_STACK_DEPTH: u32 = 4 * 65536; // 4 page stack.
 pub const INITIAL_FREE_PAGES: u16 = 2; // 2 pages come free

crates/arbos-revm/src/context.rs

@@ -3,10 +3,12 @@ use revm::{
     context::{BlockEnv, ContextTr, TxEnv},
 };
 
-use crate::config::{ArbitrumConfig, ArbitrumConfigTr};
+use crate::{
+    config::{ArbitrumConfig, ArbitrumConfigTr},
+    local_context::{ArbitrumLocalContext, ArbitrumLocalContextTr},
+};
 
 pub type ArbitrumChainInfo = ();
-pub type ArbitrumLocalContext = revm::context::LocalContext;
 
 /// Type alias for the default context type of the ArbitrumEvm.
 pub type ArbitrumContext<DB> = Context<
@@ -20,6 +22,12 @@ pub type ArbitrumContext<DB> = Context<
 >;
 
 /// Type alias for Arbitrum context
-pub trait ArbitrumContextTr: ContextTr<Cfg: ArbitrumConfigTr> {}
+pub trait ArbitrumContextTr:
+    ContextTr<Cfg: ArbitrumConfigTr, Local: ArbitrumLocalContextTr>
+{
+}
 
-impl<T> ArbitrumContextTr for T where T: ContextTr<Cfg: ArbitrumConfigTr> {}
+impl<T> ArbitrumContextTr for T where
+    T: ContextTr<Cfg: ArbitrumConfigTr, Local: ArbitrumLocalContextTr>
+{
+}

crates/arbos-revm/src/lib.rs

@@ -11,6 +11,7 @@ pub mod context;
 pub mod evm;
 pub mod handler;
 pub mod inspector;
+pub mod local_context;
 //pub mod precompiles;
 pub mod result;
 //pub mod spec;

crates/arbos-revm/src/local_context.rs

@@ -0,0 +1,73 @@
+use std::{cell::RefCell, rc::Rc};
+
+use revm::context::LocalContextTr;
+
+pub trait ArbitrumLocalContextTr: LocalContextTr {
+    fn stylus_pages_ever(&self) -> u16;
+    fn stylus_pages_open(&self) -> u16;
+    fn add_stylus_pages_open(&mut self, pages: u16);
+    fn set_stylus_pages_open(&mut self, pages: u16);
+}
+
+/// Local context that is filled by execution.
+#[derive(Clone, Debug)]
+pub struct ArbitrumLocalContext {
+    /// Interpreter shared memory buffer. A reused memory buffer for calls.
+    pub shared_memory_buffer: Rc<RefCell<Vec<u8>>>,
+    /// Stylus pages ever used in this transaction.
+    pub stylus_pages_ever: u16,
+    /// Stylus pages currently open.
+    pub stylus_pages_open: u16,
+}
+
+impl Default for ArbitrumLocalContext {
+    fn default() -> Self {
+        Self {
+            shared_memory_buffer: Rc::new(RefCell::new(Vec::with_capacity(1024 * 4))),
+            stylus_pages_ever: 0,
+            stylus_pages_open: 0,
+        }
+    }
+}
+
+impl LocalContextTr for ArbitrumLocalContext {
+    fn clear(&mut self) {
+        // Sets len to 0 but it will not shrink to drop the capacity.
+        unsafe { self.shared_memory_buffer.borrow_mut().set_len(0) };
+    }
+
+    fn shared_memory_buffer(&self) -> &Rc<RefCell<Vec<u8>>> {
+        &self.shared_memory_buffer
+    }
+}
+
+impl ArbitrumLocalContextTr for ArbitrumLocalContext {
+    fn stylus_pages_ever(&self) -> u16 {
+        self.stylus_pages_ever
+    }
+
+    fn stylus_pages_open(&self) -> u16 {
+        self.stylus_pages_open
+    }
+
+    fn add_stylus_pages_open(&mut self, pages: u16) {
+        self.stylus_pages_open = self.stylus_pages_open.saturating_add(pages);
+        if self.stylus_pages_open > self.stylus_pages_ever {
+            self.stylus_pages_ever = self.stylus_pages_open;
+        }
+    }
+
+    fn set_stylus_pages_open(&mut self, pages: u16) {
+        self.stylus_pages_open = pages;
+        if self.stylus_pages_open > self.stylus_pages_ever {
+            self.stylus_pages_ever = self.stylus_pages_open;
+        }
+    }
+}
+
+impl ArbitrumLocalContext {
+    /// Creates a new local context, initcodes are hashes and added to the mapping.
+    pub fn new() -> Self {
+        Self::default()
+    }
+}

crates/arbos-revm/src/stylus_api.rs

@@ -112,7 +112,10 @@ where
         };
 
         let mut gas = Gas::new(gas_limit);
-        _ = gas.record_cost(100);
+        if !gas.record_cost(100) {
+            gas.spend_all();
+            return (Status::OutOfGas.into(), VecReader::new(vec![]), ArbGas(gas.spent()));
+        }
 
         let first_frame_input = FrameInput::Call(Box::new(CallInputs {
             input: CallInput::Bytes(calldata),
@@ -232,7 +235,7 @@ where
             return (
                 [vec![0x00], "out of gas".as_bytes().to_vec()].concat(),
                 VecReader::new(vec![]),
-                ArbGas(gas_remaining),
+                ArbGas(0),
             );
         }
 

crates/arbos-revm/src/stylus_executor.rs

@@ -45,11 +45,11 @@ use crate::{
     ArbitrumEvm,
     config::{ArbitrumConfigTr, ArbitrumStylusConfigTr},
     constants::{
-        COST_SCALAR_PERCENT, INITIAL_CACHED_COST_SCALAR, INITIAL_FREE_PAGES,
-        INITIAL_INIT_COST_SCALAR, INITIAL_MIN_CACHED_GAS, INITIAL_MIN_INIT_GAS, INITIAL_PAGE_GAS,
-        MEMORY_EXPONENTS, MIN_CACHED_GAS_UNITS, MIN_INIT_GAS_UNITS, STYLUS_DISCRIMINANT,
+        COST_SCALAR_PERCENT, MEMORY_EXPONENTS, MIN_CACHED_GAS_UNITS, MIN_INIT_GAS_UNITS,
+        STYLUS_DISCRIMINANT,
     },
     context::ArbitrumContextTr,
+    local_context::ArbitrumLocalContextTr,
     stylus_api::StylusHandler,
 };
 
@@ -62,7 +62,7 @@ lazy_static::lazy_static! {
 type EvmApiHandler<'a> =
     Arc<Box<dyn Fn(EvmApiMethod, Vec<u8>) -> (Vec<u8>, VecReader, arbutil::evm::api::Gas) + 'a>>;
 
-pub fn build_evm_data<CTX>(context: &mut CTX, input: InputsImpl) -> EvmData
+pub fn build_evm_data<CTX>(context: &CTX, input: InputsImpl) -> EvmData
 where
     CTX: ArbitrumContextTr,
 {
@@ -111,16 +111,16 @@ pub(crate) struct StylusExecutionContext {
     calldata: Bytes,
 }
 
-pub fn stylus_call_cost(new: u16, open: u16, ever: u16) -> u64 {
+pub fn stylus_call_cost(new: u16, open: u16, ever: u16, free_pages: u16, page_gas: u16) -> u64 {
     let new_open = open.saturating_add(new);
     let new_ever = max(ever, new_open);
 
-    if new_ever < INITIAL_FREE_PAGES {
+    if new_ever < free_pages {
         return 0;
     }
 
-    let adding = new_open.saturating_sub(open).saturating_sub(INITIAL_FREE_PAGES);
-    let linear = (adding as u64).saturating_mul(INITIAL_PAGE_GAS as u64);
+    let adding = new_open.saturating_sub(open).saturating_sub(free_pages);
+    let linear = (adding as u64).saturating_mul(page_gas as u64);
     let exp = |x: u16| -> u64 {
         if x < MEMORY_EXPONENTS.len() as u16 {
             return MEMORY_EXPONENTS[x as usize] as u64;
@@ -134,17 +134,20 @@ pub fn stylus_call_cost(new: u16, open: u16, ever: u16) -> u64 {
     linear.saturating_add(expand)
 }
 
-pub fn init_gas(params: StylusData) -> u64 {
-    let base = INITIAL_MIN_INIT_GAS as u64 * MIN_INIT_GAS_UNITS;
-    let dyno = (params.init_cost as u64)
-        .saturating_mul(INITIAL_INIT_COST_SCALAR as u64 * COST_SCALAR_PERCENT);
+pub fn init_gas_cost(init_cost: u16, min_init_gas: u8, init_cost_scaler: u8) -> u64 {
+    let base = min_init_gas as u64 * MIN_INIT_GAS_UNITS;
+    let dyno = (init_cost as u64).saturating_mul(init_cost_scaler as u64 * COST_SCALAR_PERCENT);
     base.saturating_add(dyno.div_ceil(100))
 }
 
-pub fn cached_gas(params: StylusData) -> u64 {
-    let base = INITIAL_MIN_CACHED_GAS as u64 * MIN_CACHED_GAS_UNITS;
-    let dyno = (params.cached_init_cost as u64)
-        .saturating_mul(INITIAL_CACHED_COST_SCALAR as u64 * COST_SCALAR_PERCENT);
+pub fn cached_gas_cost(
+    cached_init_cost: u16,
+    min_cached_init_gas: u8,
+    cached_init_cost_scaler: u8,
+) -> u64 {
+    let base = min_cached_init_gas as u64 * MIN_CACHED_GAS_UNITS;
+    let dyno = (cached_init_cost as u64)
+        .saturating_mul(cached_init_cost_scaler as u64 * COST_SCALAR_PERCENT);
     base.saturating_add(dyno.div_ceil(100))
 }
 
@@ -303,18 +306,33 @@ where
             Vec<u8>,
         ) -> (Vec<u8>, VecReader, ArbGas),
     ) -> Option<InterpreterAction> {
-        let context = self.ctx();
-
-        let compile_config = CompileConfig::version(
-            context.cfg().stylus().stylus_version(),
-            context.cfg().stylus().debug_mode(),
-        );
-
-        let stylus_config = StylusConfig::new(
-            context.cfg().stylus().stylus_version(),
-            context.cfg().stylus().max_stack_depth(),
-            context.cfg().stylus().ink_price(),
-        );
+        let (stylus_config, compile_config, evm_data) = {
+            let context = self.ctx();
+
+            let config = context.cfg().stylus();
+
+            let stylus_config = StylusConfig::new(
+                config.stylus_version(),
+                config.max_stack_depth(),
+                config.ink_price(),
+            );
+
+            let compile_config =
+                CompileConfig::version(config.stylus_version(), config.debug_mode());
+
+            let evm_data = build_evm_data(
+                self.ctx(),
+                InputsImpl {
+                    target_address: stylus_ctx.target_address,
+                    caller_address: stylus_ctx.caller_address,
+                    input: CallInput::Bytes(stylus_ctx.calldata.clone()),
+                    call_value: stylus_ctx.call_value,
+                    bytecode_address: Some(stylus_ctx.target_address),
+                },
+            );
+
+            (stylus_config, compile_config, evm_data)
+        };
 
         let (serialized, _module, stylus_data) = {
             // Use read lock to get cached program if available
@@ -331,6 +349,8 @@ where
             if let Some((serialized, module, stylus_data)) = maybe_cached {
                 (serialized, module, stylus_data)
             } else {
+                let context = self.ctx();
+
                 let bytecode = context.journal_mut().code(stylus_ctx.bytecode_address).ok()?.data;
 
                 if !bytecode.starts_with(STYLUS_DISCRIMINANT) {
@@ -365,8 +385,27 @@ where
             bytecode_address: Some(stylus_ctx.target_address),
         };
 
-        let call_cost = stylus_call_cost(stylus_data.footprint, 0, INITIAL_FREE_PAGES)
-            + cached_gas(stylus_data);
+        let (call_cost, stylus_open_pages) = {
+            let context = self.ctx();
+            let wasm_open_pages = context.local().stylus_pages_open();
+
+            let page_grow_cost = stylus_call_cost(
+                stylus_data.footprint,
+                wasm_open_pages,
+                context.local().stylus_pages_ever(),
+                context.cfg().stylus().free_pages(),
+                context.cfg().stylus().page_gas(),
+            );
+
+            let program_cost = cached_gas_cost(
+                stylus_data.cached_init_cost,
+                context.cfg().stylus().min_cached_init_gas(),
+                context.cfg().stylus().cached_cost_scalar(),
+            );
+
+            let cost = program_cost.saturating_add(page_grow_cost);
+            (cost, wasm_open_pages)
+        };
 
         let mut gas = Gas::new(stylus_ctx.gas_limit);
         if !gas.record_cost(call_cost) {
@@ -377,7 +416,10 @@ where
             }));
         }
 
-        let evm_data = build_evm_data(self.ctx(), inputs.clone());
+        {
+            self.ctx().local_mut().add_stylus_pages_open(stylus_data.footprint);
+        }
+
         let evm_api =
             self.build_api_requestor(inputs.clone(), stylus_ctx.is_static, api_request_handler);
 
@@ -416,6 +458,8 @@ where
 
         gas.erase_cost(gas_left);
 
+        self.ctx().local_mut().set_stylus_pages_open(stylus_open_pages);
+
         Some(InterpreterAction::Return(InterpreterResult { result, output: data.into(), gas }))
     }
 

crates/cheatcodes/src/inspector.rs

@@ -40,7 +40,7 @@ use foundry_evm_core::{
     abi::Vm::stopExpectSafeMemoryCall,
     backend::{DatabaseError, DatabaseExt, RevertDiagnostic},
     constants::{CHEATCODE_ADDRESS, HARDHAT_CONSOLE_ADDRESS, MAGIC_ASSUME},
-    evm::{BlockEnv, EthEvmContext, FoundryEvm, new_evm_with_existing_context},
+    evm::{BlockEnv, EthEvmContext, FoundryEvm, LocalContext, new_evm_with_existing_context},
 };
 use foundry_evm_traces::{
     TracingInspector, TracingInspectorConfig, identifier::SignaturesIdentifier,
@@ -52,7 +52,7 @@ use rand::Rng;
 use revm::{
     Inspector, Journal,
     bytecode::opcode as op,
-    context::{JournalTr, LocalContext, TransactionType, result::EVMError},
+    context::{JournalTr, TransactionType, result::EVMError},
     context_interface::{CreateScheme, transaction::SignedAuthorization},
     handler::FrameResult,
     interpreter::{

crates/evm/core/src/env.rs

@@ -1,6 +1,6 @@
 use revm::{
     Context, Database, Journal, JournalEntry,
-    context::{JournalInner, JournalTr},
+    context::{JournalInner, JournalTr, LocalContextTr},
     primitives::hardfork::SpecId,
 };
 
@@ -71,8 +71,8 @@ impl AsEnvMut for Env {
     }
 }
 
-impl<DB: Database, J: JournalTr<Database = DB>, C> AsEnvMut
-    for Context<BlockEnv, TxEnv, CfgEnv, DB, J, C>
+impl<DB: Database, J: JournalTr<Database = DB>, C, L: LocalContextTr> AsEnvMut
+    for Context<BlockEnv, TxEnv, CfgEnv, DB, J, C, L>
 {
     fn as_env_mut(&mut self) -> EnvMut<'_> {
         EnvMut { block: &mut self.block, cfg: &mut self.cfg, tx: &mut self.tx }
@@ -87,8 +87,8 @@ pub trait ContextExt {
     ) -> (&mut Self::DB, &mut JournalInner<JournalEntry>, EnvMut<'_>);
 }
 
-impl<DB: Database, C> ContextExt
-    for Context<BlockEnv, TxEnv, CfgEnv, DB, Journal<DB, JournalEntry>, C>
+impl<DB: Database, C, L: LocalContextTr> ContextExt
+    for Context<BlockEnv, TxEnv, CfgEnv, DB, Journal<DB, JournalEntry>, C, L>
 {
     type DB = DB;
 

crates/evm/core/src/evm.rs

@@ -35,7 +35,7 @@ use arbos_revm::{ArbitrumContext, ArbitrumEvm as RevmEvm};
 pub type BlockEnv = revm::context::BlockEnv;
 pub type CfgEnv<SPEC = SpecId> = arbos_revm::config::ArbitrumConfig<SPEC>;
 pub type TxEnv = revm::context::TxEnv;
-pub type LocalContext = revm::context::LocalContext;
+pub type LocalContext = arbos_revm::local_context::ArbitrumLocalContext;
 
 pub type EthEvmContext<DB> = ArbitrumContext<DB>;
 pub type EvmEnv<SPEC = SpecId> = alloy_evm::EvmEnv<BlockEnv, CfgEnv<SPEC>>;