Skip to main content

revmc_runtime/
revm_evm.rs

1//! Generic `revm` JIT EVM.
2//!
3//! Provides [`JitEvm`] which wraps any mainnet-shaped [`EvmTr`]-based EVM and overrides
4//! execution to look up compiled functions via [`JitBackend`] before falling
5//! back to the interpreter.
6
7use crate::runtime::{JitBackend, LookupDecision, LookupRequest, RuntimeCacheKey};
8use alloy_primitives::{Address, Bytes};
9use revm_context::{Host, JournalEntry};
10use revm_context_interface::{
11    Cfg, ContextSetters, ContextTr, Database,
12    journaled_state::JournalTr,
13    result::{EVMError, HaltReason, InvalidTransaction, ResultAndState},
14};
15use revm_database_interface::DatabaseCommit;
16use revm_handler::{
17    EthFrame, EvmTr, ExecuteCommitEvm, ExecuteEvm, FrameInitOrResult, FrameTr, Handler,
18    ItemOrResult, MainnetHandler, PrecompileProvider, SystemCallCommitEvm, SystemCallEvm,
19    SystemCallTx, evm::ContextDbError,
20};
21use revm_inspector::{
22    InspectCommitEvm, InspectEvm, InspectSystemCallEvm, Inspector, InspectorEvmTr, InspectorFrame,
23    InspectorHandler, JournalExt,
24};
25use revm_interpreter::{InstructionResult, InterpreterAction, InterpreterResult, InterpreterTypes};
26use revm_primitives::{B256Map, hardfork::SpecId};
27use revm_state::EvmState;
28use std::sync::Arc;
29
30/// Mainnet EVM with JIT-compiled function dispatch.
31///
32/// Wraps an inner [`EvmTr`] and uses [`JitBackend`] to look up compiled
33/// functions before falling back to the interpreter.
34///
35/// Implements [`ExecuteEvm`] so it can be used as a drop-in replacement for
36/// the standard mainnet EVM.
37#[derive(derive_more::Debug)]
38pub struct JitEvm<EVM> {
39    inner: EVM,
40    backend: JitBackend,
41    /// Cached lookup decisions keyed by `code_hash` alone.
42    /// Invalidated when the `spec_id` changes.
43    #[debug(skip)]
44    lookup_cache: B256Map<LookupDecision>,
45    /// The `spec_id` the cache was built for; cleared on mismatch.
46    lookup_cache_spec_id: SpecId,
47}
48
49impl<EVM> JitEvm<EVM> {
50    /// Returns a reference to the inner EVM.
51    pub const fn inner(&self) -> &EVM {
52        &self.inner
53    }
54
55    /// Returns a mutable reference to the inner EVM.
56    pub fn inner_mut(&mut self) -> &mut EVM {
57        &mut self.inner
58    }
59
60    /// Consumes the JIT EVM and returns the inner EVM.
61    pub fn into_inner(self) -> EVM {
62        self.inner
63    }
64
65    /// Returns a reference to the JIT backend.
66    pub const fn backend(&self) -> &JitBackend {
67        &self.backend
68    }
69
70    /// Returns a mutable reference to the JIT backend.
71    pub fn backend_mut(&mut self) -> &mut JitBackend {
72        &mut self.backend
73    }
74
75    /// Replaces the JIT backend and clears the lookup cache.
76    pub fn set_backend(&mut self, backend: JitBackend) {
77        self.backend = backend;
78        self.lookup_cache.clear();
79    }
80}
81
82impl<EVM: EvmTr> JitEvm<EVM>
83where
84    EVM::Context: ContextTr,
85{
86    /// Creates a new JIT EVM with a disabled backend.
87    ///
88    /// Equivalent to `JitEvm::new(inner, JitBackend::disabled())`.
89    pub fn disabled(inner: EVM) -> Self {
90        Self::new(inner, JitBackend::disabled())
91    }
92
93    /// Creates a new JIT EVM from an inner EVM and backend.
94    pub fn new(inner: EVM, backend: JitBackend) -> Self {
95        let spec_id: SpecId = inner.ctx_ref().cfg().spec().into();
96        Self {
97            inner,
98            backend,
99            lookup_cache: B256Map::with_capacity_and_hasher(16, Default::default()),
100            lookup_cache_spec_id: spec_id,
101        }
102    }
103
104    fn invalidate_cache(&mut self) {
105        let spec_id: SpecId = self.inner.ctx_ref().cfg().spec().into();
106        if spec_id != self.lookup_cache_spec_id {
107            self.lookup_cache.clear();
108            self.lookup_cache_spec_id = spec_id;
109            // } else {
110            //     self.lookup_cache.retain(|_, v| matches!(v, LookupDecision::Compiled(_)));
111        }
112    }
113}
114
115impl<EVM> core::ops::Deref for JitEvm<EVM> {
116    type Target = EVM;
117
118    #[inline]
119    fn deref(&self) -> &EVM {
120        &self.inner
121    }
122}
123
124impl<EVM> core::ops::DerefMut for JitEvm<EVM> {
125    #[inline]
126    fn deref_mut(&mut self) -> &mut EVM {
127        &mut self.inner
128    }
129}
130
131impl<EVM> EvmTr for JitEvm<EVM>
132where
133    EVM: EvmTr<
134            Frame = EthFrame,
135            Precompiles: PrecompileProvider<EVM::Context, Output = InterpreterResult>,
136        >,
137{
138    type Context = EVM::Context;
139    type Instructions = EVM::Instructions;
140    type Precompiles = EVM::Precompiles;
141    type Frame = EVM::Frame;
142
143    #[inline]
144    fn all(
145        &self,
146    ) -> (
147        &Self::Context,
148        &Self::Instructions,
149        &Self::Precompiles,
150        &revm_context_interface::FrameStack<Self::Frame>,
151    ) {
152        self.inner.all()
153    }
154
155    #[inline]
156    fn all_mut(
157        &mut self,
158    ) -> (
159        &mut Self::Context,
160        &mut Self::Instructions,
161        &mut Self::Precompiles,
162        &mut revm_context_interface::FrameStack<Self::Frame>,
163    ) {
164        self.inner.all_mut()
165    }
166
167    #[inline]
168    fn frame_init(
169        &mut self,
170        frame_input: <Self::Frame as FrameTr>::FrameInit,
171    ) -> Result<
172        ItemOrResult<&mut Self::Frame, <Self::Frame as FrameTr>::FrameResult>,
173        ContextDbError<Self::Context>,
174    > {
175        self.inner.frame_init(frame_input)
176    }
177
178    #[inline]
179    fn frame_run(
180        &mut self,
181    ) -> Result<FrameInitOrResult<Self::Frame>, ContextDbError<Self::Context>> {
182        let spec_id: SpecId = self.inner.ctx_ref().cfg().spec().into();
183        let frame = self.inner.frame_stack().get();
184        let code_hash = frame.interpreter.bytecode.get_or_calculate_hash();
185
186        let decision = self.lookup_cache.entry(code_hash).or_insert_with(|| {
187            let code = frame.interpreter.bytecode.original_bytes();
188            self.backend.lookup(LookupRequest { key: RuntimeCacheKey { code_hash, spec_id }, code })
189        });
190
191        Ok(match decision {
192            LookupDecision::Compiled(program) => {
193                let (ctx, _, _, frame_stack) = self.inner.all_mut();
194                let frame = frame_stack.get();
195                let action =
196                    unsafe { program.func.call_with_interpreter(&mut frame.interpreter, ctx) };
197                frame.process_next_action::<_, ContextDbError<Self::Context>>(ctx, action).inspect(
198                    |i| {
199                        if i.is_result() {
200                            frame.set_finished(true);
201                        }
202                    },
203                )?
204            }
205            LookupDecision::Interpret(_) => self.inner.frame_run()?,
206        })
207    }
208
209    #[inline]
210    fn frame_return_result(
211        &mut self,
212        result: <Self::Frame as FrameTr>::FrameResult,
213    ) -> Result<Option<<Self::Frame as FrameTr>::FrameResult>, ContextDbError<Self::Context>> {
214        self.inner.frame_return_result(result)
215    }
216}
217
218impl<EVM> ExecuteEvm for JitEvm<EVM>
219where
220    EVM: EvmTr<
221            Frame = EthFrame,
222            Context: ContextTr<Journal: JournalTr<State = EvmState>> + ContextSetters,
223            Precompiles: PrecompileProvider<EVM::Context, Output = InterpreterResult>,
224        >,
225{
226    type ExecutionResult = revm_context_interface::result::ExecutionResult<HaltReason>;
227    type State = EvmState;
228    type Error = EVMError<<<EVM::Context as ContextTr>::Db as Database>::Error, InvalidTransaction>;
229    type Tx = <EVM::Context as ContextTr>::Tx;
230    type Block = <EVM::Context as ContextTr>::Block;
231
232    #[inline]
233    fn transact_one(&mut self, tx: Self::Tx) -> Result<Self::ExecutionResult, Self::Error> {
234        self.ctx_mut().set_tx(tx);
235        self.invalidate_cache();
236        MainnetHandler::default().run(self)
237    }
238
239    #[inline]
240    fn finalize(&mut self) -> Self::State {
241        self.ctx_mut().journal_mut().finalize()
242    }
243
244    #[inline]
245    fn set_block(&mut self, block: Self::Block) {
246        self.ctx_mut().set_block(block);
247    }
248
249    #[inline]
250    fn replay(&mut self) -> Result<ResultAndState<HaltReason>, Self::Error> {
251        self.invalidate_cache();
252        MainnetHandler::default().run(self).map(|result| {
253            let state = self.ctx_mut().journal_mut().finalize();
254            ResultAndState::new(result, state)
255        })
256    }
257}
258
259impl<EVM> ExecuteCommitEvm for JitEvm<EVM>
260where
261    Self: ExecuteEvm<State = EvmState>,
262    EVM: EvmTr<Context: ContextTr<Db: DatabaseCommit>>,
263{
264    #[inline]
265    fn commit(&mut self, state: Self::State) {
266        self.ctx_mut().db_mut().commit(state);
267    }
268}
269
270impl<EVM> SystemCallEvm for JitEvm<EVM>
271where
272    EVM: EvmTr<
273            Frame = EthFrame,
274            Context: ContextTr<Journal: JournalTr<State = EvmState>, Tx: SystemCallTx>
275                         + ContextSetters,
276            Precompiles: PrecompileProvider<EVM::Context, Output = InterpreterResult>,
277        >,
278{
279    #[inline]
280    fn system_call_one_with_caller(
281        &mut self,
282        caller: Address,
283        system_contract_address: Address,
284        data: Bytes,
285    ) -> Result<Self::ExecutionResult, Self::Error> {
286        self.ctx_mut().set_tx(
287            <<EVM::Context as ContextTr>::Tx as SystemCallTx>::new_system_tx_with_caller(
288                caller,
289                system_contract_address,
290                data,
291            ),
292        );
293        self.invalidate_cache();
294        MainnetHandler::default().run_system_call(self)
295    }
296}
297
298impl<EVM> SystemCallCommitEvm for JitEvm<EVM>
299where
300    Self: SystemCallEvm<State = EvmState> + ExecuteCommitEvm,
301    EVM: EvmTr<Context: ContextTr<Db: DatabaseCommit>>,
302{
303    #[inline]
304    fn system_call_with_caller_commit(
305        &mut self,
306        caller: Address,
307        system_contract_address: Address,
308        data: Bytes,
309    ) -> Result<Self::ExecutionResult, Self::Error> {
310        self.system_call_with_caller(caller, system_contract_address, data).map(|output| {
311            self.ctx_mut().db_mut().commit(output.state);
312            output.result
313        })
314    }
315}
316
317impl<EVM> InspectorEvmTr for JitEvm<EVM>
318where
319    EVM: EvmTr<
320            Frame = EthFrame,
321            Precompiles: PrecompileProvider<EVM::Context, Output = InterpreterResult>,
322        > + InspectorEvmTr,
323{
324    type Inspector = <EVM as InspectorEvmTr>::Inspector;
325
326    #[inline]
327    fn all_inspector(
328        &self,
329    ) -> (
330        &Self::Context,
331        &Self::Instructions,
332        &Self::Precompiles,
333        &revm_context_interface::FrameStack<Self::Frame>,
334        &Self::Inspector,
335    ) {
336        self.inner.all_inspector()
337    }
338
339    #[inline]
340    fn all_mut_inspector(
341        &mut self,
342    ) -> (
343        &mut Self::Context,
344        &mut Self::Instructions,
345        &mut Self::Precompiles,
346        &mut revm_context_interface::FrameStack<Self::Frame>,
347        &mut Self::Inspector,
348    ) {
349        self.inner.all_mut_inspector()
350    }
351
352    #[inline]
353    fn inspect_frame_run(
354        &mut self,
355    ) -> Result<FrameInitOrResult<Self::Frame>, ContextDbError<Self::Context>> {
356        let spec_id: SpecId = self.inner.ctx_ref().cfg().spec().into();
357        let frame = self.inner.frame_stack().get();
358        let code_hash = frame.interpreter.bytecode.get_or_calculate_hash();
359
360        let decision = self.lookup_cache.entry(code_hash).or_insert_with(|| {
361            let code = frame.interpreter.bytecode.original_bytes();
362            self.backend.lookup(LookupRequest { key: RuntimeCacheKey { code_hash, spec_id }, code })
363        });
364
365        let program = match decision {
366            LookupDecision::Compiled(program) => Arc::clone(program),
367            LookupDecision::Interpret(_) => return self.inner.inspect_frame_run(),
368        };
369
370        // Set up the on_log callback to forward logs to the inspector during JIT execution.
371        //
372        // SAFETY: We dereference raw pointers to the inspector and context inside the closure.
373        // This creates aliasing &mut references with `ecx.host` (which borrows the context),
374        // but the inspector's `log()` implementation only accesses the inspector itself,
375        // not the journaled state already borrowed by the host.
376        let (ctx, inspector) = self.ctx_inspector();
377        let ctx_ptr: *mut EVM::Context = ctx;
378        let inspector_ptr: *mut <EVM as InspectorEvmTr>::Inspector = inspector;
379        let mut on_log = move |log: &revm_primitives::Log| unsafe {
380            (*inspector_ptr).log(&mut *ctx_ptr, log.clone());
381        };
382
383        let (ctx, _, _, frame_stack) = self.inner.all_mut();
384        let frame = frame_stack.get();
385        let action = unsafe {
386            program.func.call_with_interpreter_with(&mut frame.interpreter, ctx, |ecx| {
387                // SAFETY: `on_log` lives on the stack and outlives the JIT call.
388                // The closure captures raw pointers whose types may not be
389                // `'static`, so we erase the lifetime via pointer cast.
390                ecx.on_log = Some(core::mem::transmute::<
391                    &mut dyn FnMut(&revm_primitives::Log),
392                    &mut (dyn FnMut(&revm_primitives::Log) + '_),
393                >(&mut on_log));
394            })
395        };
396
397        // Handle selfdestruct.
398        let (ctx, inspector) = self.ctx_inspector();
399        if let InterpreterAction::Return(result) = &action
400            && result.result == InstructionResult::SelfDestruct
401        {
402            inspect_selfdestruct(ctx, inspector);
403        }
404
405        let (ctx, _, _, frame_stack) = self.inner.all_mut();
406        let frame = frame_stack.get();
407        let mut result = frame.process_next_action::<_, ContextDbError<Self::Context>>(ctx, action);
408
409        if let Ok(ItemOrResult::Result(frame_result)) = &mut result {
410            let (ctx, inspector, frame) = self.ctx_inspector_frame();
411            if let Some(frame) = frame.eth_frame() {
412                revm_inspector::handler::frame_end(ctx, inspector, &frame.input, frame_result);
413                frame.set_finished(true);
414            }
415        }
416        result
417    }
418}
419
420impl<EVM> InspectEvm for JitEvm<EVM>
421where
422    EVM: EvmTr<
423            Frame = EthFrame,
424            Context: ContextTr<Journal: JournalTr<State = EvmState> + JournalExt> + ContextSetters,
425            Precompiles: PrecompileProvider<EVM::Context, Output = InterpreterResult>,
426        > + InspectorEvmTr,
427{
428    type Inspector = <EVM as InspectorEvmTr>::Inspector;
429
430    #[inline]
431    fn set_inspector(&mut self, inspector: Self::Inspector) {
432        *self.inner.inspector() = inspector;
433    }
434
435    #[inline]
436    fn inspect_one_tx(&mut self, tx: Self::Tx) -> Result<Self::ExecutionResult, Self::Error> {
437        self.inner.ctx_mut().set_tx(tx);
438        self.invalidate_cache();
439        MainnetHandler::default().inspect_run(self)
440    }
441}
442
443impl<EVM> InspectCommitEvm for JitEvm<EVM> where Self: InspectEvm + ExecuteCommitEvm {}
444
445impl<EVM> InspectSystemCallEvm for JitEvm<EVM>
446where
447    EVM: EvmTr<
448            Frame = EthFrame,
449            Context: ContextTr<
450                Journal: JournalTr<State = EvmState> + JournalExt,
451                Tx: SystemCallTx,
452            > + ContextSetters,
453            Precompiles: PrecompileProvider<EVM::Context, Output = InterpreterResult>,
454        > + InspectorEvmTr,
455{
456    #[inline]
457    fn inspect_one_system_call_with_caller(
458        &mut self,
459        caller: Address,
460        system_contract_address: Address,
461        data: Bytes,
462    ) -> Result<Self::ExecutionResult, Self::Error> {
463        self.inner.ctx_mut().set_tx(
464            <<EVM::Context as ContextTr>::Tx as SystemCallTx>::new_system_tx_with_caller(
465                caller,
466                system_contract_address,
467                data,
468            ),
469        );
470        self.invalidate_cache();
471        MainnetHandler::default().inspect_run_system_call(self)
472    }
473}
474
475#[inline(never)]
476#[cold]
477fn inspect_selfdestruct<CTX, IT>(context: &mut CTX, inspector: &mut impl Inspector<CTX, IT>)
478where
479    CTX: ContextTr<Journal: JournalExt> + Host,
480    IT: InterpreterTypes,
481{
482    if let Some(
483        JournalEntry::AccountDestroyed {
484            address: contract, target: to, had_balance: balance, ..
485        }
486        | JournalEntry::BalanceTransfer { from: contract, to, balance, .. },
487    ) = context.journal_mut().journal().last()
488    {
489        inspector.selfdestruct(*contract, *to, *balance);
490    }
491}
492
493#[cfg(test)]
494#[cfg(feature = "llvm")]
495mod tests {
496    use super::*;
497    use crate::runtime::{JitBackend, RuntimeConfig};
498    use alloy_primitives::{Address, Bytes, TxKind, U256};
499    use revm_bytecode::opcode as op;
500    use revm_context::TxEnv;
501    use revm_context_interface::result::{ExecutionResult, Output};
502    use revm_database::{CacheDB, EmptyDB};
503    use revm_handler::MainBuilder;
504
505    fn blocking_backend() -> JitBackend {
506        JitBackend::new(RuntimeConfig { blocking: true, ..Default::default() }).unwrap()
507    }
508
509    type TestInnerEvm = revm_handler::MainnetEvm<
510        revm_context::Context<
511            revm_context::BlockEnv,
512            TxEnv,
513            revm_context::CfgEnv,
514            CacheDB<EmptyDB>,
515        >,
516    >;
517
518    fn test_jit_evm(backend: JitBackend) -> JitEvm<TestInnerEvm> {
519        let inner = revm_context::Context::new(CacheDB::default(), SpecId::CANCUN).build_mainnet();
520        JitEvm::new(inner, backend)
521    }
522
523    fn deploy_contract_with_nonce(
524        evm: &mut JitEvm<TestInnerEvm>,
525        bytecode: &[u8],
526        nonce: u64,
527    ) -> Address {
528        let len = bytecode.len();
529        assert!(len <= 255);
530        let offset = 10u8;
531        let mut deploy_code = vec![
532            op::PUSH1,
533            len as u8,
534            op::PUSH1,
535            offset,
536            op::PUSH0,
537            op::CODECOPY,
538            op::PUSH1,
539            len as u8,
540            op::PUSH0,
541            op::RETURN,
542        ];
543        deploy_code.extend_from_slice(bytecode);
544
545        let tx = TxEnv {
546            kind: TxKind::Create,
547            nonce,
548            data: Bytes::copy_from_slice(&deploy_code),
549            gas_limit: 1_000_000,
550            ..Default::default()
551        };
552
553        let result = evm.transact_commit(tx).unwrap();
554        match result {
555            ExecutionResult::Success { output, .. } => match output {
556                Output::Create(_, Some(addr)) => addr,
557                other => panic!("expected Create output, got: {other:?}"),
558            },
559            other => panic!("expected Success, got: {other:?}"),
560        }
561    }
562
563    fn deploy_contract(evm: &mut JitEvm<TestInnerEvm>, bytecode: &[u8]) -> Address {
564        deploy_contract_with_nonce(evm, bytecode, 0)
565    }
566
567    #[test]
568    fn jit_evm_simple_return() {
569        let backend = blocking_backend();
570
571        // PUSH1 0x42 PUSH0 MSTORE PUSH1 0x20 PUSH0 RETURN
572        let runtime_code: &[u8] = &[0x60, 0x42, 0x5f, 0x52, 0x60, 0x20, 0x5f, 0xf3];
573
574        let mut evm = test_jit_evm(backend.clone());
575        let contract_addr = deploy_contract(&mut evm, runtime_code);
576
577        let tx = TxEnv {
578            kind: TxKind::Call(contract_addr),
579            nonce: 1,
580            gas_limit: 1_000_000,
581            ..Default::default()
582        };
583
584        let result = evm.transact(tx).unwrap();
585        match result.result {
586            ExecutionResult::Success { output, .. } => {
587                let data = match &output {
588                    Output::Call(bytes) => bytes,
589                    other => panic!("expected Call output, got: {other:?}"),
590                };
591                assert_eq!(U256::from_be_slice(data), U256::from(0x42));
592            }
593            other => panic!("expected Success, got: {other:?}"),
594        }
595
596        let stats = backend.stats();
597        assert!(stats.compilations_succeeded > 0, "expected compilations, got: {stats:?}");
598        assert!(stats.resident_entries > 0, "expected resident entries, got: {stats:?}");
599    }
600
601    #[test]
602    fn jit_evm_add() {
603        let backend = blocking_backend();
604
605        // PUSH1 1 PUSH1 2 ADD PUSH0 MSTORE PUSH1 0x20 PUSH0 RETURN
606        let runtime_code: &[u8] =
607            &[0x60, 0x01, 0x60, 0x02, 0x01, 0x5f, 0x52, 0x60, 0x20, 0x5f, 0xf3];
608
609        let mut evm = test_jit_evm(backend.clone());
610        let contract_addr = deploy_contract(&mut evm, runtime_code);
611
612        let tx = TxEnv {
613            kind: TxKind::Call(contract_addr),
614            nonce: 1,
615            gas_limit: 1_000_000,
616            ..Default::default()
617        };
618
619        let result = evm.transact(tx).unwrap();
620        match result.result {
621            ExecutionResult::Success { output, .. } => {
622                let data = match &output {
623                    Output::Call(bytes) => bytes,
624                    other => panic!("expected Call output, got: {other:?}"),
625                };
626                assert_eq!(U256::from_be_slice(data), U256::from(3));
627            }
628            other => panic!("expected Success, got: {other:?}"),
629        }
630
631        let stats = backend.stats();
632        assert!(stats.compilations_succeeded > 0, "expected compilations, got: {stats:?}");
633        assert!(stats.resident_entries > 0, "expected resident entries, got: {stats:?}");
634    }
635
636    #[test]
637    fn jit_evm_fallback_empty_code() {
638        let backend = blocking_backend();
639        let mut evm = test_jit_evm(backend);
640
641        let tx = TxEnv {
642            kind: TxKind::Call(Address::with_last_byte(0xEE)),
643            gas_limit: 1_000_000,
644            ..Default::default()
645        };
646
647        let result = evm.transact(tx).unwrap();
648        assert!(
649            matches!(result.result, ExecutionResult::Success { .. }),
650            "expected Success for empty-code call, got: {:?}",
651            result.result,
652        );
653    }
654
655    /// Non-blocking mode: JIT compiles in background and results eventually appear.
656    #[test]
657    fn jit_evm_non_blocking() {
658        use crate::runtime::RuntimeTuning;
659
660        let config = RuntimeConfig {
661            enabled: true,
662            tuning: RuntimeTuning {
663                jit_hot_threshold: 1,
664                jit_worker_count: 1,
665                ..Default::default()
666            },
667            ..Default::default()
668        };
669        let backend = JitBackend::new(config).unwrap();
670
671        let runtime_code: &[u8] = &[0x60, 0x42, 0x5f, 0x52, 0x60, 0x20, 0x5f, 0xf3];
672
673        let mut evm = test_jit_evm(backend.clone());
674        let contract_addr = deploy_contract(&mut evm, runtime_code);
675
676        // First call: should succeed (interpreter fallback or JIT).
677        let tx = TxEnv {
678            kind: TxKind::Call(contract_addr),
679            nonce: 1,
680            gas_limit: 1_000_000,
681            ..Default::default()
682        };
683        let result = evm.transact(tx).unwrap();
684        assert!(matches!(result.result, ExecutionResult::Success { .. }));
685
686        // Wait for JIT to compile in background.
687        let code_hash = alloy_primitives::keccak256(runtime_code);
688        let deadline = std::time::Instant::now() + std::time::Duration::from_secs(30);
689        loop {
690            if backend.get_compiled(code_hash, SpecId::CANCUN).is_some() {
691                break;
692            }
693            assert!(std::time::Instant::now() < deadline, "timed out waiting for JIT");
694            std::thread::sleep(std::time::Duration::from_millis(50));
695        }
696
697        // Second call: should use JIT-compiled function.
698        let mut evm = test_jit_evm(backend);
699        let contract_addr = deploy_contract(&mut evm, runtime_code);
700
701        let tx = TxEnv {
702            kind: TxKind::Call(contract_addr),
703            nonce: 1,
704            gas_limit: 1_000_000,
705            ..Default::default()
706        };
707        let result = evm.transact(tx).unwrap();
708        match result.result {
709            ExecutionResult::Success { output, .. } => {
710                let data = match &output {
711                    Output::Call(bytes) => bytes,
712                    other => panic!("expected Call output, got: {other:?}"),
713                };
714                assert_eq!(U256::from_be_slice(data), U256::from(0x42));
715            }
716            other => panic!("expected Success, got: {other:?}"),
717        }
718    }
719
720    /// CALL into another contract: JIT handles the nested call frame.
721    #[test]
722    fn jit_evm_nested_call() {
723        let backend = blocking_backend();
724        let mut evm = test_jit_evm(backend);
725
726        // Inner contract: PUSH1 0x42 PUSH0 MSTORE PUSH1 0x20 PUSH0 RETURN
727        let inner_code: &[u8] = &[0x60, 0x42, 0x5f, 0x52, 0x60, 0x20, 0x5f, 0xf3];
728        let inner_addr = deploy_contract_with_nonce(&mut evm, inner_code, 0);
729
730        // Outer contract: CALL inner, then copy return data and return it.
731        let mut outer_code = vec![
732            op::PUSH0, // retLen
733            op::PUSH0, // retOff
734            op::PUSH0, // argsLen
735            op::PUSH0, // argsOff
736            op::PUSH0, // value
737        ];
738        // PUSH20 <inner_addr>
739        outer_code.push(op::PUSH20);
740        outer_code.extend_from_slice(inner_addr.as_slice());
741        outer_code.extend_from_slice(&[
742            op::PUSH2,
743            0xFF,
744            0xFF, // gas
745            op::CALL,
746            op::RETURNDATASIZE,
747            op::PUSH0,
748            op::PUSH0,
749            op::RETURNDATACOPY,
750            op::RETURNDATASIZE,
751            op::PUSH0,
752            op::RETURN,
753        ]);
754        let outer_addr = deploy_contract_with_nonce(&mut evm, &outer_code, 1);
755
756        let tx = TxEnv {
757            kind: TxKind::Call(outer_addr),
758            nonce: 2,
759            gas_limit: 1_000_000,
760            ..Default::default()
761        };
762        let result = evm.transact(tx).unwrap();
763        match result.result {
764            ExecutionResult::Success { output, .. } => {
765                let data = match &output {
766                    Output::Call(bytes) => bytes,
767                    other => panic!("expected Call output, got: {other:?}"),
768                };
769                assert_eq!(U256::from_be_slice(data), U256::from(0x42));
770            }
771            other => panic!("expected Success, got: {other:?}"),
772        }
773    }
774
775    /// CREATE2 factory: deploy the same runtime code multiple times with different salts.
776    /// All deployments share the same code hash and should hit the JIT cache.
777    #[test]
778    fn jit_evm_create2_factory() {
779        let backend = blocking_backend();
780        let mut evm = test_jit_evm(backend.clone());
781
782        // Runtime code: PUSH1 0xAB PUSH0 MSTORE PUSH1 0x20 PUSH0 RETURN
783        let runtime_code: &[u8] = &[0x60, 0xAB, 0x5f, 0x52, 0x60, 0x20, 0x5f, 0xf3];
784
785        // Build init code that deploys `runtime_code` (same pattern as deploy_contract_with_nonce).
786        let rt_len = runtime_code.len();
787        let init_offset = 10u8;
788        let mut init_code = vec![
789            op::PUSH1,
790            rt_len as u8,
791            op::PUSH1,
792            init_offset,
793            op::PUSH0,
794            op::CODECOPY,
795            op::PUSH1,
796            rt_len as u8,
797            op::PUSH0,
798            op::RETURN,
799        ];
800        init_code.extend_from_slice(runtime_code);
801        let init_len = init_code.len();
802
803        // Factory contract: copies embedded init_code to memory, reads salt from
804        // calldata, then CREATE2. Returns the deployed address.
805        let mut factory_code = vec![
806            op::PUSH1,
807            0, // placeholder for init_len
808            op::PUSH1,
809            0, // placeholder for code_offset
810            op::PUSH0,
811            op::CODECOPY, // mem[0..init_len] = init_code
812            op::PUSH0,
813            op::CALLDATALOAD, // salt from calldata[0..32]
814            op::PUSH1,
815            0,         // placeholder for init_len
816            op::PUSH0, // offset
817            op::PUSH0, // value
818            op::CREATE2,
819            op::PUSH0,
820            op::MSTORE,
821            op::PUSH1,
822            0x20,
823            op::PUSH0,
824            op::RETURN,
825        ];
826        let code_offset = factory_code.len() as u8;
827        factory_code[1] = init_len as u8;
828        factory_code[3] = code_offset;
829        factory_code[9] = init_len as u8;
830        factory_code.extend_from_slice(&init_code);
831
832        let factory_addr = deploy_contract(&mut evm, &factory_code);
833
834        // Deploy 3 instances via the factory with different salts.
835        let mut deployed_addrs = Vec::new();
836        for (i, salt) in (1u64..=3).enumerate() {
837            let mut calldata = [0u8; 32];
838            calldata[24..].copy_from_slice(&salt.to_be_bytes());
839
840            let tx = TxEnv {
841                kind: TxKind::Call(factory_addr),
842                nonce: 1 + i as u64,
843                data: Bytes::copy_from_slice(&calldata),
844                gas_limit: 1_000_000,
845                ..Default::default()
846            };
847            let result = evm.transact_commit(tx).unwrap();
848            match result {
849                ExecutionResult::Success { output, .. } => {
850                    let data = match &output {
851                        Output::Call(bytes) => bytes,
852                        other => panic!("expected Call output, got: {other:?}"),
853                    };
854                    let addr = Address::from_slice(&data[12..32]);
855                    assert_ne!(addr, Address::ZERO, "CREATE2 should return a non-zero address");
856                    deployed_addrs.push(addr);
857                }
858                other => panic!("expected Success, got: {other:?}"),
859            }
860        }
861
862        // All deployed addresses should be different (different salts).
863        deployed_addrs.sort();
864        deployed_addrs.dedup();
865        assert_eq!(
866            deployed_addrs.len(),
867            3,
868            "CREATE2 with different salts should yield different addresses"
869        );
870
871        // Call each deployed contract — all should return 0xAB.
872        for (i, &addr) in deployed_addrs.iter().enumerate() {
873            let tx = TxEnv {
874                kind: TxKind::Call(addr),
875                nonce: 4 + i as u64,
876                gas_limit: 1_000_000,
877                ..Default::default()
878            };
879            let result = evm.transact_commit(tx).unwrap();
880            match result {
881                ExecutionResult::Success { output, .. } => {
882                    let data = match &output {
883                        Output::Call(bytes) => bytes,
884                        other => panic!("expected Call output, got: {other:?}"),
885                    };
886                    assert_eq!(U256::from_be_slice(data), U256::from(0xABu64));
887                }
888                other => panic!("expected Success, got: {other:?}"),
889            }
890        }
891
892        let stats = backend.stats();
893        assert!(stats.compilations_succeeded > 0, "expected compilations, got: {stats:?}");
894    }
895
896    /// CREATE deploys a contract whose runtime code gets JIT-compiled and called.
897    #[test]
898    fn jit_evm_create_and_call() {
899        let backend = blocking_backend();
900        let mut evm = test_jit_evm(backend);
901
902        // Runtime code: PUSH1 0x99 PUSH0 MSTORE PUSH1 0x20 PUSH0 RETURN
903        let runtime_code: &[u8] = &[0x60, 0x99, 0x5f, 0x52, 0x60, 0x20, 0x5f, 0xf3];
904
905        // Deploy via helper (uses CREATE internally).
906        let contract_addr = deploy_contract(&mut evm, runtime_code);
907
908        // Call the deployed contract.
909        let tx = TxEnv {
910            kind: TxKind::Call(contract_addr),
911            nonce: 1,
912            gas_limit: 1_000_000,
913            ..Default::default()
914        };
915        let result = evm.transact(tx).unwrap();
916        match result.result {
917            ExecutionResult::Success { output, .. } => {
918                let data = match &output {
919                    Output::Call(bytes) => bytes,
920                    other => panic!("expected Call output, got: {other:?}"),
921                };
922                assert_eq!(U256::from_be_slice(data), U256::from(0x99u64));
923            }
924            other => panic!("expected Success, got: {other:?}"),
925        }
926    }
927}