From cc8682f83bfa73b10c6293e8c387c32dbba18e98 Mon Sep 17 00:00:00 2001
From: Dan Cline <6798349+Rjected@users.noreply.github.com>
Date: Sat, 16 Mar 2024 09:32:36 -0400
Subject: [PATCH] feat: add convert_boxed and insert_boxed for InstructionTable
 (#1194)

* feat: add into_box and insert_boxed for InstructionTable

Also adds an example for an instruction with captured context

* fix tests with clones

* fix valgrind

* fix valgrind again

* convert to boxed in insert_boxed

* update docs for insert_boxed

* make convert_boxed more concise

* Update crates/interpreter/src/instructions/opcode.rs

* Update crates/interpreter/src/instructions/opcode.rs
---
 crates/interpreter/src/instructions/opcode.rs | 35 ++++++++++++
 crates/revm/src/builder.rs                    | 57 ++++++++++++++++++-
 2 files changed, 91 insertions(+), 1 deletion(-)

diff --git a/crates/interpreter/src/instructions/opcode.rs b/crates/interpreter/src/instructions/opcode.rs
index a55ea1fd06..687637b4bc 100644
--- a/crates/interpreter/src/instructions/opcode.rs
+++ b/crates/interpreter/src/instructions/opcode.rs
@@ -42,6 +42,26 @@ impl<H: Host> InstructionTables<'_, H> {
 }
 
 impl<'a, H: Host + 'a> InstructionTables<'a, H> {
+    /// Inserts a boxed instruction into the table with the specified index.
+    ///
+    /// This will convert the table into the [BoxedInstructionTable] variant if it is currently a
+    /// plain instruction table, before inserting the instruction.
+    #[inline]
+    pub fn insert_boxed(&mut self, opcode: u8, instruction: BoxedInstruction<'a, H>) {
+        // first convert the table to boxed variant
+        self.convert_boxed();
+
+        // now we can insert the instruction
+        match self {
+            Self::Plain(_) => {
+                unreachable!("we already converted the table to boxed variant");
+            }
+            Self::Boxed(table) => {
+                table[opcode as usize] = Box::new(instruction);
+            }
+        }
+    }
+
     /// Inserts the instruction into the table with the specified index.
     #[inline]
     pub fn insert(&mut self, opcode: u8, instruction: Instruction<H>) {
@@ -54,6 +74,21 @@ impl<'a, H: Host + 'a> InstructionTables<'a, H> {
             }
         }
     }
+
+    /// Converts the current instruction table to a boxed variant. If the table is already boxed,
+    /// this is a no-op.
+    #[inline]
+    pub fn convert_boxed(&mut self) {
+        match self {
+            Self::Plain(table) => {
+                *self = Self::Boxed(core::array::from_fn(|i| {
+                    let instruction: BoxedInstruction<'a, H> = Box::new(table[i]);
+                    instruction
+                }));
+            }
+            Self::Boxed(_) => {}
+        };
+    }
 }
 
 /// Make instruction table.
diff --git a/crates/revm/src/builder.rs b/crates/revm/src/builder.rs
index e6afed749e..1b8d827983 100644
--- a/crates/revm/src/builder.rs
+++ b/crates/revm/src/builder.rs
@@ -449,7 +449,62 @@ mod test {
         Context, ContextPrecompile, ContextStatefulPrecompile, Evm, InMemoryDB, InnerEvmContext,
     };
     use revm_interpreter::{Host, Interpreter};
-    use std::sync::Arc;
+    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(|mut table| {
+                    // now we can finally insert
+                    table.insert_boxed(0xEF, custom_instruction);
+                    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() {