feat: stack var (WIP)

This commit is contained in:
2025-08-09 08:12:53 +08:00
parent fd182b6233
commit d8ad7fe904
36 changed files with 1521 additions and 1058 deletions

View File

@@ -1,89 +1,49 @@
//! The central evaluation context for the nixjit interpreter.
//!
//! This module defines the `Context` struct, which holds all the state
//! necessary for the evaluation of a Nix expression. It manages the
//! Intermediate Representations (IRs), scopes, evaluation stack, and
//! the Just-In-Time (JIT) compiler.
//!
//! The `Context` implements various traits (`DowngradeContext`, `ResolveContext`, etc.)
//! to provide the necessary services for each stage of the compilation and
//! evaluation pipeline.
use std::cell::{OnceCell, RefCell};
use std::rc::Rc;
use std::{marker::PhantomPinned, ops::Deref};
use derive_more::Unwrap;
use hashbrown::{HashMap, HashSet};
use bumpalo::{Bump, boxed::Box};
use hashbrown::HashMap;
use itertools::Itertools;
use petgraph::graph::{DiGraph, NodeIndex};
use petgraph::{
dot::{Config, Dot},
graphmap::DiGraphMap,
};
use nixjit_builtins::{
Builtins, BuiltinsContext,
builtins::{CONSTS_LEN, GLOBAL_LEN, SCOPED_LEN},
};
use nixjit_error::{Error, Result};
use nixjit_eval::{EvalContext, Evaluate, Value};
use nixjit_hir::{Downgrade, DowngradeContext, Hir};
use nixjit_ir::{ArgIdx, Const, ExprId, Param, PrimOp, PrimOpId};
use nixjit_lir::{Lir, LookupResult, Resolve, ResolveContext};
use nixjit_eval::{Args, EvalContext, Value};
use nixjit_hir::{DowngradeContext, Hir};
use nixjit_ir::{AttrSet, Const, ExprId, Param, PrimOpId, StackIdx};
use nixjit_lir::{Lir, ResolveContext};
use nixjit_jit::{JITCompiler, JITContext, JITFunc};
use replace_with::replace_with_and_return;
use crate::downgrade::DowngradeCtx;
use crate::eval::EvalCtx;
use crate::resolve::ResolveCtx;
/// Represents a lexical scope during name resolution.
enum Scope {
/// A `with` expression scope.
With,
/// A `let` binding scope, mapping variable names to their expression IDs.
Let(HashMap<String, ExprId>),
/// A function argument scope. `Some` holds the name of the argument set if present.
Arg(Option<String>),
mod downgrade;
mod eval;
mod resolve;
#[derive(Debug)]
struct Pin<'bump, T> {
ptr: Box<'bump, T>,
_marker: PhantomPinned,
}
/// Represents an expression at different stages of compilation.
#[derive(Debug, Unwrap)]
enum Ir {
/// An expression in the High-Level Intermediate Representation (HIR).
Hir(Hir),
/// An expression in the Low-Level Intermediate Representation (LIR).
Lir(Lir),
impl<T> Deref for Pin<'_, T> {
type Target = T;
fn deref(&self) -> &Self::Target {
self.ptr.as_ref()
}
}
impl Ir {
unsafe fn unwrap_hir_ref_unchecked(&self) -> &Hir {
if let Self::Hir(hir) = self {
hir
} else {
unsafe { core::hint::unreachable_unchecked() }
}
}
unsafe fn unwrap_hir_mut_unchecked(&mut self) -> &mut Hir {
#[cfg(debug_assertions)]
if let Self::Hir(hir) = self {
hir
} else {
unsafe { core::hint::unreachable_unchecked() }
}
#[cfg(not(debug_assertions))]
if let Self::Hir(hir) = self {
hir
} else {
unsafe { core::hint::unreachable_unchecked() }
}
}
unsafe fn unwrap_lir_ref_unchecked(&self) -> &Lir {
#[cfg(debug_assertions)]
if let Self::Lir(lir) = self {
lir
} else {
unsafe { core::hint::unreachable_unchecked() }
}
#[cfg(not(debug_assertions))]
if let Self::Lir(lir) = self {
lir
} else {
panic!()
impl<'bump, T> Pin<'bump, T> {
fn new_in(x: T, bump: &'bump Bump) -> Self {
Self {
ptr: Box::new_in(x, bump),
_marker: PhantomPinned,
}
}
}
@@ -92,103 +52,111 @@ impl Ir {
///
/// This struct orchestrates the entire Nix expression evaluation process,
/// from parsing and semantic analysis to interpretation and JIT compilation.
pub struct Context {
/// Arena for all expressions, which can be either HIR or LIR.
/// `RefCell` is used for interior mutability to allow on-demand resolution.
irs: Vec<RefCell<Ir>>,
/// Tracks whether an `ExprId` has been resolved from HIR to LIR.
resolved: Vec<bool>,
/// The stack of lexical scopes used for name resolution.
scopes: Vec<Scope>,
/// The number of arguments in the current function call scope.
args_count: usize,
/// A table of primitive operation implementations.
primops: Vec<fn(&mut Context, Vec<Value>) -> Result<Value>>,
pub struct Context<'bump> {
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>,
/// A dependency graph between expressions.
graph: DiGraph<ExprId, ()>,
/// Maps an `ExprId` to its corresponding `NodeIndex` in the dependency graph.
nodes: Vec<NodeIndex>,
graph: DiGraphMap<ExprId, StackIdx>,
/// The call stack for function evaluation, where each frame holds arguments.
stack: Vec<Vec<Value>>,
/// A stack of namespaces for `with` expressions during evaluation.
with_scopes: Vec<Rc<HashMap<String, Value>>>,
/// A table of primitive operation implementations.
primops: [(usize, fn(&mut Self, Args) -> Result<Value>); GLOBAL_LEN + SCOPED_LEN],
/// The Just-In-Time (JIT) compiler.
jit: JITCompiler<Self>,
/// A cache for JIT-compiled functions, indexed by `ExprId`.
compiled: Vec<OnceCell<JITFunc<Self>>>,
bump: &'bump Bump,
}
impl Default for Context {
fn default() -> Self {
impl<'bump> Context<'bump> {
pub fn new(bump: &'bump Bump) -> Self {
let Builtins {
consts,
global,
scoped,
} = Builtins::new();
let global_scope = Scope::Let(
consts
.iter()
.enumerate()
.map(|(id, (k, _))| (k.to_string(), unsafe { ExprId::from(id) }))
.chain(global.iter().enumerate().map(|(idx, (k, _, _))| {
(k.to_string(), unsafe { ExprId::from(idx + CONSTS_LEN) })
}))
.chain(core::iter::once(("builtins".to_string(), unsafe {
ExprId::from(CONSTS_LEN + GLOBAL_LEN + SCOPED_LEN)
})))
.collect(),
);
let primops = global
let global_scope = consts
.iter()
.map(|&(_, _, f)| f)
.chain(scoped.iter().map(|&(_, _, f)| f))
.collect();
let irs = consts
.into_iter()
.map(|(_, val)| Ir::Lir(Lir::Const(Const { val })))
.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) })),
)
.chain(core::iter::once(("builtins", unsafe {
ExprId::from_raw(CONSTS_LEN + GLOBAL_LEN + SCOPED_LEN)
})))
.collect();
let primops = global
.iter()
.map(|&(_, arity, f)| (arity, f))
.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) })))
.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, arity, _))| {
Ir::Lir(Lir::PrimOp(PrimOp {
name,
arity,
id: unsafe { PrimOpId::from(idx) },
}))
}),
)
.map(RefCell::new)
.map(|(idx, (name, _))| (name.to_string(), unsafe { ExprId::from_raw(idx) }))
.chain(global.into_iter().enumerate().map(|(idx, (name, ..))| {
(name.to_string(), unsafe {
ExprId::from_raw(idx + CONSTS_LEN)
})
}))
.chain(scoped.into_iter().enumerate().map(|(idx, (name, ..))| {
(name.to_string(), unsafe {
ExprId::from_raw(idx + CONSTS_LEN + GLOBAL_LEN)
})
}))
.chain(core::iter::once(("builtins".to_string(), unsafe {
ExprId::from_raw(CONSTS_LEN + GLOBAL_LEN + SCOPED_LEN + 1)
})))
.collect(),
..AttrSet::default()
})))
.chain(core::iter::once(Lir::Thunk(unsafe {
ExprId::from_raw(CONSTS_LEN + GLOBAL_LEN + SCOPED_LEN)
})))
.map(|lir| Pin::new_in(lir, bump))
.collect();
Self {
irs,
resolved: Vec::new(),
scopes: vec![global_scope],
args_count: 0,
primops,
hirs: Vec::new(),
lirs,
funcs: HashMap::new(),
graph: DiGraph::new(),
nodes: Vec::new(),
stack: Vec::new(),
with_scopes: Vec::new(),
global_scope,
repl_scope: HashMap::new(),
graph: DiGraphMap::new(),
jit: JITCompiler::new(),
compiled: Vec::new(),
primops,
bump,
}
}
}
impl Context {
/// Creates a new, default `Context`.
pub fn new() -> Self {
Self::default()
pub fn downgrade_ctx<'a>(&'a mut self) -> DowngradeCtx<'a, 'bump> {
DowngradeCtx::new(self)
}
pub fn resolve_ctx<'a>(&'a mut self) -> ResolveCtx<'a, 'bump> {
ResolveCtx::new(self)
}
pub fn eval_ctx<'a>(&'a mut self) -> EvalCtx<'a, 'bump> {
EvalCtx::new(self)
}
/// The main entry point for evaluating a Nix expression string.
///
/// This function performs the following steps:
@@ -196,227 +164,40 @@ impl Context {
/// 2. Downgrades the AST to the High-Level IR (HIR).
/// 3. Resolves the HIR to the Low-Level IR (LIR).
/// 4. Evaluates the LIR to produce a final `Value`.
pub fn eval(mut self, expr: &str) -> Result<nixjit_value::Value> {
pub fn eval(&mut self, expr: &str) -> Result<nixjit_value::Value> {
let root = rnix::Root::parse(expr);
if !root.errors().is_empty() {
return Err(Error::ParseError(
root.errors().iter().map(|err| err.to_string()).join(";"),
return Err(Error::parse_error(
root.errors().iter().map(|err| err.to_string()).join("; "),
));
}
let root = root.tree().expr().unwrap().downgrade(&mut self)?;
self.resolve(root)?;
Ok(EvalContext::eval(&mut self, root)?.to_public())
}
}
impl DowngradeContext for Context {
fn new_expr(&mut self, expr: Hir) -> ExprId {
let id = unsafe { ExprId::from(self.irs.len()) };
self.irs.push(Ir::Hir(expr).into());
self.nodes.push(self.graph.add_node(id));
self.resolved.push(false);
self.compiled.push(OnceCell::new());
id
}
fn with_expr<T>(&self, id: ExprId, f: impl FnOnce(&Hir, &Self) -> T) -> T {
unsafe {
let idx = id.raw();
f(&self.irs[idx].borrow().unwrap_hir_ref_unchecked(), self)
}
}
fn with_expr_mut<T>(&mut self, id: ExprId, f: impl FnOnce(&mut Hir, &mut Self) -> T) -> T {
unsafe {
let idx = id.raw();
let self_mut = &mut *(self as *mut Self);
f(
&mut self
.irs
.get_unchecked_mut(idx)
.borrow_mut()
.unwrap_hir_mut_unchecked(),
self_mut,
)
}
}
}
impl ResolveContext for Context {
fn lookup(&self, name: &str) -> LookupResult {
let mut arg_idx = 0;
let mut has_with = false;
for scope in self.scopes.iter().rev() {
match scope {
Scope::Let(scope) => {
if let Some(&expr) = scope.get(name) {
return LookupResult::Expr(expr);
}
}
Scope::Arg(ident) => {
if ident.as_deref() == Some(name) {
return LookupResult::Arg(unsafe { ArgIdx::from(arg_idx) });
}
arg_idx += 1;
}
Scope::With => has_with = true,
}
}
if has_with {
LookupResult::Unknown
} else {
LookupResult::NotFound
let root = self
.downgrade_ctx()
.downgrade_root(root.tree().expr().unwrap())?;
self.resolve_ctx().resolve_root(root)?;
println!(
"{:?}",
Dot::with_config(&self.graph, &[Config::EdgeNoLabel])
);
for (idx, ir) in self.lirs.iter().enumerate() {
println!("{:?} {:#?}", unsafe { ExprId::from_raw(idx) }, &ir);
}
Ok(self.eval_ctx().eval_root(root)?.to_public())
}
fn new_dep(&mut self, expr: ExprId, dep: ExprId) {
unsafe {
let expr = expr.raw();
let dep = dep.raw();
let expr = *self.nodes.get_unchecked(expr);
let dep = *self.nodes.get_unchecked(dep);
self.graph.add_edge(expr, dep, ());
}
}
fn resolve(&mut self, expr: ExprId) -> Result<()> {
unsafe {
let idx = expr.raw();
let self_mut = &mut *(self as *mut Self);
replace_with_and_return(
&mut *self.irs.get_unchecked(idx).borrow_mut(),
|| {
Ir::Hir(Hir::Const(Const {
val: nixjit_value::Const::Null,
}))
},
|ir| {
let Ir::Hir(hir) = ir else {
return (Ok(()), ir);
};
match hir.resolve(self_mut) {
Ok(lir) => (Ok(()), Ir::Lir(lir)),
Err(err) => (
Err(err),
Ir::Hir(Hir::Const(Const {
val: nixjit_value::Const::Null,
})),
),
}
},
)?;
pub fn add_binding(&mut self, ident: &str, expr: &str) -> Result<()> {
let root = rnix::Root::parse(expr);
if !root.errors().is_empty() {
return Err(Error::parse_error(
root.errors().iter().map(|err| err.to_string()).join("; "),
));
}
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);
Ok(())
}
fn new_func(&mut self, body: ExprId, param: Param) {
self.funcs.insert(body, param);
}
fn with_let_env<'a, T>(
&mut self,
bindings: impl Iterator<Item = (&'a String, &'a ExprId)>,
f: impl FnOnce(&mut Self) -> T,
) -> T {
let mut scope = HashMap::new();
for (name, expr) in bindings {
scope.insert(name.clone(), *expr);
}
self.scopes.push(Scope::Let(scope));
let res = f(self);
self.scopes.pop();
res
}
fn with_with_env<T>(&mut self, f: impl FnOnce(&mut Self) -> T) -> (bool, T) {
self.scopes.push(Scope::With);
let res = f(self);
self.scopes.pop();
(true, res)
}
fn with_param_env<T>(&mut self, ident: Option<String>, f: impl FnOnce(&mut Self) -> T) -> T {
self.scopes.push(Scope::Arg(ident));
self.args_count += 1;
let res = f(self);
self.args_count -= 1;
self.scopes.pop();
res
}
}
impl EvalContext for Context {
fn eval(&mut self, expr: ExprId) -> Result<nixjit_eval::Value> {
let idx = unsafe { expr.raw() };
let lir = unsafe {
&*(self
.irs
.get_unchecked(idx)
.borrow()
.unwrap_lir_ref_unchecked() as *const Lir)
};
println!("{:#?}", self.irs);
lir.eval(self)
}
fn pop_frame(&mut self) -> Vec<nixjit_eval::Value> {
self.stack.pop().unwrap()
}
fn lookup_stack<'a>(&'a self, offoset: usize) -> &'a Value {
todo!()
}
fn lookup_with<'a>(&'a self, ident: &str) -> Option<&'a nixjit_eval::Value> {
for scope in self.with_scopes.iter().rev() {
if let Some(val) = scope.get(ident) {
return Some(val);
}
}
None
}
fn lookup_arg<'a>(&'a self, idx: ArgIdx) -> &'a Value {
unsafe {
let values = self.stack.last().unwrap_unchecked();
dbg!(values, idx);
&values[values.len() - idx.raw() - 1]
}
}
fn with_with_env<T>(
&mut self,
namespace: std::rc::Rc<HashMap<String, nixjit_eval::Value>>,
f: impl FnOnce(&mut Self) -> T,
) -> T {
self.with_scopes.push(namespace);
let res = f(self);
self.with_scopes.pop();
res
}
fn with_args_env<T>(
&mut self,
args: Vec<nixjit_eval::Value>,
f: impl FnOnce(&mut Self) -> T,
) -> (Vec<Value>, T) {
self.stack.push(args);
let res = f(self);
let frame = self.stack.pop().unwrap();
(frame, res)
}
fn call_primop(&mut self, id: nixjit_ir::PrimOpId, args: Vec<Value>) -> Result<Value> {
unsafe { (self.primops.get_unchecked(id.raw()))(self, args) }
}
}
impl JITContext for Context {
fn enter_with(&mut self, namespace: std::rc::Rc<HashMap<String, nixjit_eval::Value>>) {
self.with_scopes.push(namespace);
}
fn exit_with(&mut self) {
self.with_scopes.pop();
}
}
impl BuiltinsContext for Context {}
impl BuiltinsContext for Context<'_> {}