feat(arbos-revm): track stylus pages

Cargo.toml

@@ -405,8 +405,8 @@ idna_adapter = "=1.1.0"
 # alloy-transport-ws = { git = "https://github.com/alloy-rs/alloy", rev = "7fab7ee" }
 
 ## alloy-evm
-# alloy-evm = { git = "https://github.com/alloy-rs/evm.git", rev = "7762adc" }
-# alloy-op-evm = { git = "https://github.com/alloy-rs/evm.git", rev = "7762adc" }
+alloy-evm = { git = "https://github.com/bernard-wagner/alloy-evm.git", rev = "213ad42" }
+alloy-op-evm = { git = "https://github.com/bernard-wagner/alloy-evm.git", rev = "213ad42" }
 
 ## revm
 revm = { git = "https://github.com/bluealloy/revm.git", rev = "409c2b3" }

crates/arbos-revm/src/context.rs

@@ -3,15 +3,16 @@ use revm::{
     context::{BlockEnv, CfgEnv, ContextTr, TxEnv},
 };
 
-use crate::chain::{ArbitrumChainInfo, ArbitrumChainInfoTr};
+use crate::{chain::{ArbitrumChainInfo, ArbitrumChainInfoTr}, local_context::{ArbitrumLocalContext, ArbitrumLocalContextTr}};
 
 /// Type alias for the default context type of the ArbitrumEvm.
-pub type ArbitrumContext<DB> = Context<BlockEnv, TxEnv, CfgEnv, DB, Journal<DB>, ArbitrumChainInfo>;
+pub type ArbitrumContext<DB> = Context<BlockEnv, TxEnv, CfgEnv, DB, Journal<DB>, ArbitrumChainInfo, ArbitrumLocalContext>;
 
 /// Type alias for Arbitrum context
 pub trait ArbitrumContextTr: ContextTr<
     // Cfg: Cfg<Spec = ArbitrumSpecId>,
     Chain: ArbitrumChainInfoTr,
+    Local: ArbitrumLocalContextTr,
 >
 {
 }
@@ -20,6 +21,7 @@ impl<T> ArbitrumContextTr for T where
     T: ContextTr<
         //Cfg: Cfg<Spec = ArbitrumSpecId>,
         Chain: ArbitrumChainInfoTr,
+        Local: ArbitrumLocalContextTr,
     >
 {
 }

crates/arbos-revm/src/lib.rs

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

crates/arbos-revm/src/local_context.rs

@@ -0,0 +1,74 @@
+use std::{cell::RefCell, rc::Rc};
+
+use revm::context::LocalContextTr;
+
+pub trait ArbitrumLocalContextTr: LocalContextTr {
+    fn stylus_pages_ever(&self) -> u64;
+    fn stylus_pages_open(&self) -> u64;
+    fn add_stylus_pages_open(&mut self, pages: u64);
+    fn set_stylus_pages_open(&mut self, pages: u64);
+}
+
+/// 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: u64,
+    /// Stylus pages currently open.
+    pub stylus_pages_open: u64,
+}
+
+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) -> u64 {
+        self.stylus_pages_ever
+    }
+
+    fn stylus_pages_open(&self) -> u64 {
+        self.stylus_pages_open
+    }
+
+    fn add_stylus_pages_open(&mut self, pages: u64) {
+        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: u64) {
+        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()
+    }
+}