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,147 @@
+use alloy_primitives::{hex, Bytes, B256};
+use core::panic;
+use reth_primitives::Bytecode;
+use reth_revm::{
+    handler::register::EvmHandler,
+    interpreter::{InterpreterAction, SharedMemory, EMPTY_SHARED_MEMORY},
+    Context as RevmContext, Database, Frame,
+};
+use revmc::{primitives::SpecId, EvmCompilerFn, EvmLlvmBackend, OptimizationLevel};
+use std::{
+    collections::HashMap,
+    sync::{
+        mpsc::{Receiver, Sender},
+        Arc, Mutex,
+    },
+    thread,
+};
+
+#[derive(Debug)]
+pub struct ExternalContext {
+    sender: Sender<(B256, Bytes)>,
+    // TODO: cache shouldn't be here (and should definitely not be wrapped in a mutex)
+    cache: Arc<Mutex<HashMap<B256, EvmCompilerFn>>>,
+}
+
+impl ExternalContext {
+    pub fn new(spec_id: SpecId) -> Self {
+        let cache = Arc::new(Mutex::new(HashMap::new()));
+        let (sender, receiver) = std::sync::mpsc::channel();
+
+        // TODO: graceful shutdown
+        thread::spawn({
+            let cache = cache.clone();
+
+            move || {
+                let context = Box::leak(Box::new(revmc::llvm::Context::create()));
+                // TODO: fail properly here.
+                let backend =
+                    EvmLlvmBackend::new(context, false, OptimizationLevel::Aggressive).unwrap();
+                let mut compiler = revmc::EvmCompiler::new(backend);
+
+                while let Ok((hash, code)) = receiver.recv() {
+                    // 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");
+
+                    cache.lock().unwrap().insert(hash, result);
+
+                    unsafe { compiler.clear().expect("could not clear") };
+                }
+            }
+        });
+
+        Self { sender, cache }
+    }
+
+    pub fn get_compiled_fn(&self, hash: B256, code: Bytes) -> Option<EvmCompilerFn> {
+        let Some(f) = self.cache.lock().unwrap().get(&hash).cloned() else {
+            self.sender.send((hash, code)).unwrap();
+            return None;
+        };
+
+        Some(f)
+    }
+}
+
+pub fn register_compiler_handler<DB>(handler: &mut EvmHandler<'_, ExternalContext, DB>)
+where
+    DB: Database,
+{
+    let f = handler.execution.execute_frame.clone();
+
+    handler.execution.execute_frame = Arc::new(move |frame, memory, table, context| {
+        let Some(action) = execute_frame(frame, memory, context) else {
+            dbg!("fallback");
+            return f(frame, memory, table, context);
+        };
+
+        Ok(action)
+    });
+}
+
+fn execute_frame<DB: Database>(
+    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,
+        None => unreachable_no_hash(),
+    };
+
+    // should be cheap enough to clone because it's backed by bytes::Bytes
+    let code = interpreter.contract.bytecode.bytes();
+
+    // TODO: put rules here for whether or not to compile the function
+    let f = context.external.get_compiled_fn(hash, code)?;
+
+    // let f = match library.get_function(hash) {
+    //     Ok(Some(f)) => f,
+    //     Ok(None) => return None,
+    //     // Shouldn't happen.
+    //     Err(err) => {
+    //         unlikely_log_get_function_error(err, &hash);
+    //         return None;
+    //     }
+    // };
+
+    // interpreter.shared_memory =
+    //     std::mem::replace(memory, reth_revm::interpreter::EMPTY_SHARED_MEMORY);
+    // let result = unsafe { f.call_with_interpreter(interpreter, context) };
+    // *memory = interpreter.take_memory();
+    // Some(result)
+
+    interpreter.shared_memory = std::mem::replace(memory, EMPTY_SHARED_MEMORY);
+
+    let result = unsafe { f.call_with_interpreter(interpreter, context) };
+    *memory = interpreter.take_memory();
+
+    dbg!("executed", &hash);
+
+    Some(result)
+}
+
+#[cold]
+#[inline(never)]
+const fn unreachable_no_hash() -> ! {
+    panic!("unreachable: bytecode hash is not set in the interpreter")
+}
+
+#[cold]
+#[inline(never)]
+const fn unreachable_misconfigured() -> ! {
+    panic!("unreachable: AOT EVM is misconfigured")
+}
+
+#[cold]
+#[inline(never)]
+fn unlikely_log_get_function_error(err: impl std::error::Error, hash: &B256) {
+    tracing::error!(%err, %hash, "failed getting function from shared library");
+}

crates/node/src/evm.rs

@@ -33,6 +33,8 @@ use reth_revm::{
 use revm_precompile::secp256r1;
 use std::sync::Arc;
 
+use crate::compiler::{self, register_compiler_handler};
+
 /// Custom EVM configuration
 #[derive(Debug, Clone)]
 pub struct OdysseyEvmConfig {
@@ -221,12 +223,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 +251,14 @@ impl ConfigureEvm for OdysseyEvmConfig {
             .build()
     }
 
-    fn default_external_context<'a>(&self) -> Self::DefaultExternalContext<'a> {}
+    fn default_external_context<'a>(&self) -> Self::DefaultExternalContext<'a> {
+        let Self { chain_spec } = self;
+
+        // TODO: not sure if this is correct, I don't see a call-site for this method
+        let spec_id = revm_spec(chain_spec, &Head::default());
+
+        compiler::ExternalContext::new(spec_id)
+    }
 }
 
 /// 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;