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

@@ -0,0 +1,61 @@
use std::cell::RefCell;
use nixjit_error::Result;
use nixjit_hir::{Downgrade, DowngradeContext, Hir};
use nixjit_ir::ExprId;
use super::Context;
pub struct DowngradeCtx<'ctx, 'bump> {
ctx: &'ctx mut Context<'bump>,
irs: Vec<RefCell<Hir>>,
}
impl<'ctx, 'bump> DowngradeCtx<'ctx, 'bump> {
pub fn new(ctx: &'ctx mut Context<'bump>) -> Self {
Self {
ctx,
irs: Vec::new(),
}
}
}
impl DowngradeCtx<'_, '_> {
fn get_ir(&self, id: ExprId) -> &RefCell<Hir> {
let idx = unsafe { id.raw() } - self.ctx.lirs.len() - self.ctx.hirs.len();
if cfg!(debug_assertions) {
self.irs.get(idx).unwrap()
} else {
unsafe { self.irs.get_unchecked(idx) }
}
}
}
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) }
}
fn with_expr_mut<T>(&mut self, id: ExprId, f: impl FnOnce(&mut Hir, &mut Self) -> T) -> T {
unsafe {
let self_mut = &mut *(self as *mut Self);
f(&mut self.get_ir(id).borrow_mut(), self_mut)
}
}
fn downgrade_root(mut self, root: rnix::ast::Expr) -> Result<ExprId> {
let id = root.downgrade(&mut self)?;
self.ctx
.hirs
.extend(self.irs.into_iter().map(RefCell::into_inner));
for (idx, ir) in self.ctx.hirs.iter().enumerate() {
println!(
"{:?} {:#?}",
unsafe { ExprId::from_raw(idx + self.ctx.lirs.len()) },
&ir
);
}
Ok(id)
}
}

View File

@@ -0,0 +1,144 @@
use std::rc::Rc;
use hashbrown::HashMap;
use itertools::Itertools;
use nixjit_error::Result;
use nixjit_eval::{Args, EvalContext, Evaluate, StackFrame, Value};
use nixjit_ir::ExprId;
use nixjit_jit::JITContext;
use nixjit_lir::Lir;
use super::Context;
pub struct EvalCtx<'ctx, 'bump> {
ctx: &'ctx mut Context<'bump>,
stack: Vec<StackFrame>,
with_scopes: Vec<Rc<HashMap<String, Value>>>,
}
impl<'ctx, 'bump> EvalCtx<'ctx, 'bump> {
pub fn new(ctx: &'ctx mut Context<'bump>) -> Self {
Self {
ctx,
stack: Vec::new(),
with_scopes: Vec::new(),
}
}
fn eval_deps(&mut self, expr: ExprId, arg: Option<Value>) -> Result<()> {
let deps = self
.ctx
.graph
.edges(expr)
.sorted_by_key(|(.., idx)| **idx)
.map(|(_, dep, idx)| (dep, *idx))
.collect_vec();
let mut frame = (0..deps.len())
.map(|_| Value::Blackhole)
.collect::<StackFrame>();
dbg!(&deps, &self.stack);
for (dep, idx) in deps {
unsafe {
if matches!(
&**self.ctx.lirs.get_unchecked(dep.raw()),
Lir::Arg(_)
) {
*frame.get_unchecked_mut(idx.raw()) = arg.as_ref().unwrap().clone();
continue;
}
}
let dep = self.eval(dep)?;
unsafe {
*frame.get_unchecked_mut(idx.raw()) = dep;
}
}
*self.stack.last_mut().unwrap() = frame;
dbg!(&self.stack);
Ok(())
}
}
impl EvalContext for EvalCtx<'_, '_> {
fn eval_root(mut self, expr: ExprId) -> Result<Value> {
self.stack.push(StackFrame::new());
self.eval_deps(expr, None)?;
self.eval(expr)
}
fn eval(&mut self, expr: ExprId) -> Result<Value> {
let idx = unsafe { expr.raw() };
let lir = unsafe { &*(&**self.ctx.lirs.get_unchecked(idx) as *const Lir) };
lir.eval(self)
}
fn call(&mut self, func: ExprId, arg: Option<Value>, frame: StackFrame) -> Result<Value> {
self.stack.push(frame);
if let Err(err) = self.eval_deps(func, arg) {
self.stack.pop();
return Err(err)
}
let ret = self.eval(func);
self.stack.pop();
ret
}
fn lookup_with<'a>(&'a self, ident: &str) -> Option<&'a Value> {
for scope in self.with_scopes.iter().rev() {
if let Some(val) = scope.get(ident) {
return Some(val);
}
}
None
}
fn lookup_stack(&self, idx: nixjit_ir::StackIdx) -> &Value {
if cfg!(debug_assertions) {
self.stack
.last()
.unwrap()
.get(unsafe { idx.raw() })
.unwrap()
} else {
unsafe {
self.stack
.last()
.unwrap_unchecked()
.get_unchecked(idx.raw())
}
}
}
fn capture_stack(&self) -> &StackFrame {
self.stack.last().unwrap()
}
fn call_primop(&mut self, id: nixjit_ir::PrimOpId, args: Args) -> Result<Value> {
unsafe { (self.ctx.primops.get_unchecked(id.raw()).1)(self.ctx, args) }
}
fn get_primop_arity(&self, id: nixjit_ir::PrimOpId) -> usize {
unsafe { self.ctx.primops.get_unchecked(id.raw()).0 }
}
fn with_with_env<T>(
&mut self,
namespace: Rc<HashMap<String, Value>>,
f: impl FnOnce(&mut Self) -> T,
) -> T {
self.with_scopes.push(namespace);
let res = f(self);
self.with_scopes.pop();
res
}
}
impl JITContext for EvalCtx<'_, '_> {
fn enter_with(&mut self, namespace: Rc<HashMap<String, Value>>) {
self.with_scopes.push(namespace);
}
fn exit_with(&mut self) {
self.with_scopes.pop();
}
}

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<'_> {}

