Skip to main content

zkvmc_jit/
cache.rs

1use std::ptr::NonNull;
2
3use rustc_hash::FxHashMap;
4use zkvmc_context::{Compiler, Context, JitFn};
5
6/// A JIT cache that stores compiled functions.
7pub struct JitCache<C> {
8    inner: C,
9    cache: FxHashMap<u32, JitFn>,
10}
11
12impl<C> JitCache<C> {
13    pub fn new(inner: C) -> Self {
14        Self {
15            inner,
16            cache: FxHashMap::default(),
17        }
18    }
19}
20
21impl<C> Compiler for JitCache<C>
22where
23    C: Compiler,
24{
25    fn compile(&mut self, ctx: NonNull<Context>) -> JitFn {
26        *self
27            .cache
28            .entry(unsafe { ctx.as_ref().pc })
29            .or_insert_with(|| self.inner.compile(ctx))
30    }
31}