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

feat: WIP

Browse Source
This commit is contained in:
imxyy_soope_
2025-08-22 21:40:52 +08:00
parent 3c7722a3d2
commit 65bcfcb47b
9 changed files with 312 additions and 154 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
evaluator/nixjit_context/src/downgrade.rs
View File

@@ -34,7 +34,7 @@ impl DowngradeCtx<'_, '_> {
impl DowngradeContext for DowngradeCtx<'_, '_> {
fn new_expr(&mut self, expr: Hir) -> ExprId {
self.irs.push(expr.into());
unsafe { ExprId::from_raw(self.ctx.lirs.len() + self.ctx.hirs.len() + self.irs.len() - 1) }
self.ctx.alloc_id()
}
fn with_expr_mut<T>(&mut self, id: ExprId, f: impl FnOnce(&mut Hir, &mut Self) -> T) -> T {

35
evaluator/nixjit_context/src/eval.rs
View File

@@ -1,3 +1,7 @@
#[cfg(debug_assertions)]
use std::cell::OnceCell;
#[cfg(not(debug_assertions))]
use std::mem::MaybeUninit;
use std::rc::Rc;
use hashbrown::HashMap;
@@ -8,12 +12,40 @@ use nixjit_eval::{Args, EvalContext, Evaluate, StackFrame, Value};
use nixjit_ir::ExprId;
use nixjit_jit::JITContext;
use nixjit_lir::Lir;
use petgraph::visit::{Topo, Walker};
use petgraph::prelude::DiGraph;
use super::Context;
#[derive(Default)]
struct ValueCache(
#[cfg(debug_assertions)]
Option<Value>,
#[cfg(not(debug_assertions))]
MaybeUninit<Value>
);
impl ValueCache {
fn insert(&mut self, val: Value) {
#[cfg(debug_assertions)]
{
assert!(self.0.is_none());
let _ = self.0.insert(val);
}
#[cfg(not(debug_assertions))]
self.0.write(val);
}
}
impl Drop for ValueCache {
fn drop(&mut self) {
#[cfg(not(debug_assertions))]
self.0.assume_init_drop();
}
}
pub struct EvalCtx<'ctx, 'bump> {
ctx: &'ctx mut Context<'bump>,
graph: DiGraph<ValueCache, ()>,
stack: Vec<StackFrame>,
with_scopes: Vec<Rc<HashMap<String, Value>>>,
}
@@ -22,6 +54,7 @@ impl<'ctx, 'bump> EvalCtx<'ctx, 'bump> {
pub fn new(ctx: &'ctx mut Context<'bump>) -> Self {
Self {
ctx,
graph: DiGraph::new(),
stack: Vec::new(),
with_scopes: Vec::new(),
}

89
evaluator/nixjit_context/src/lib.rs
View File

@@ -1,4 +1,7 @@
use std::{marker::PhantomPinned, ops::{Deref, DerefMut}};
use std::cell::Cell;
use std::marker::PhantomPinned;
use std::ops::{Deref, DerefMut};
use std::ptr::NonNull;
use bumpalo::{Bump, boxed::Box};
use hashbrown::HashMap;
@@ -10,12 +13,12 @@ use petgraph::{
use nixjit_builtins::{
Builtins, BuiltinsContext,
builtins::{CONSTS_LEN, GLOBAL_LEN, SCOPED_LEN},
builtins::{GLOBAL_LEN, SCOPED_LEN},
};
use nixjit_error::{Error, Result};
use nixjit_eval::{Args, EvalContext, Value};
use nixjit_hir::{DowngradeContext, Hir};
use nixjit_ir::{AttrSet, Const, ExprId, Param, PrimOpId, StackIdx};
use nixjit_ir::{AttrSet, ExprId, Param, PrimOpId, StackIdx};
use nixjit_lir::{Lir, ResolveContext};
use crate::downgrade::DowngradeCtx;
@@ -65,13 +68,14 @@ impl<'bump, T> Pin<'bump, T> {
/// This struct orchestrates the entire Nix expression evaluation process,
/// from parsing and semantic analysis to interpretation and JIT compilation.
pub struct Context<'bump> {
ir_count: usize,
hirs: Vec<Hir>,
lirs: Vec<Pin<'bump, Lir>>,
/// Maps a function's body `ExprId` to its parameter definition.
funcs: HashMap<ExprId, Param>,
repl_scope: HashMap<String, ExprId>,
global_scope: HashMap<&'static str, ExprId>,
repl_scope: NonNull<HashMap<String, ExprId>>,
global_scope: NonNull<HashMap<&'static str, ExprId>>,
/// A dependency graph between expressions.
graph: DiGraphMap<ExprId, StackIdx>,
@@ -82,25 +86,24 @@ pub struct Context<'bump> {
bump: &'bump Bump,
}
impl Drop for Context<'_> {
fn drop(&mut self) {
unsafe {
self.repl_scope.drop_in_place();
self.global_scope.drop_in_place();
}
}
}
impl<'bump> Context<'bump> {
pub fn new(bump: &'bump Bump) -> Self {
let Builtins {
consts,
global,
scoped,
} = Builtins::new();
let global_scope = consts
let Builtins { global, scoped } = Builtins::new();
let global_scope = global
.iter()
.enumerate()
.map(|(id, (k, _))| (*k, unsafe { ExprId::from_raw(id) }))
.chain(
global
.iter()
.enumerate()
.map(|(idx, (k, _, _))| (*k, unsafe { ExprId::from_raw(idx + CONSTS_LEN) })),
)
.map(|(idx, (k, _, _))| (*k, unsafe { ExprId::from_raw(idx) }))
.chain(core::iter::once(("builtins", unsafe {
ExprId::from_raw(CONSTS_LEN + GLOBAL_LEN + SCOPED_LEN)
ExprId::from_raw(GLOBAL_LEN + SCOPED_LEN)
})))
.collect();
let primops = global
@@ -109,47 +112,41 @@ impl<'bump> Context<'bump> {
.chain(scoped.iter().map(|&(_, arity, f)| (arity, f)))
.collect_array()
.unwrap();
let lirs = consts
.into_iter()
.map(|(_, val)| Lir::Const(Const { val }))
.chain((0..global.len()).map(|idx| Lir::PrimOp(unsafe { PrimOpId::from_raw(idx) })))
let lirs = (0..global.len()).map(|idx| Lir::PrimOp(unsafe { PrimOpId::from_raw(idx) }))
.chain(
(0..scoped.len())
.map(|idx| Lir::PrimOp(unsafe { PrimOpId::from_raw(idx + GLOBAL_LEN) })),
)
.chain(core::iter::once(Lir::AttrSet(AttrSet {
stcs: consts
.into_iter()
.enumerate()
.map(|(idx, (name, _))| (name.to_string(), unsafe { ExprId::from_raw(idx) }))
.chain(global.into_iter().enumerate().map(|(idx, (name, ..))| {
stcs: global.into_iter().enumerate().map(|(idx, (name, ..))| {
(name.to_string(), unsafe {
ExprId::from_raw(idx + CONSTS_LEN)
ExprId::from_raw(idx)
})
}))
})
.chain(scoped.into_iter().enumerate().map(|(idx, (name, ..))| {
(name.to_string(), unsafe {
ExprId::from_raw(idx + CONSTS_LEN + GLOBAL_LEN)
ExprId::from_raw(idx + GLOBAL_LEN)
})
}))
.chain(core::iter::once(("builtins".to_string(), unsafe {
ExprId::from_raw(CONSTS_LEN + GLOBAL_LEN + SCOPED_LEN + 1)
ExprId::from_raw(GLOBAL_LEN + SCOPED_LEN + 1)
})))
.collect(),
..AttrSet::default()
})))
.chain(core::iter::once(Lir::Thunk(unsafe {
ExprId::from_raw(CONSTS_LEN + GLOBAL_LEN + SCOPED_LEN)
ExprId::from_raw(GLOBAL_LEN + SCOPED_LEN)
})))
.map(|lir| Pin::new_in(lir, bump))
.collect();
Self {
ir_count: 0,
hirs: Vec::new(),
lirs,
funcs: HashMap::new(),
global_scope,
repl_scope: HashMap::new(),
global_scope: NonNull::from(bump.alloc(global_scope)),
repl_scope: NonNull::from(bump.alloc(HashMap::new())),
graph: DiGraphMap::new(),
primops,
@@ -207,9 +204,29 @@ impl<'bump> Context<'bump> {
let root_expr = root.tree().expr().unwrap();
let expr_id = self.downgrade_ctx().downgrade_root(root_expr)?;
self.resolve_ctx().resolve_root(expr_id)?;
self.repl_scope.insert(ident.to_string(), expr_id);
unsafe { self.repl_scope.as_mut() }.insert(ident.to_string(), expr_id);
Ok(())
}
}
impl Context<'_> {
fn alloc_id(&mut self) -> ExprId {
self.ir_count += 1;
unsafe { ExprId::from_raw(self.ir_count - 1) }
}
fn add_dep(&mut self, from: ExprId, to: ExprId, count: &Cell<usize>) -> StackIdx {
if let Some(&idx) = self.graph.edge_weight(from, to) {
idx
} else {
let idx = count.get();
count.set(idx + 1);
let idx = unsafe { StackIdx::from_raw(idx) };
assert_ne!(from, to);
self.graph.add_edge(from, to, idx);
idx
}
}
}
impl BuiltinsContext for Context<'_> {}

97
evaluator/nixjit_context/src/resolve.rs
View File

@@ -1,11 +1,11 @@
use std::cell::RefCell;
use std::cell::{Cell, RefCell};
use derive_more::Unwrap;
use hashbrown::HashMap;
use nixjit_error::Result;
use nixjit_hir::Hir;
use nixjit_ir::{Const, ExprId, Param, StackIdx};
use nixjit_ir::{Const, ExprId, Param, PrimOpId, StackIdx};
use nixjit_lir::{Lir, LookupResult, Resolve, ResolveContext};
use replace_with::replace_with_and_return;
@@ -30,33 +30,22 @@ enum Ir {
Lir(Lir),
}
impl Ir {
unsafe fn unwrap_hir_unchecked(self) -> Hir {
if let Self::Hir(hir) = self {
hir
} else {
unsafe { core::hint::unreachable_unchecked() }
}
}
}
pub struct ResolveCtx<'ctx, 'bump> {
ctx: &'ctx mut Context<'bump>,
irs: Vec<Pin<'bump, RefCell<Ir>>>,
scopes: Vec<Scope<'ctx>>,
has_with: bool,
with_used: bool,
closures: Vec<(ExprId, Option<ExprId>, usize)>,
closures: Vec<(ExprId, Option<ExprId>, Cell<usize>)>,
current_expr: Option<ExprId>,
}
impl<'ctx, 'bump> ResolveCtx<'ctx, 'bump> {
pub fn new(ctx: &'ctx mut Context<'bump>) -> Self {
let ctx_mut = unsafe { &mut *(ctx as *mut Context) };
Self {
scopes: vec![
Scope::Builtins(&ctx.global_scope),
Scope::Repl(&ctx.repl_scope),
Scope::Builtins(unsafe { ctx.global_scope.as_ref() }),
Scope::Repl(unsafe { ctx.repl_scope.as_ref() }),
],
has_with: false,
with_used: false,
@@ -65,7 +54,7 @@ impl<'ctx, 'bump> ResolveCtx<'ctx, 'bump> {
.map(|hir| Ir::Hir(hir).into())
.map(|ir| Pin::new_in(ir, ctx.bump))
.collect(),
ctx: ctx_mut,
ctx,
closures: Vec::new(),
current_expr: None,
}
@@ -80,33 +69,10 @@ impl<'ctx, 'bump> ResolveCtx<'ctx, 'bump> {
}
}
fn get_ir_mut(&mut self, id: ExprId) -> &mut RefCell<Ir> {
let idx = unsafe { id.raw() } - self.ctx.lirs.len();
if cfg!(debug_assertions) {
self.irs.get_mut(idx).unwrap()
} else {
unsafe { self.irs.get_unchecked_mut(idx) }
}
}
fn add_dep(&mut self, from: ExprId, to: ExprId, count: &mut usize) -> StackIdx {
if let Some(&idx) = self.ctx.graph.edge_weight(from, to) {
idx
} else {
*count += 1;
let idx = unsafe { StackIdx::from_raw(*count - 1) };
assert_ne!(from, to);
self.ctx.graph.add_edge(from, to, idx);
idx
}
}
fn new_lir(&mut self, lir: Lir) -> ExprId {
self.irs.push(Pin::new_in(
RefCell::new(Ir::Lir(lir)),
self.ctx.bump,
));
unsafe { ExprId::from_raw(self.ctx.lirs.len() + self.irs.len() - 1) }
self.irs
.push(Pin::new_in(RefCell::new(Ir::Lir(lir)), self.ctx.bump));
self.ctx.alloc_id()
}
}
@@ -115,7 +81,7 @@ impl ResolveContext for ResolveCtx<'_, '_> {
let prev_expr = self.current_expr.replace(expr);
let result = unsafe {
let ctx = &mut *(self as *mut Self);
let ir = &mut self.get_ir_mut(expr);
let ir = self.get_ir(expr);
if !matches!(ir.try_borrow().as_deref(), Ok(Ir::Hir(_))) {
return Ok(());
}
@@ -126,7 +92,7 @@ impl ResolveContext for ResolveCtx<'_, '_> {
val: nixjit_value::Const::Null,
}))
},
|ir| match ir.unwrap_hir_unchecked().resolve(ctx) {
|ir| match ir.unwrap_hir().resolve(ctx) {
Ok(lir) => (Ok(()), Ir::Lir(lir)),
Err(err) => (
Err(err),
@@ -141,12 +107,8 @@ impl ResolveContext for ResolveCtx<'_, '_> {
result
}
fn resolve_call(&mut self, func: ExprId, arg: ExprId) -> Result<()> {
todo!()
}
fn resolve_root(mut self, expr: ExprId) -> Result<()> {
self.closures.push((expr, None, 0));
self.closures.push((expr, None, Cell::new(0)));
let ret = self.resolve(expr);
if ret.is_ok() {
self.ctx.lirs.extend(
@@ -167,17 +129,17 @@ impl ResolveContext for ResolveCtx<'_, '_> {
for scope in self.scopes.iter().rev() {
match scope {
Scope::Builtins(scope) => {
if let Some(&expr) = scope.get(&name) {
return LookupResult::Expr(expr);
if let Some(&primop) = scope.get(&name) {
return LookupResult::PrimOp(primop);
}
}
Scope::Let(scope) | &Scope::Repl(scope) => {
if let Some(&dep) = scope.get(name) {
let (expr, _, deps) = unsafe { &mut *(self as *mut Self) }
let (expr, _, deps) = self
.closures
.last_mut()
.last()
.unwrap();
let idx = self.add_dep(*expr, dep, deps);
let idx = self.ctx.add_dep(*expr, dep, deps);
return LookupResult::Stack(idx);
}
}
@@ -186,19 +148,18 @@ impl ResolveContext for ResolveCtx<'_, '_> {
// This is an outer function's parameter, treat as dependency
// We need to find the corresponding parameter expression to create dependency
// For now, we need to handle this case by creating a dependency to the parameter
let mut iter = unsafe { &mut *(self as *mut Self) }
.closures
.iter_mut()
.rev()
.take(closure_depth + 1)
.rev();
let mut iter = self.closures.iter().rev().take(closure_depth + 1).rev();
let Some((func, Some(arg), count)) = iter.next() else {
unreachable!()
};
let mut cur = self.add_dep(*func, *arg, count);
let mut cur = self.ctx.add_dep(*func, *arg, count);
for (func, _, count) in iter {
let idx = self.new_lir(Lir::StackRef(cur));
cur = self.add_dep(*func, idx, count);
self.irs.push(Pin::new_in(
RefCell::new(Ir::Lir(Lir::StackRef(cur))),
self.ctx.bump,
));
let idx = self.ctx.alloc_id();
cur = self.ctx.add_dep(*func, idx, count);
}
return LookupResult::Stack(cur);
}
@@ -215,12 +176,10 @@ impl ResolveContext for ResolveCtx<'_, '_> {
}
fn lookup_arg(&mut self) -> StackIdx {
let Some((func, Some(arg), count)) =
unsafe { &mut *(self as *mut Self) }.closures.last_mut()
else {
let Some((func, Some(arg), count)) = self.closures.last() else {
unreachable!()
};
self.add_dep(*func, *arg, count)
self.ctx.add_dep(*func, *arg, count)
}
fn new_func(&mut self, body: ExprId, param: Param) {
@@ -255,7 +214,7 @@ impl ResolveContext for ResolveCtx<'_, '_> {
f: impl FnOnce(&mut Self) -> T,
) -> T {
let arg = self.new_lir(Lir::Arg(nixjit_ir::Arg));
self.closures.push((func, Some(arg), 0));
self.closures.push((func, Some(arg), Cell::new(0)));
self.scopes.push(Scope::Arg(ident));
let res = f(self);
self.scopes.pop();