View File

@@ -0,0 +1,263 @@
use std::cell::RefCell;
use std::pin::Pin;
use bumpalo::boxed::Box;
use derive_more::Unwrap;
use hashbrown::HashMap;
use nixjit_error::Result;
use nixjit_hir::Hir;
use nixjit_ir::{Const, ExprId, Param, StackIdx};
use nixjit_lir::{Lir, LookupResult, Resolve, ResolveContext};
use replace_with::replace_with_and_return;
use super::Context;
#[derive(Clone)]
enum Scope<'ctx> {
/// 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>),
Builtins(&'ctx HashMap<&'static str, ExprId>),
Repl(&'ctx HashMap<String, ExprId>)
}
/// 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 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<Box<'bump, RefCell<Ir>>>>,
scopes: Vec<Scope<'ctx>>,
has_with: bool,
with_used: bool,
closures: Vec<(ExprId, Option<ExprId>, 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)
],
has_with: false,
with_used: false,
irs: core::mem::take(&mut ctx.hirs)
.into_iter()
.map(|hir| Ir::Hir(hir).into())
.map(|ir| Box::new_in(ir, ctx.bump))
.map(Pin::new)
.collect(),
ctx: ctx_mut,
closures: Vec::new(),
current_expr: None,
}
}
fn get_ir(&self, id: ExprId) -> &RefCell<Ir> {
let idx = unsafe { id.raw() } - self.ctx.lirs.len();
if cfg!(debug_assertions) {
self.irs.get(idx).unwrap()
} else {
unsafe { self.irs.get_unchecked(idx) }
}
}
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(Box::new_in(
RefCell::new(Ir::Lir(lir)),
self.ctx.bump,
)));
unsafe { ExprId::from_raw(self.ctx.lirs.len() + self.irs.len() - 1) }
}
}
impl ResolveContext for ResolveCtx<'_, '_> {
fn resolve(&mut self, expr: ExprId) -> Result<()> {
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);
if !matches!(ir.try_borrow().as_deref(), Ok(Ir::Hir(_))) {
return Ok(());
}
replace_with_and_return(
&mut *ir.borrow_mut(),
|| {
Ir::Hir(Hir::Const(Const {
val: nixjit_value::Const::Null,
}))
},
|ir| match ir.unwrap_hir_unchecked().resolve(ctx) {
Ok(lir) => (Ok(()), Ir::Lir(lir)),
Err(err) => (
Err(err),
Ir::Hir(Hir::Const(Const {
val: nixjit_value::Const::Null,
})),
),
},
)
};
self.current_expr = prev_expr;
result
}
fn resolve_root(mut self, expr: ExprId) -> Result<()> {
self.closures.push((expr, None, 0));
let ret = self.resolve(expr);
if ret.is_ok() {
self.ctx.lirs.extend(
self.irs
.into_iter()
.map(|pin| unsafe { core::mem::transmute::<Pin<_>, Box<_>>(pin) })
.map(Box::into_inner)
.map(RefCell::into_inner)
.map(Ir::unwrap_lir)
.map(|lir| crate::Pin::new_in(lir, self.ctx.bump))
);
}
ret
}
fn lookup(&mut self, name: &str) -> LookupResult {
let mut closure_depth = 0;
// Then search from outer to inner scopes for dependencies
for scope in self.scopes.iter().rev() {
match scope {
Scope::Builtins(scope) => {
if let Some(&expr) = scope.get(&name) {
return LookupResult::Expr(expr);
}
}
Scope::Let(scope) | &Scope::Repl(scope) => {
if let Some(&dep) = scope.get(name) {
let (expr, _, deps) = unsafe { &mut *(self as *mut Self) }
.closures
.last_mut()
.unwrap();
let idx = self.add_dep(*expr, dep, deps);
return LookupResult::Stack(idx);
}
}
Scope::Arg(ident) => {
if ident.as_deref() == Some(name) {
// 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 Some((func, Some(arg), count)) = iter.next() else {
unreachable!()
};
let mut cur = self.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);
}
return LookupResult::Stack(cur);
}
closure_depth += 1;
}
}
}
if self.has_with {
self.with_used = true;
LookupResult::Unknown
} else {
LookupResult::NotFound
}
}
fn lookup_arg(&mut self) -> StackIdx {
let Some((func, Some(arg), count)) = unsafe { &mut *(self as *mut Self) }.closures.last_mut() else {
unreachable!()
};
self.add_dep(*func, *arg, count)
}
fn new_func(&mut self, body: ExprId, param: Param) {
self.ctx.funcs.insert(body, param);
}
fn with_let_env<T>(
&mut self,
bindings: HashMap<String, ExprId>,
f: impl FnOnce(&mut Self) -> T,
) -> T {
self.scopes.push(Scope::Let(bindings));
let res = f(self);
self.scopes.pop();
res
}
fn with_with_env<T>(&mut self, f: impl FnOnce(&mut Self) -> T) -> (bool, T) {
let has_with = self.has_with;
let with_used = self.with_used;
self.has_with = true;
self.with_used = false;
let res = f(self);
self.has_with = has_with;
(core::mem::replace(&mut self.with_used, with_used), res)
}
fn with_param_env<T>(
&mut self,
func: ExprId,
ident: Option<String>,
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.scopes.push(Scope::Arg(ident));
let res = f(self);
self.scopes.pop();
self.closures.pop();
res
}
}