Skip to main content

just_engine/runner/jit/
mod.rs

1//! JIT compilation module for the JavaScript engine.
2//!
3//! Provides a bytecode compiler and stack-based VM as an alternative
4//! execution path to the tree-walking interpreter. The pipeline is:
5//!
6//! ```text
7//! JavaScript source → Parser → AST → Compiler → Bytecode → VM → Result
8//! ```
9//!
10//! The bytecode representation eliminates per-node dispatch overhead
11//! and improves cache locality compared to recursive AST evaluation.
12
13pub mod bytecode;
14pub mod compiler;
15pub mod reg_bytecode;
16pub mod reg_compiler;
17pub mod reg_jit;
18pub mod reg_vm;
19pub mod vm;
20
21use crate::parser::ast::ProgramData;
22use crate::runner::ds::error::JErrorType;
23use crate::runner::ds::value::JsValue;
24use crate::runner::plugin::registry::BuiltInRegistry;
25use crate::runner::plugin::types::EvalContext;
26
27use self::bytecode::Chunk;
28use self::compiler::Compiler;
29use self::reg_bytecode::RegChunk;
30use self::reg_compiler::RegCompiler;
31use self::reg_jit::RegJit;
32use self::reg_vm::{RegVm, RegVmResult};
33use self::vm::{Vm, VmResult};
34
35/// Compile an AST program into stack-based bytecode.
36///
37/// Transforms the AST into a linear sequence of bytecode instructions
38/// for execution by the stack-based VM.
39///
40/// # Examples
41///
42/// ```
43/// use just::parser::JsParser;
44/// use just::runner::jit;
45///
46/// let code = "var x = 5 + 3;";
47/// let ast = JsParser::parse_to_ast_from_str(code).unwrap();
48/// let chunk = jit::compile(&ast);
49/// ```
50pub fn compile(program: &ProgramData) -> Chunk {
51    let compiler = Compiler::new();
52    compiler.compile_program(program)
53}
54
55/// Execute register bytecode with Cranelift JIT compilation.
56///
57/// Uses the Cranelift code generator to compile numeric-heavy code paths
58/// to native machine code for maximum performance.
59///
60/// # Examples
61///
62/// ```
63/// use just::parser::JsParser;
64/// use just::runner::plugin::types::EvalContext;
65/// use just::runner::jit::reg_compiler::RegCompiler;
66/// use just::runner::jit;
67///
68/// let code = "var sum = 0; for (var i = 0; i < 100; i++) { sum = sum + i; }";
69/// let ast = JsParser::parse_to_ast_from_str(code).unwrap();
70/// let compiler = RegCompiler::new();
71/// let chunk = compiler.compile_program(&ast);
72///
73/// let ctx = EvalContext::new();
74/// let (result, ctx) = jit::execute_reg_jit(&chunk, ctx).unwrap();
75/// ```
76pub fn execute_reg_jit(chunk: &RegChunk, mut ctx: EvalContext) -> Result<(JsValue, EvalContext), JErrorType> {
77    let mut jit = RegJit::new()?;
78    let (result, regs) = jit.execute(chunk)?;
79
80    for local in &chunk.locals {
81        let name = chunk.get_name(local.name_idx);
82        let val = f64_to_jsvalue(regs[local.reg as usize]);
83        if !ctx.has_var_binding(name) {
84            let _ = ctx.create_var_binding(name);
85            let _ = ctx.initialize_var_binding(name, val.clone());
86        } else {
87            let _ = ctx.set_var_binding(name, val.clone());
88        }
89    }
90
91    Ok((f64_to_jsvalue(result), ctx))
92}
93
94/// Execute register bytecode with JIT when possible; fall back to RegVm on failure.
95pub fn execute_reg_jit_or_vm(
96    chunk: &RegChunk,
97    mut ctx: EvalContext,
98    registry: &BuiltInRegistry,
99) -> Result<(JsValue, EvalContext), JErrorType> {
100    if let Ok(mut jit) = RegJit::new() {
101        if let Ok((result, regs)) = jit.execute(chunk) {
102            for local in &chunk.locals {
103                let name = chunk.get_name(local.name_idx);
104                let val = f64_to_jsvalue(regs[local.reg as usize]);
105                if !ctx.has_var_binding(name) {
106                    let _ = ctx.create_var_binding(name);
107                    let _ = ctx.initialize_var_binding(name, val.clone());
108                } else {
109                    let _ = ctx.set_var_binding(name, val.clone());
110                }
111            }
112            return Ok((f64_to_jsvalue(result), ctx));
113        }
114    }
115
116    let mut vm = RegVm::new(chunk, ctx, registry);
117    let result = match vm.run() {
118        RegVmResult::Ok(val) => Ok(val),
119        RegVmResult::Error(e) => Err(e),
120    };
121    let ctx_out = vm.into_ctx();
122    result.map(|val| (val, ctx_out))
123}
124
125/// Compile an AST program into register-based bytecode.
126///
127/// Register-based bytecode is more suitable for JIT compilation
128/// compared to stack-based bytecode.
129///
130/// # Examples
131///
132/// ```
133/// use just::parser::JsParser;
134/// use just::runner::jit;
135///
136/// let code = "var x = 5 + 3;";
137/// let ast = JsParser::parse_to_ast_from_str(code).unwrap();
138/// let chunk = jit::compile_reg(&ast);
139/// ```
140pub fn compile_reg(program: &ProgramData) -> RegChunk {
141    let compiler = RegCompiler::new();
142    compiler.compile_program(program)
143}
144
145/// Execute stack-based bytecode.
146///
147/// Runs the bytecode through the stack-based VM.
148///
149/// # Examples
150///
151/// ```
152/// use just::parser::JsParser;
153/// use just::runner::plugin::types::EvalContext;
154/// use just::runner::plugin::registry::BuiltInRegistry;
155/// use just::runner::jit;
156///
157/// let code = "var x = Math.abs(-42);";
158/// let ast = JsParser::parse_to_ast_from_str(code).unwrap();
159/// let chunk = jit::compile(&ast);
160///
161/// let mut ctx = EvalContext::new();
162/// ctx.install_core_builtins(BuiltInRegistry::with_core());
163///
164/// let result = jit::execute(&chunk, ctx).unwrap();
165/// ```
166pub fn execute(chunk: &Chunk, ctx: EvalContext) -> Result<JsValue, JErrorType> {
167    let mut vm = Vm::new(chunk, ctx);
168    match vm.run() {
169        VmResult::Ok(val) => Ok(val),
170        VmResult::Error(e) => Err(e),
171    }
172}
173
174/// Execute a compiled register bytecode chunk.
175pub fn execute_reg(
176    chunk: &RegChunk,
177    ctx: EvalContext,
178    registry: &BuiltInRegistry,
179) -> Result<JsValue, JErrorType> {
180    let mut vm = RegVm::new(chunk, ctx, registry);
181    match vm.run() {
182        RegVmResult::Ok(val) => Ok(val),
183        RegVmResult::Error(e) => Err(e),
184    }
185}
186
187/// Compile and execute a program in one step.
188///
189/// Convenience function that combines [`compile`] and [`execute`].
190///
191/// # Examples
192///
193/// ```
194/// use just::parser::JsParser;
195/// use just::runner::plugin::types::EvalContext;
196/// use just::runner::plugin::registry::BuiltInRegistry;
197/// use just::runner::jit;
198///
199/// let code = "var sum = 0; for (var i = 0; i < 10; i++) { sum = sum + i; } sum";
200/// let ast = JsParser::parse_to_ast_from_str(code).unwrap();
201///
202/// let mut ctx = EvalContext::new();
203/// ctx.install_core_builtins(BuiltInRegistry::with_core());
204///
205/// let result = jit::compile_and_run(&ast, ctx).unwrap();
206/// ```
207pub fn compile_and_run(
208    program: &ProgramData,
209    ctx: EvalContext,
210) -> Result<JsValue, JErrorType> {
211    let chunk = compile(program);
212    execute(&chunk, ctx)
213}
214
215/// Compile and execute a program with the register VM.
216pub fn compile_and_run_reg(
217    program: &ProgramData,
218    ctx: EvalContext,
219    registry: &BuiltInRegistry,
220) -> Result<JsValue, JErrorType> {
221    let chunk = compile_reg(program);
222    execute_reg(&chunk, ctx, registry)
223}
224
225/// Compile and execute, returning the context for variable inspection.
226///
227/// Returns both the result and the evaluation context, allowing you to
228/// inspect variables after execution.
229///
230/// # Examples
231///
232/// ```
233/// use just::parser::JsParser;
234/// use just::runner::plugin::registry::BuiltInRegistry;
235/// use just::runner::jit;
236///
237/// let code = "var x = 42; var y = x * 2;";
238/// let ast = JsParser::parse_to_ast_from_str(code).unwrap();
239///
240/// let (result, ctx) = jit::compile_and_run_with_ctx(&ast);
241/// let y = ctx.get_binding("y").unwrap();
242/// ```
243pub fn compile_and_run_with_ctx(
244    program: &ProgramData,
245) -> (Result<JsValue, JErrorType>, EvalContext) {
246    let chunk = compile(program);
247    let mut ctx = EvalContext::new();
248    ctx.install_core_builtins(BuiltInRegistry::with_core());
249    let mut vm = Vm::new(&chunk, ctx);
250    let result = match vm.run() {
251        VmResult::Ok(val) => Ok(val),
252        VmResult::Error(e) => Err(e),
253    };
254    let ctx_out = vm.into_ctx();
255    (result, ctx_out)
256}
257
258/// Compile and execute with the numeric JIT prototype.
259pub fn compile_and_run_reg_jit_with_ctx(
260    program: &ProgramData,
261) -> (Result<JsValue, JErrorType>, EvalContext) {
262    let chunk = compile_reg(program);
263    let ctx = EvalContext::new();
264    match execute_reg_jit(&chunk, ctx) {
265        Ok((value, ctx_out)) => (Ok(value), ctx_out),
266        Err(err) => (Err(err), EvalContext::new()),
267    }
268}
269
270fn f64_to_jsvalue(n: f64) -> JsValue {
271    if n.is_nan() {
272        JsValue::Number(crate::runner::ds::value::JsNumberType::NaN)
273    } else if n == f64::INFINITY {
274        JsValue::Number(crate::runner::ds::value::JsNumberType::PositiveInfinity)
275    } else if n == f64::NEG_INFINITY {
276        JsValue::Number(crate::runner::ds::value::JsNumberType::NegativeInfinity)
277    } else if n.fract() == 0.0 && n >= i64::MIN as f64 && n <= i64::MAX as f64 {
278        JsValue::Number(crate::runner::ds::value::JsNumberType::Integer(n as i64))
279    } else {
280        JsValue::Number(crate::runner::ds::value::JsNumberType::Float(n))
281    }
282}
283
284/// Compile and execute with register VM, returning the EvalContext.
285pub fn compile_and_run_reg_with_ctx(
286    program: &ProgramData,
287    registry: &BuiltInRegistry,
288) -> (Result<JsValue, JErrorType>, EvalContext) {
289    let chunk = compile_reg(program);
290    let mut ctx = EvalContext::new();
291    ctx.install_core_builtins(BuiltInRegistry::with_core());
292    let mut vm = RegVm::new(&chunk, ctx, registry);
293    let result = match vm.run() {
294        RegVmResult::Ok(val) => Ok(val),
295        RegVmResult::Error(e) => Err(e),
296    };
297    let ctx_out = vm.into_ctx();
298    (result, ctx_out)
299}