Skip to main content

byteflow/jit/
trace.rs

1use std::collections::HashMap;
2use std::ops::Range;
3use std::sync::Arc;
4
5use crate::Chunk;
6
7use super::frame::JitEntry;
8
9/// Maximum instructions recorded into a single trace.
10pub const MAX_TRACE_LENGTH: usize = 256;
11
12/// How many times `(function, pc)` must run before compilation is attempted.
13pub const HOT_THRESHOLD: u32 = 1_000;
14
15#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
16pub struct TraceKey {
17    pub function: u32,
18    pub entry_pc: u32,
19}
20
21#[derive(Debug, Clone)]
22pub struct TraceSpan {
23    pub function: u32,
24    pub range: Range<u32>,
25}
26
27#[derive(Debug)]
28pub struct CompiledTrace {
29    pub span: TraceSpan,
30    pub entry: JitEntry,
31}
32
33pub struct TraceCache {
34    traces: HashMap<TraceKey, CompiledTrace>,
35}
36
37impl TraceCache {
38    pub fn new() -> Self {
39        Self {
40            traces: HashMap::new(),
41        }
42    }
43
44    pub fn get(&self, key: &TraceKey) -> Option<&CompiledTrace> {
45        self.traces.get(key)
46    }
47
48    pub fn insert(&mut self, key: TraceKey, trace: CompiledTrace) {
49        self.traces.insert(key, trace);
50    }
51
52    pub fn clear(&mut self) {
53        self.traces.clear();
54    }
55}
56
57impl Default for TraceCache {
58    fn default() -> Self {
59        Self::new()
60    }
61}
62
63#[derive(Default)]
64pub struct HotCounter {
65    counters: HashMap<TraceKey, u32>,
66}
67
68impl HotCounter {
69    pub fn hit(&mut self, key: TraceKey, threshold: u32) -> bool {
70        let counter = self.counters.entry(key).or_insert(0);
71        *counter = counter.saturating_add(1);
72        *counter >= threshold
73    }
74
75    pub fn reset(&mut self, key: TraceKey) {
76        self.counters.remove(&key);
77    }
78
79    pub fn clear(&mut self) {
80        self.counters.clear();
81    }
82}
83
84pub struct JitContext {
85    pub chunk: Arc<Chunk>,
86    pub cache: TraceCache,
87    pub hot: HotCounter,
88    pub hot_threshold: u32,
89    module: cranelift_jit::JITModule,
90}
91
92impl JitContext {
93    pub fn new(chunk: Arc<Chunk>) -> Result<Self, super::error::CompileError> {
94        use cranelift_codegen::settings;
95
96        let flags = settings::Flags::new(settings::builder());
97        let isa = cranelift_native::builder()
98            .map_err(|e| super::error::CompileError::Backend(e.to_string()))?
99            .finish(flags)
100            .map_err(|e| super::error::CompileError::Backend(e.to_string()))?;
101        let builder = cranelift_jit::JITBuilder::with_isa(isa, cranelift_module::default_libcall_names());
102        let module = cranelift_jit::JITModule::new(builder);
103        Ok(Self {
104            chunk,
105            cache: TraceCache::new(),
106            hot: HotCounter::default(),
107            hot_threshold: HOT_THRESHOLD,
108            module,
109        })
110    }
111
112    pub fn module_mut(&mut self) -> &mut cranelift_jit::JITModule {
113        &mut self.module
114    }
115
116    /// Drop compiled traces and hot counters after the bytecode image changes.
117    pub fn reload(&mut self, chunk: Arc<Chunk>) {
118        self.chunk = chunk;
119        self.cache.clear();
120        self.hot.clear();
121    }
122}