Add Optional JIT compilation

Cargo.toml

@@ -208,3 +208,6 @@ futures = "0.3"
 
 # misc-testing
 rstest = "0.18.2"
+
+[patch.crates-io]
+revmc = { path = "../revmc/crates/revmc" }

crates/node/Cargo.toml

@@ -31,6 +31,11 @@ alloy-primitives.workspace = true
 serde_json.workspace = true
 tracing.workspace = true
 eyre.workspace = true
+revmc = "0.1.0"
+
 
 [lints]
 workspace = true
+
+[build-dependencies]
+revmc-build = "0.1.0"

crates/node/build.rs

@@ -0,0 +1,5 @@
+#![allow(missing_docs)]
+
+fn main() {
+    revmc_build::emit();
+}

crates/node/src/compiler.rs

@@ -0,0 +1,162 @@
+/*! The compiler module is responsible for compiling EVM bytecode to machine code using LLVM. */
+
+use alloy_primitives::{hex, Bytes, B256};
+use core::panic;
+use reth_revm::{
+    handler::register::EvmHandler,
+    interpreter::{InterpreterAction, SharedMemory},
+    Context as RevmContext, Database, Frame,
+};
+use revmc::{
+    llvm::Context as LlvmContext, primitives::SpecId, EvmCompiler, EvmCompilerFn, EvmLlvmBackend,
+    OptimizationLevel,
+};
+use std::{
+    collections::HashMap,
+    sync::{mpsc::Sender, Arc, Mutex},
+    thread,
+};
+
+/// The [Compiler] struct is a client for passing functions to the compiler thread. It also contains a cache of compiled functions
+#[derive(Debug, Clone)]
+pub struct Compiler {
+    sender: Sender<(SpecId, B256, Bytes)>,
+    fn_cache: Arc<Mutex<HashMap<B256, Option<EvmCompilerFn>>>>,
+}
+
+// TODO: probably shouldn't have a default for something that spawns a thread?
+impl Default for Compiler {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl Compiler {
+    /// Create a new compiler instance. This spawns a new compiler thread and the returned struct contains a [Sender](std::sync::mpsc::Sender) for sending functions to the compiler thread,
+    /// as well as a cache to compiled functions
+    pub fn new() -> Self {
+        let fn_cache = Arc::new(Mutex::new(HashMap::new()));
+        let (sender, receiver) = std::sync::mpsc::channel();
+
+        // TODO: graceful shutdown
+        thread::spawn({
+            let fn_cache = fn_cache.clone();
+
+            move || {
+                let ctx = LlvmContext::create();
+                // let mut compilers = Vec::new();
+
+                while let Ok((spec_id, hash, code)) = receiver.recv() {
+                    fn_cache.lock().unwrap().insert(hash, None);
+
+                    // TODO: fail properly here.
+                    let backend =
+                        EvmLlvmBackend::new(&ctx, false, OptimizationLevel::Aggressive).unwrap();
+                    let compiler = Box::leak(Box::new(EvmCompiler::new(backend)));
+
+                    // Do we have to allocate here? Not sure there's a better option
+                    let name = hex::encode(hash);
+                    dbg!("compiled", &name);
+
+                    let result =
+                        unsafe { compiler.jit(&name, &code, spec_id) }.expect("catastrophe");
+
+                    fn_cache.lock().unwrap().insert(hash, Some(result));
+
+                    // compilers.push(compiler);
+                }
+            }
+        });
+
+        Self { sender, fn_cache }
+    }
+
+    // TODO:
+    // For safety, we should also borrow the EvmCompiler that holds the actual module with code to
+    // make sure that it's not dropped while before or during the function call.
+    fn get_compiled_fn(&self, spec_id: SpecId, hash: B256, code: Bytes) -> Option<EvmCompilerFn> {
+        match self.fn_cache.lock().unwrap().get(&hash) {
+            Some(maybe_f) => *maybe_f,
+            None => {
+                // TODO: put rules here for whether or not to compile the function
+                self.sender.send((spec_id, hash, code)).unwrap();
+                None
+            }
+        }
+    }
+}
+
+/// The [ExternalContext] struct is a container for the [Compiler] struct.
+#[derive(Debug)]
+pub struct ExternalContext {
+    compiler: Compiler,
+}
+
+impl ExternalContext {
+    /// Create a new [ExternalContext] instance from a given [Compiler] instance.
+    pub const fn new(compiler: Compiler) -> Self {
+        Self { compiler }
+    }
+
+    /// Get a compiled function if one exists, otherwise send the bytecode to the compiler to be compiled.
+    pub fn get_compiled_fn(
+        &self,
+        spec_id: SpecId,
+        hash: B256,
+        code: Bytes,
+    ) -> Option<EvmCompilerFn> {
+        self.compiler.get_compiled_fn(spec_id, hash, code)
+    }
+}
+
+/// Registers the compiler handler with the EVM handler.
+pub fn register_compiler_handler<DB>(handler: &mut EvmHandler<'_, ExternalContext, DB>)
+where
+    DB: Database,
+{
+    let f = handler.execution.execute_frame.clone();
+    let spec_id = handler.cfg.spec_id;
+
+    handler.execution.execute_frame = Arc::new(move |frame, memory, table, context| {
+        let Some(action) = execute_frame(spec_id, frame, memory, context) else {
+            dbg!("fallback");
+            return f(frame, memory, table, context);
+        };
+
+        Ok(action)
+    });
+}
+
+fn execute_frame<DB: Database>(
+    spec_id: SpecId,
+    frame: &mut Frame,
+    memory: &mut SharedMemory,
+    context: &mut RevmContext<ExternalContext, DB>,
+) -> Option<InterpreterAction> {
+    // let library = context.external.get_or_load_library(context.evm.spec_id())?;
+    let interpreter = frame.interpreter_mut();
+
+    let hash = match interpreter.contract.hash {
+        Some(hash) => hash,
+        // TODO: is this an issue with EOF?
+        None => unreachable_no_hash(),
+    };
+
+    // should be cheap enough to clone because it's backed by bytes::Bytes
+    let code = interpreter.contract.bytecode.bytes();
+
+    let f = context.external.get_compiled_fn(spec_id, hash, code)?;
+
+    // Safety: as long as the function is still in the cache, this is safe to call
+    let result = unsafe { f.call_with_interpreter_and_memory(interpreter, memory, context) };
+
+    dbg!("EXECUTED", &hash);
+
+    Some(result)
+}
+
+#[cold]
+#[inline(never)]
+const fn unreachable_no_hash() -> ! {
+    panic!("unreachable: bytecode hash is not set in the interpreter")
+}

crates/node/src/evm.rs

@@ -33,16 +33,20 @@ use reth_revm::{
 use revm_precompile::secp256r1;
 use std::sync::Arc;
 
+use crate::compiler::{self, register_compiler_handler, Compiler};
+
 /// Custom EVM configuration
 #[derive(Debug, Clone)]
 pub struct OdysseyEvmConfig {
     chain_spec: Arc<OpChainSpec>,
+    compiler: Compiler,
 }
 
 impl OdysseyEvmConfig {
     /// Creates a new Odyssey EVM configuration with the given chain spec.
-    pub const fn new(chain_spec: Arc<OpChainSpec>) -> Self {
-        Self { chain_spec }
+    pub fn new(chain_spec: Arc<OpChainSpec>) -> Self {
+        let compiler = Compiler::new();
+        Self { chain_spec, compiler }
     }
 
     /// Sets the precompiles to the EVM handler
@@ -221,12 +225,14 @@ impl ConfigureEvmEnv for OdysseyEvmConfig {
 }
 
 impl ConfigureEvm for OdysseyEvmConfig {
-    type DefaultExternalContext<'a> = ();
+    type DefaultExternalContext<'a> = compiler::ExternalContext;
 
     fn evm<DB: Database>(&self, db: DB) -> Evm<'_, Self::DefaultExternalContext<'_>, DB> {
         EvmBuilder::default()
             .with_db(db)
             .optimism()
+            .reset_handler_with_external_context(self.default_external_context())
+            .append_handler_register(register_compiler_handler)
             // add additional precompiles
             .append_handler_register(Self::set_precompiles)
             .build()
@@ -247,7 +253,9 @@ impl ConfigureEvm for OdysseyEvmConfig {
             .build()
     }
 
-    fn default_external_context<'a>(&self) -> Self::DefaultExternalContext<'a> {}
+    fn default_external_context<'a>(&self) -> Self::DefaultExternalContext<'a> {
+        compiler::ExternalContext::new(self.compiler.clone())
+    }
 }
 
 /// Determine the revm spec ID from the current block and reth chainspec.

crates/node/src/lib.rs

@@ -16,5 +16,6 @@
 #![warn(unused_crate_dependencies)]
 
 pub mod chainspec;
+pub mod compiler;
 pub mod evm;
 pub mod node;