feat: add convert_boxed and insert_boxed for InstructionTable

crates/interpreter/src/instructions/opcode.rs

@@ -42,6 +42,24 @@ impl<H: Host> InstructionTables<'_, H> {
 }
 
 impl<'a, H: Host + 'a> InstructionTables<'a, H> {
+    /// Inserts a boxed instruction into the table with the specified index.
+    ///
+    /// Returns an error if the table is not boxed.
+    #[inline]
+    pub fn insert_boxed(
+        &mut self,
+        opcode: u8,
+        instruction: BoxedInstruction<'a, H>,
+    ) -> Result<(), InsertError> {
+        match self {
+            Self::Plain(_) => Err(InsertError::NotBoxed),
+            Self::Boxed(table) => {
+                table[opcode as usize] = Box::new(instruction);
+                Ok(())
+            }
+        }
+    }
+
     /// Inserts the instruction into the table with the specified index.
     #[inline]
     pub fn insert(&mut self, opcode: u8, instruction: Instruction<H>) {
@@ -54,6 +72,28 @@ impl<'a, H: Host + 'a> InstructionTables<'a, H> {
             }
         }
     }
+
+    /// Converts the instruction table to a boxed variant. If the table is already boxed, this does
+    /// nothing and returns the table unchanged.
+    #[inline]
+    pub fn to_boxed(self) -> InstructionTables<'a, H> {
+        let boxed_tables = match self {
+            Self::Plain(table) => core::array::from_fn(|i| {
+                let instruction: BoxedInstruction<'a, H> = Box::new(table[i]);
+                instruction
+            }),
+            Self::Boxed(table) => table,
+        };
+
+        Self::Boxed(boxed_tables)
+    }
+}
+
+/// An error that can occur when trying to insert instructions into the instruction table.
+#[derive(Debug)]
+pub enum InsertError {
+    /// Tried to insert a boxed instruction into an instruction table which was not boxed.
+    NotBoxed,
 }
 
 /// Make instruction table.

crates/revm/src/builder.rs

@@ -448,8 +448,63 @@ mod test {
         },
         Context, ContextPrecompile, ContextStatefulPrecompile, Evm, InMemoryDB, InnerEvmContext,
     };
-    use revm_interpreter::{Host, Interpreter};
-    use std::sync::Arc;
+    use revm_interpreter::{opcode::InstructionTables, Host, Interpreter};
+    use std::{cell::RefCell, rc::Rc, sync::Arc};
+
+    /// Custom evm context
+    #[derive(Default, Clone, Debug)]
+    pub(crate) struct CustomContext {
+        pub(crate) inner: Rc<RefCell<u8>>,
+    }
+
+    #[test]
+    fn simple_add_stateful_instruction() {
+        let code = Bytecode::new_raw([0xEF, 0x00].into());
+        let code_hash = code.hash_slow();
+        let to_addr = address!("ffffffffffffffffffffffffffffffffffffffff");
+
+        // initialize the custom context and make sure it's zero
+        let custom_context = CustomContext::default();
+        assert_eq!(*custom_context.inner.borrow(), 0);
+
+        let to_capture = custom_context.clone();
+        let mut evm = Evm::builder()
+            .with_db(InMemoryDB::default())
+            .modify_db(|db| {
+                db.insert_account_info(to_addr, AccountInfo::new(U256::ZERO, 0, code_hash, code))
+            })
+            .modify_tx_env(|tx| tx.transact_to = TransactTo::Call(to_addr))
+            // we need to use handle register box to capture the custom context in the handle
+            // register
+            .append_handler_register_box(Box::new(move |handler| {
+                let custom_context = to_capture.clone();
+
+                // we need to use a box to capture the custom context in the instruction
+                let custom_instruction = Box::new(
+                    move |_interp: &mut Interpreter, _host: &mut Evm<'_, (), InMemoryDB>| {
+                        // modify the value
+                        let mut inner = custom_context.inner.borrow_mut();
+                        *inner += 1;
+                    },
+                );
+
+                // need to make esure the instruction table is a boxed instruction table so that we
+                // can insert the custom instruction as a boxed instruction
+                let mut table = handler.take_instruction_table();
+                table = table.map(InstructionTables::to_boxed).map(|mut table| {
+                    // now we can finally insert
+                    table.insert_boxed(0xEF, custom_instruction).unwrap();
+                    table
+                });
+                handler.instruction_table = table;
+            }))
+            .build();
+
+        let _result_and_state = evm.transact().unwrap();
+
+        // ensure the custom context was modified
+        assert_eq!(*custom_context.inner.borrow(), 1);
+    }
 
     #[test]
     fn simple_add_instruction() {