fix: WIP

evaluator/nixjit_context/src/lib.rs

@@ -1,30 +1,50 @@
 use core::mem::MaybeUninit;
 use std::cell::{OnceCell, RefCell};
+use std::rc::Rc;
 
 use hashbrown::{HashMap, HashSet};
-use itertools::Itertools;
+use petgraph::algo::toposort;
+use petgraph::graph::{DiGraph, NodeIndex};
 
-use nixjit_error::Result;
-use nixjit_eval::EvalContext;
+use nixjit_error::{Error, Result};
+use nixjit_eval::{EvalContext, Evaluate, Value};
 use nixjit_hir::{DowngradeContext, Hir};
-use nixjit_ir::{ExprId, Param};
-use nixjit_lir::{Lir, ResolveContext};
-use nixjit_value::Value;
+use nixjit_ir::{ArgIdx, ExprId, Param};
+use nixjit_lir::{Lir, LookupResult, Resolve, ResolveContext};
 
 use nixjit_jit::{JITCompiler, JITContext, JITFunc};
 
+#[derive(Debug)]
+struct Frame {
+    values: Vec<Value<Context>>,
+    left: usize,
+}
+
 #[derive(Default)]
 pub struct Context {
     hirs: Vec<RefCell<Hir>>,
     lirs: Vec<MaybeUninit<Lir>>,
+    resolved: Vec<bool>,
+    scopes: Vec<HashMap<String, LookupResult>>,
+    funcs: HashMap<ExprId, Param>,
+    graph: DiGraph<ExprId, ()>,
+    nodes: Vec<NodeIndex>,
+
+    stack: Vec<Frame>,
+    with_scopes: Vec<Rc<HashMap<String, Value<Self>>>>,
+
     jit: JITCompiler<Self>,
     compiled: Vec<OnceCell<JITFunc<Self>>>,
 }
 
 impl Drop for Context {
     fn drop(&mut self) {
-        unsafe {
-            self.lirs.iter_mut().for_each(|lir| lir.assume_init_drop());
+        for (i, lir) in self.lirs.iter_mut().enumerate() {
+            if self.resolved[i] {
+                unsafe {
+                    lir.assume_init_drop();
+                }
+            }
         }
     }
 }
@@ -37,20 +57,22 @@ impl Context {
 
 impl DowngradeContext for Context {
     fn new_expr(&mut self, expr: Hir) -> ExprId {
-        let id = unsafe { core::mem::transmute(self.hirs.len() + self.lirs.len()) };
+        let id = ExprId::from(self.hirs.len());
         self.hirs.push(expr.into());
+        self.lirs.push(MaybeUninit::uninit());
+        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: usize = core::mem::transmute(id);
-            f(&self.hirs.get_unchecked(idx).borrow(), self)
-        }
+        let idx = usize::from(id);
+        f(&self.hirs[idx].borrow(), self)
     }
     fn with_expr_mut<T>(&mut self, id: ExprId, f: impl FnOnce(&mut Hir, &mut Self) -> T) -> T {
+        let idx = usize::from(id);
         unsafe {
-            let self_mut = &mut *(self as *mut Context);
-            let idx: usize = core::mem::transmute(id);
+            let self_mut = &mut *(self as *mut Self);
             f(&mut self.hirs.get_unchecked_mut(idx).borrow_mut(), self_mut)
         }
     }
@@ -58,88 +80,163 @@ impl DowngradeContext for Context {
 
 impl ResolveContext for Context {
     fn lookup(&self, name: &str) -> nixjit_lir::LookupResult {
-        todo!()
+        for scope in self.scopes.iter().rev() {
+            if let Some(val) = scope.get(name) {
+                return val.clone();
+            }
+        }
+        LookupResult::NotFound
     }
+
     fn new_dep(&mut self, expr: ExprId, dep: ExprId) {
-        todo!()
+        let expr = *self.nodes.get(usize::from(expr)).unwrap();
+        let dep = *self.nodes.get(usize::from(dep)).unwrap();
+        self.graph.add_edge(expr, dep, ());
     }
+
     fn resolve(&mut self, expr: ExprId) -> Result<()> {
-        todo!()
+        let idx = usize::from(expr);
+        if self.resolved[idx] {
+            return Ok(());
+        }
+
+        let hir = self.hirs[idx].replace(nixjit_hir::Hir::Const(nixjit_ir::Const::from(false)));
+        let lir = hir.resolve(self)?;
+        self.lirs[idx].write(lir);
+        self.resolved[idx] = true;
+        Ok(())
     }
+
     fn new_func(&mut self, body: ExprId, param: Param) {
-        todo!()
+        self.funcs.insert(body, param);
     }
+
     fn with_let_env<'a, T>(
         &mut self,
         bindings: impl IntoIterator<Item = (&'a String, &'a ExprId)>,
         f: impl FnOnce(&mut Self) -> T,
     ) -> T {
-        todo!()
+        let mut scope = HashMap::new();
+        for (name, expr) in bindings {
+            scope.insert(name.clone(), LookupResult::Expr(*expr));
+        }
+        self.scopes.push(scope);
+        let res = f(self);
+        self.scopes.pop();
+        res
     }
+
     fn with_with_env<T>(&mut self, f: impl FnOnce(&mut Self) -> T) -> (bool, T) {
-        todo!()
+        self.scopes.push(HashMap::new());
+        let res = f(self);
+        self.scopes.pop();
+        (true, res)
     }
+
     fn with_param_env<'a, T>(
         &mut self,
         ident: Option<&'a str>,
         f: impl FnOnce(&mut Self) -> T,
     ) -> T {
-        todo!()
+        let mut scope = HashMap::new();
+        if let Some(ident) = ident {
+            scope.insert(ident.to_string(), LookupResult::Arg(ArgIdx::from(0)));
+        }
+        self.scopes.push(scope);
+        let res = f(self);
+        self.scopes.pop();
+        res
     }
 }
 
 impl EvalContext for Context {
-    fn eval(&mut self, expr: ExprId) -> Result<nixjit_eval::Value<Self>>
-    where
-        Self: Sized,
-    {
-        todo!()
+    fn eval(&mut self, expr: ExprId) -> Result<nixjit_eval::Value<Self>> {
+        let idx = usize::from(expr);
+        let lir = unsafe { &*(self.lirs[idx].assume_init_ref() as *const Lir) };
+        lir.eval(self)
     }
-    fn pop_frame(&mut self) -> Vec<nixjit_eval::Value<Self>>
-    where
-        Self: Sized,
-    {
-        todo!()
+
+    fn pop_frame(&mut self) -> Vec<nixjit_eval::Value<Self>> {
+        self.stack.pop().unwrap().values
     }
-    fn lookup_with<'a>(&'a self, ident: &str) -> Option<&'a nixjit_eval::Value<Self>>
-    where
-        Self: Sized,
-    {
-        todo!()
+
+    fn lookup_with<'a>(&'a self, ident: &str) -> Option<&'a nixjit_eval::Value<Self>> {
+        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, offset: usize) -> Option<&'a Value<Self>> {
+        self.stack.last().and_then(|frame| frame.values.get(offset))
+    }
+
     fn with_with_env<T>(
         &mut self,
         namespace: std::rc::Rc<HashMap<String, nixjit_eval::Value<Self>>>,
         f: impl FnOnce(&mut Self) -> T,
-    ) -> T
-    where
-        Self: Sized,
-    {
-        todo!()
+    ) -> 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<Self>>,
         f: impl FnOnce(&mut Self) -> T,
-    ) -> (Vec<nixjit_eval::Value<Self>>, T)
-    where
-        Self: Sized,
-    {
-        todo!()
+    ) -> (Vec<nixjit_eval::Value<Self>>, T) {
+        self.stack.push(Frame {
+            left: args.len(),
+            values: args,
+        });
+        let res = f(self);
+        (self.stack.pop().unwrap().values, res)
+    }
+
+    fn consume_arg(&mut self, func: ExprId) -> Result<bool> {
+        let Some(frame) = self.stack.last_mut() else {
+            return Ok(false);
+        };
+        if frame.left == 0 {
+            return Ok(false);
+        }
+        frame.left -= 1;
+        let param = self.funcs.get(&func).unwrap();
+        if let Some(required) = &param.required {
+            let attrs = frame.values[frame.values.len() - frame.left - 1]
+                .as_ref()
+                .try_unwrap_attr_set()
+                .map_err(|_| Error::EvalError(format!("expected a set but found ...")))?;
+            if required.iter().any(|attr| attrs.get(attr).is_none())
+                || param.allowed.as_ref().map_or(false, |allowed| {
+                    attrs.iter().any(|(attr, _)| allowed.get(attr).is_none())
+                })
+            {
+                return Err(Error::EvalError(format!("TODO")));
+            }
+        }
+        Ok(true)
     }
 }
 
 impl JITContext for Context {
     fn lookup_arg(&self, offset: usize) -> &nixjit_eval::Value<Self> {
-        todo!()
+        &self.stack.last().unwrap().values[offset]
     }
+
     fn lookup_stack(&self, offset: usize) -> &nixjit_eval::Value<Self> {
         todo!()
     }
+
     fn enter_with(&mut self, namespace: std::rc::Rc<HashMap<String, nixjit_eval::Value<Self>>>) {
-        todo!()
+        self.with_scopes.push(namespace);
     }
+
     fn exit_with(&mut self) {
-        todo!()
+        self.with_scopes.pop();
     }
 }