Skip to content

Commit 4d01e68

Browse files
0xhenrique0xhenrique
committed
Add --jit flag to opt-in
1 parent fbca271 commit 4d01e68

File tree

2 files changed

+55
-29
lines changed

2 files changed

+55
-29
lines changed

‎src/main.rs‎

Lines changed: 19 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -7,13 +7,15 @@ use std::process;
77
struct Config {
88
filename: Option<String>,
99
arena_enabled: bool,
10+
jit_enabled: bool,
1011
}
1112

1213
fn parse_args() -> Config {
1314
let args: Vec<String> = env::args().collect();
1415
let mut config = Config {
1516
filename: None,
1617
arena_enabled: false, // Arena is opt-in via --arena flag
18+
jit_enabled: false, // JIT is opt-in via --jit flag
1719
};
1820

1921
let mut i = 1;
@@ -22,6 +24,9 @@ fn parse_args() -> Config {
2224
"--arena" => {
2325
config.arena_enabled = true;
2426
}
27+
"--jit" => {
28+
config.jit_enabled = true;
29+
}
2530
"--help" | "-h" => {
2631
print_usage();
2732
process::exit(0);
@@ -46,6 +51,7 @@ fn print_usage() {
4651
eprintln!();
4752
eprintln!("Options:");
4853
eprintln!(" --arena Enable arena allocation for cons cells (experimental)");
54+
eprintln!(" --jit Enable JIT compilation for hot functions (experimental)");
4955
eprintln!(" --help Show this help message");
5056
eprintln!();
5157
eprintln!("If FILE is provided, execute it. Otherwise, start the REPL.");
@@ -57,16 +63,16 @@ fn main() {
5763
set_arena_enabled(config.arena_enabled);
5864

5965
if let Some(filename) = config.filename {
60-
if let Err(e) = run_file(&filename, config.arena_enabled) {
66+
if let Err(e) = run_file(&filename, config.arena_enabled, config.jit_enabled) {
6167
eprintln!("Error: {}", e);
6268
process::exit(1);
6369
}
6470
} else {
65-
run_repl(config.arena_enabled);
71+
run_repl(config.arena_enabled, config.jit_enabled);
6672
}
6773
}
6874

69-
fn run_file(filename: &str, arena_enabled: bool) -> Result<Value, String> {
75+
fn run_file(filename: &str, arena_enabled: bool, jit_enabled: bool) -> Result<Value, String> {
7076
let contents = fs::read_to_string(filename)
7177
.map_err(|e| format!("Could not read file '{}': {}", filename, e))?;
7278

@@ -88,6 +94,9 @@ fn run_file(filename: &str, arena_enabled: bool) -> Result<Value, String> {
8894

8995
// Execute via bytecode VM
9096
let mut vm = standard_vm();
97+
if jit_enabled {
98+
vm.enable_jit();
99+
}
91100
let result = vm.run(chunk)?;
92101

93102
// Arena cleanup
@@ -103,11 +112,14 @@ fn run_file(filename: &str, arena_enabled: bool) -> Result<Value, String> {
103112
Ok(final_result)
104113
}
105114

106-
fn run_repl(arena_enabled: bool) {
115+
fn run_repl(arena_enabled: bool, jit_enabled: bool) {
107116
println!("Lisp VM v0.1.0 (bytecode)");
108117
if arena_enabled {
109118
println!("Arena allocation: enabled");
110119
}
120+
if jit_enabled {
121+
println!("JIT compilation: enabled");
122+
}
111123
println!("Type :q or :quit to exit.");
112124
println!();
113125

@@ -117,6 +129,9 @@ fn run_repl(arena_enabled: bool) {
117129

118130
// Single VM instance to maintain globals across expressions
119131
let mut vm = standard_vm();
132+
if jit_enabled {
133+
vm.enable_jit();
134+
}
120135

121136
let stdin = io::stdin();
122137
let mut stdout = io::stdout();

‎src/vm.rs‎

Lines changed: 36 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -57,10 +57,12 @@ pub struct VM {
5757
compiled_fn_cache: HashMap<SymbolKey, Rc<Chunk>>,
5858
// For native functions, we only need the function pointer
5959
native_fn_cache: HashMap<SymbolKey, fn(&[Value]) -> Result<Value, String>>,
60-
// JIT compiler (lazily initialized)
60+
// JIT compiler (lazily initialized when jit_enabled is true)
6161
jit: Option<JitCompiler>,
6262
// Call counts for JIT compilation triggering (chunk pointer -> count)
6363
call_counts: HashMap<usize, u32>,
64+
// JIT enabled flag (opt-in via --jit flag)
65+
jit_enabled: bool,
6466
}
6567

6668
impl VM {
@@ -74,6 +76,7 @@ impl VM {
7476
native_fn_cache: HashMap::new(),
7577
jit: None,
7678
call_counts: HashMap::new(),
79+
jit_enabled: false,
7780
}
7881
}
7982

@@ -87,9 +90,15 @@ impl VM {
8790
native_fn_cache: HashMap::new(),
8891
jit: None,
8992
call_counts: HashMap::new(),
93+
jit_enabled: false,
9094
}
9195
}
9296

97+
/// Enable JIT compilation (opt-in via --jit flag)
98+
pub fn enable_jit(&mut self) {
99+
self.jit_enabled = true;
100+
}
101+
93102
pub fn define_global(&mut self, name: &str, value: Value) {
94103
// Avoid String allocation if key already exists
95104
{
@@ -431,33 +440,35 @@ impl VM {
431440
}
432441
}
433442

434-
// Try JIT execution first (fast path)
435-
if let Some(jit) = &self.jit {
436-
if let Some(jit_code) = jit.get_compiled(chunk_ptr) {
437-
let result = unsafe {
438-
jit_code.execute(&mut self.registers[new_base..new_base + cf_num_registers])
439-
};
440-
if let Ok(value) = result {
441-
// JIT succeeded - store result and continue
442-
unsafe { *self.registers.get_unchecked_mut(base + dest as usize) = value };
443-
continue;
443+
// Try JIT execution first (fast path) - only if JIT is enabled
444+
if self.jit_enabled {
445+
if let Some(jit) = &self.jit {
446+
if let Some(jit_code) = jit.get_compiled(chunk_ptr) {
447+
let result = unsafe {
448+
jit_code.execute(&mut self.registers[new_base..new_base + cf_num_registers])
449+
};
450+
if let Ok(value) = result {
451+
// JIT succeeded - store result and continue
452+
unsafe { *self.registers.get_unchecked_mut(base + dest as usize) = value };
453+
continue;
454+
}
444455
}
445456
}
446-
}
447457

448-
// JIT not available or failed - fall back to interpreter
449-
// Update call count for potential future JIT compilation
450-
{
451-
let count = self.call_counts.entry(chunk_ptr).or_insert(0);
452-
*count += 1;
453-
if *count == JIT_THRESHOLD && is_jit_compatible(&cf_chunk) {
454-
// Trigger JIT compilation (lazily, will be used on next call)
455-
if self.jit.is_none() {
456-
self.jit = JitCompiler::new().ok();
457-
}
458-
if let Some(jit) = &mut self.jit {
459-
let name = format!("jit_func_{:x}", chunk_ptr);
460-
let _ = jit.compile_function(&cf_chunk, &name);
458+
// JIT not available or failed - fall back to interpreter
459+
// Update call count for potential future JIT compilation
460+
{
461+
let count = self.call_counts.entry(chunk_ptr).or_insert(0);
462+
*count += 1;
463+
if *count == JIT_THRESHOLD && is_jit_compatible(&cf_chunk) {
464+
// Trigger JIT compilation (lazily, will be used on next call)
465+
if self.jit.is_none() {
466+
self.jit = JitCompiler::new().ok();
467+
}
468+
if let Some(jit) = &mut self.jit {
469+
let name = format!("jit_func_{:x}", chunk_ptr);
470+
let _ = jit.compile_function(&cf_chunk, &name);
471+
}
461472
}
462473
}
463474
}

0 commit comments

Comments
 (0)