imxyy1soope1/nixjit
1
0
Fork 0
Code Issues Pull Requests Actions Projects 1 Releases Activity

chore: comment

Browse Source
This commit is contained in:
imxyy_soope_
2025-08-07 21:00:32 +08:00
parent f946cb2fd1
commit 67cdcfea33
24 changed files with 734 additions and 105 deletions
Download Patch File Download Diff File Expand all files Collapse all files

58
evaluator/nixjit_jit/src/compile.rs
View File

@@ -1,13 +1,29 @@
//! This module defines the `JITCompile` trait and its implementations for
//! various IR types. It provides the translation from LIR to Cranelift IR.
use cranelift::codegen::ir::{self, StackSlot};
use cranelift::prelude::*;
use nixjit_eval::{EvalContext, Value};
use nixjit_eval::Value;
use nixjit_ir::*;
use nixjit_lir::Lir;
use super::{Context, JITContext};
/// A trait for compiling IR nodes to Cranelift IR.
///
/// This trait defines how different IR nodes should be compiled to
/// Cranelift IR instructions that can be executed by the JIT compiler.
pub trait JITCompile<Ctx: JITContext> {
/// Compiles the IR node to Cranelift IR.
///
/// # Arguments
/// * `ctx` - The compilation context
/// * `engine` - The evaluation context value
/// * `env` - The environment value
///
/// # Returns
/// A stack slot containing the compiled result
fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot;
}
@@ -24,6 +40,9 @@ impl<Ctx: JITContext> JITCompile<Ctx> for Lir {
}
impl<Ctx: JITContext> JITCompile<Ctx> for AttrSet {
/// Compiles an attribute set to Cranelift IR.
///
/// This creates a new attribute set and compiles all static attributes into it.
fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot {
let attrs = ctx.create_attrs();
for (k, v) in self.stcs.iter() {
@@ -35,6 +54,9 @@ impl<Ctx: JITContext> JITCompile<Ctx> for AttrSet {
}
impl<Ctx: JITContext> JITCompile<Ctx> for List {
/// Compiles a list to Cranelift IR.
///
/// This creates a new list by compiling all items and storing them in an array.
fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot {
let array = ctx.alloc_array(self.items.len());
for (i, item) in self.items.iter().enumerate() {
@@ -67,6 +89,11 @@ impl<Ctx: JITContext> JITCompile<Ctx> for HasAttr {
}
impl<Ctx: JITContext> JITCompile<Ctx> for BinOp {
/// Compiles a binary operation to Cranelift IR.
///
/// This implementation handles various binary operations like addition, subtraction,
/// division, logical AND/OR, and equality checks. It generates code that checks
/// the types of operands and performs the appropriate operation.
fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot {
use BinOpKind::*;
let lhs = self.lhs.compile(ctx, engine, env);
@@ -328,6 +355,9 @@ impl<Ctx: JITContext> JITCompile<Ctx> for UnOp {
}
impl<Ctx: JITContext> JITCompile<Ctx> for Attr {
/// Compiles an attribute key to Cranelift IR.
///
/// An attribute can be either a static string or a dynamic expression.
fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot {
use Attr::*;
match self {
@@ -338,6 +368,10 @@ impl<Ctx: JITContext> JITCompile<Ctx> for Attr {
}
impl<Ctx: JITContext> JITCompile<Ctx> for Select {
/// Compiles an attribute selection to Cranelift IR.
///
/// This compiles the expression to select from, builds the attribute path,
/// and calls the select helper function.
fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot {
let val = self.expr.compile(ctx, engine, env);
let attrpath = ctx.alloc_array(self.attrpath.len());
@@ -366,6 +400,10 @@ impl<Ctx: JITContext> JITCompile<Ctx> for Select {
}
impl<Ctx: JITContext> JITCompile<Ctx> for If {
/// Compiles an if-expression to Cranelift IR.
///
/// This generates code that evaluates the condition, checks that it's a boolean,
/// and then jumps to the appropriate branch (true or false).
fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot {
let cond = self.cond.compile(ctx, engine, env);
let cond_type = ctx.builder.ins().stack_load(types::I64, cond, 0);
@@ -424,6 +462,10 @@ impl<Ctx: JITContext> JITCompile<Ctx> for If {
}
impl<Ctx: JITContext> JITCompile<Ctx> for Call {
/// Compiles a function call to Cranelift IR.
///
/// This compiles the function expression and all arguments, builds an argument array,
/// and calls the call helper function.
fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot {
let func = self.func.compile(ctx, engine, env);
let args = ctx.alloc_array(self.args.len());
@@ -452,6 +494,10 @@ impl<Ctx: JITContext> JITCompile<Ctx> for Call {
}
impl<Ctx: JITContext> JITCompile<Ctx> for With {
/// Compiles a `with` expression to Cranelift IR.
///
/// This enters a new `with` scope with the compiled namespace, compiles the body expression,
/// and then exits the `with` scope.
fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot {
let namespace = self.namespace.compile(ctx, engine, env);
ctx.enter_with(env, namespace);
@@ -475,6 +521,10 @@ impl<Ctx: JITContext> JITCompile<Ctx> for ConcatStrings {
}
impl<Ctx: JITContext> JITCompile<Ctx> for Const {
/// Compiles a constant value to Cranelift IR.
///
/// This handles boolean, integer, float, and null constants by storing
/// their values and type tags in a stack slot.
fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot {
use nixjit_value::Const::*;
let slot = ctx.alloca();
@@ -507,12 +557,18 @@ impl<Ctx: JITContext> JITCompile<Ctx> for Const {
}
impl<Ctx: JITContext> JITCompile<Ctx> for Str {
/// Compiles a string literal to Cranelift IR.
///
/// This creates a string value from the string literal.
fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot {
ctx.create_string(&self.val)
}
}
impl<Ctx: JITContext> JITCompile<Ctx> for Var {
/// Compiles a variable lookup to Cranelift IR.
///
/// This looks up a variable by its symbol in the current environment.
fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot {
ctx.lookup(env, &self.sym)
}

63
evaluator/nixjit_jit/src/helpers.rs
View File

@@ -1,3 +1,8 @@
//! Helper functions for the JIT compiler.
//!
//! These functions are called from JIT-compiled code to perform operations
//! that are difficult or unsafe to do directly in the generated IR.
use core::{slice, str};
use std::alloc::Layout;
use std::alloc::alloc;
@@ -10,6 +15,10 @@ use nixjit_eval::{AttrSet, EvalContext, List, Value};
use super::JITContext;
/// Helper function to call a function with arguments.
///
/// This function is called from JIT-compiled code to perform function calls.
/// It takes a function value and an array of arguments, and executes the call.
pub extern "C" fn helper_call<Ctx: JITContext>(
func: &mut Value,
args_ptr: *mut Value,
@@ -22,6 +31,9 @@ pub extern "C" fn helper_call<Ctx: JITContext>(
func.call(args.into_iter().map(Ok), ctx).unwrap();
}
/// Helper function to look up a value in the evaluation stack.
///
/// This function is called from JIT-compiled code to access values in the evaluation stack.
pub extern "C" fn helper_lookup_stack<Ctx: JITContext>(
ctx: &Ctx,
offset: usize,
@@ -30,6 +42,9 @@ pub extern "C" fn helper_lookup_stack<Ctx: JITContext>(
ret.write(ctx.lookup_stack(offset).clone());
}
/// Helper function to look up a function argument.
///
/// This function is called from JIT-compiled code to access function arguments.
pub extern "C" fn helper_lookup_arg<Ctx: JITContext>(
ctx: &Ctx,
offset: usize,
@@ -38,6 +53,10 @@ pub extern "C" fn helper_lookup_arg<Ctx: JITContext>(
ret.write(JITContext::lookup_arg(ctx, offset).clone());
}
/// Helper function to look up a variable by name.
///
/// This function is called from JIT-compiled code to perform variable lookups
/// in the current scope and `with` expression scopes.
pub extern "C" fn helper_lookup<Ctx: JITContext>(
ctx: &Ctx,
sym_ptr: *const u8,
@@ -56,6 +75,10 @@ pub extern "C" fn helper_lookup<Ctx: JITContext>(
}
}
/// Helper function to perform attribute selection.
///
/// This function is called from JIT-compiled code to select attributes from
/// an attribute set using a path of attribute names.
pub extern "C" fn helper_select<Ctx: JITContext>(
val: &mut Value,
path_ptr: *mut Value,
@@ -70,6 +93,10 @@ pub extern "C" fn helper_select<Ctx: JITContext>(
.unwrap();
}
/// Helper function to perform attribute selection with a default value.
///
/// This function is called from JIT-compiled code to select attributes from
/// an attribute set, returning a default value if the selection fails.
pub extern "C" fn helper_select_with_default<Ctx: JITContext>(
val: &mut Value,
path_ptr: *mut Value,
@@ -88,10 +115,17 @@ pub extern "C" fn helper_select_with_default<Ctx: JITContext>(
.unwrap();
}
/// Helper function to check equality between two values.
///
/// This function is called from JIT-compiled code to perform equality comparisons.
pub extern "C" fn helper_eq<Ctx: JITContext>(lhs: &mut Value, rhs: &Value) {
lhs.eq(rhs);
}
/// Helper function to create a string value.
///
/// This function is called from JIT-compiled code to create string values
/// from raw byte arrays.
pub unsafe extern "C" fn helper_create_string<Ctx: JITContext>(
ptr: *const u8,
len: usize,
@@ -104,6 +138,10 @@ pub unsafe extern "C" fn helper_create_string<Ctx: JITContext>(
}
}
/// Helper function to create a list value.
///
/// This function is called from JIT-compiled code to create list values
/// from arrays of values.
pub unsafe extern "C" fn helper_create_list<Ctx: JITContext>(
ptr: *mut Value,
len: usize,
@@ -116,12 +154,19 @@ pub unsafe extern "C" fn helper_create_list<Ctx: JITContext>(
}
}
/// Helper function to create an attribute set.
///
/// This function is called from JIT-compiled code to create a new, empty attribute set.
pub unsafe extern "C" fn helper_create_attrs<Ctx: JITContext>(
ret: &mut MaybeUninit<HashMap<String, Value>>,
) {
ret.write(HashMap::new());
}
/// Helper function to add an attribute to an attribute set.
///
/// This function is called from JIT-compiled code to insert a key-value pair
/// into an attribute set.
pub unsafe extern "C" fn helper_push_attr<Ctx: JITContext>(
attrs: &mut HashMap<String, Value>,
sym_ptr: *const u8,
@@ -136,6 +181,10 @@ pub unsafe extern "C" fn helper_push_attr<Ctx: JITContext>(
}
}
/// Helper function to finalize an attribute set.
///
/// This function is called from JIT-compiled code to convert a HashMap into
/// a proper attribute set value.
pub unsafe extern "C" fn helper_finalize_attrs<Ctx: JITContext>(
attrs: NonNull<HashMap<String, Value>>,
ret: &mut MaybeUninit<Value>,
@@ -145,6 +194,10 @@ pub unsafe extern "C" fn helper_finalize_attrs<Ctx: JITContext>(
));
}
/// Helper function to enter a `with` expression scope.
///
/// This function is called from JIT-compiled code to enter a new `with` scope
/// with the given namespace.
pub unsafe extern "C" fn helper_enter_with<Ctx: JITContext>(
ctx: &mut Ctx,
namespace: NonNull<Value>,
@@ -152,14 +205,24 @@ pub unsafe extern "C" fn helper_enter_with<Ctx: JITContext>(
ctx.enter_with(unsafe { namespace.read() }.unwrap_attr_set().into_inner());
}
/// Helper function to exit a `with` expression scope.
///
/// This function is called from JIT-compiled code to exit the current `with` scope.
pub unsafe extern "C" fn helper_exit_with<Ctx: JITContext>(ctx: &mut Ctx) {
ctx.exit_with();
}
/// Helper function to allocate an array of values.
///
/// This function is called from JIT-compiled code to allocate memory for
/// arrays of values, such as function arguments or list elements.
pub unsafe extern "C" fn helper_alloc_array<Ctx: JITContext>(len: usize) -> *mut u8 {
unsafe { alloc(Layout::array::<Value>(len).unwrap()) }
}
/// Helper function for debugging.
///
/// This function is called from JIT-compiled code to print a value for debugging purposes.
pub extern "C" fn helper_dbg<Ctx: JITContext>(value: &Value) {
println!("{value:?}")
}

40
evaluator/nixjit_jit/src/lib.rs
View File

@@ -1,3 +1,14 @@
//! The Just-In-Time (JIT) compilation module for nixjit.
//!
//! This module provides functionality to compile Low-Level IR (LIR) expressions
//! into optimized machine code using Cranelift. The JIT compiler translates
//! Nix expressions into efficient native code for faster evaluation.
//!
//! The main components are:
//! - `JITCompiler`: The core compiler that manages the compilation process
//! - `JITContext`: A trait that provides the execution context for JIT-compiled code
//! - `Context`: An internal compilation context used during code generation
use std::marker::PhantomData;
use std::ops::Deref;
use std::rc::Rc;
@@ -18,15 +29,33 @@ mod helpers;
pub use compile::JITCompile;
use helpers::*;
pub trait JITContext: EvalContext + Sized {
/// A trait that provides the execution context for JIT-compiled code.
///
/// This trait extends `EvalContext` with additional methods needed
/// for JIT compilation, such as stack and argument lookups, and
/// managing `with` expression scopes.
pub trait JITContext: EvalContext {
/// Looks up a value in the evaluation stack by offset.
fn lookup_stack(&self, offset: usize) -> &Value;
/// Looks up a function argument by offset.
fn lookup_arg(&self, offset: usize) -> &Value;
/// Enters a `with` expression scope with the given namespace.
fn enter_with(&mut self, namespace: Rc<HashMap<String, Value>>);
/// Exits the current `with` expression scope.
fn exit_with(&mut self);
}
/// Type alias for a JIT-compiled function.
///
/// This represents a function pointer to JIT-compiled code that takes
/// a context pointer and a mutable value pointer as arguments.
type F<Ctx> = unsafe extern "C" fn(*const Ctx, *mut Value);
/// A JIT-compiled function.
///
/// This struct holds a function pointer to the compiled code and
/// a set of strings that were used during compilation, which need
/// to be kept alive for the function to work correctly.
pub struct JITFunc<Ctx: JITContext> {
func: F<Ctx>,
strings: HashSet<String>,
@@ -39,10 +68,18 @@ impl<Ctx: JITContext> Deref for JITFunc<Ctx> {
}
}
/// The internal compilation context used during code generation.
///
/// This context holds references to the compiler, the Cranelift function builder,
/// and manages resources like stack slots and string literals during compilation.
struct Context<'comp, 'ctx, Ctx: JITContext> {
/// Reference to the JIT compiler.
pub compiler: &'comp mut JITCompiler<Ctx>,
/// The Cranelift function builder used to generate IR.
pub builder: FunctionBuilder<'ctx>,
/// Stack slots available for reuse.
free_slots: Vec<StackSlot>,
/// String literals used during compilation.
strings: HashSet<String>,
}
@@ -374,6 +411,7 @@ impl<'comp, 'ctx, Ctx: JITContext> Context<'comp, 'ctx, Ctx> {
}
}
/// The main JIT compiler that manages the compilation process.
pub struct JITCompiler<Ctx: JITContext> {
ctx: codegen::Context,
module: JITModule,