feat: a lot

evaluator/nixjit_context/src/lib.rs

@@ -1,50 +1,162 @@
-use core::mem::MaybeUninit;
 use std::cell::{OnceCell, RefCell};
 use std::rc::Rc;
 
+use derive_more::Unwrap;
 use hashbrown::{HashMap, HashSet};
-use petgraph::algo::toposort;
+use itertools::Itertools;
 use petgraph::graph::{DiGraph, NodeIndex};
 
+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::{DowngradeContext, Hir};
-use nixjit_ir::{ArgIdx, ExprId, Param};
+use nixjit_hir::{Downgrade, DowngradeContext, Hir};
+use nixjit_ir::{ArgIdx, Const, ExprId, Param, PrimOp, PrimOpId};
 use nixjit_lir::{Lir, LookupResult, Resolve, ResolveContext};
 
 use nixjit_jit::{JITCompiler, JITContext, JITFunc};
+use replace_with::replace_with_and_return;
 
-#[derive(Debug)]
-struct Frame {
-    values: Vec<Value<Context>>,
-    left: usize,
+enum Scope {
+    With,
+    Let(HashMap<String, ExprId>),
+    Arg(Option<String>),
+}
+
+#[derive(Debug, Unwrap)]
+enum Ir {
+    Hir(Hir),
+    Lir(Lir),
+}
+
+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_hir_unchecked(self) -> Hir {
+        if cfg!(debug_assertions) {
+            self.unwrap_hir()
+        } else 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!()
+        }
+    }
 }
 
-#[derive(Default)]
 pub struct Context {
-    hirs: Vec<RefCell<Hir>>,
-    lirs: Vec<MaybeUninit<Lir>>,
+    irs: Vec<RefCell<Ir>>,
     resolved: Vec<bool>,
-    scopes: Vec<HashMap<String, LookupResult>>,
+    scopes: Vec<Scope>,
+    args_count: usize,
+    primops: Vec<fn(&mut Context, Vec<Value>) -> Result<Value>>,
     funcs: HashMap<ExprId, Param>,
     graph: DiGraph<ExprId, ()>,
     nodes: Vec<NodeIndex>,
 
-    stack: Vec<Frame>,
-    with_scopes: Vec<Rc<HashMap<String, Value<Self>>>>,
+    stack: Vec<Vec<Value>>,
+    with_scopes: Vec<Rc<HashMap<String, Value>>>,
 
     jit: JITCompiler<Self>,
     compiled: Vec<OnceCell<JITFunc<Self>>>,
 }
 
-impl Drop for Context {
-    fn drop(&mut self) {
-        for (i, lir) in self.lirs.iter_mut().enumerate() {
-            if self.resolved[i] {
-                unsafe {
-                    lir.assume_init_drop();
-                }
-            }
+impl Default for Context {
+    fn default() -> 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
+            .iter()
+            .map(|&(_, _, f)| f)
+            .chain(scoped.iter().map(|&(_, _, f)| f))
+            .collect();
+        let irs = consts
+            .into_iter()
+            .map(|(_, val)| Ir::Lir(Lir::Const(Const { val })))
+            .chain(
+                global
+                    .into_iter()
+                    .enumerate()
+                    .map(|(idx, (name, arity, _))| {
+                        Ir::Lir(Lir::PrimOp(PrimOp {
+                            name,
+                            arity,
+                            id: unsafe { PrimOpId::from(idx) },
+                        }))
+                    }),
+            )
+            .map(RefCell::new)
+            .collect();
+        Self {
+            irs,
+            resolved: Vec::new(),
+            scopes: vec![global_scope],
+            args_count: 0,
+            primops,
+            funcs: HashMap::new(),
+            graph: DiGraph::new(),
+            nodes: Vec::new(),
+
+            stack: Vec::new(),
+            with_scopes: Vec::new(),
+
+            jit: JITCompiler::new(),
+            compiled: Vec::new(),
         }
     }
 }
@@ -53,114 +165,166 @@ impl Context {
     pub fn new() -> Self {
         Self::default()
     }
+
+    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(";"),
+            ));
+        }
+        let root = root.tree().expr().unwrap().downgrade(&mut self)?;
+        self.resolve(&root)?;
+        Ok(EvalContext::eval(&mut self, &root)?.to_public(&mut HashSet::new()))
+    }
 }
 
 impl DowngradeContext for Context {
     fn new_expr(&mut self, expr: Hir) -> ExprId {
-        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));
+        let id = unsafe { ExprId::from(self.irs.len()) };
+        self.irs.push(Ir::Hir(expr).into());
+        self.nodes.push(self.graph.add_node(unsafe { id.clone() }));
         self.resolved.push(false);
         self.compiled.push(OnceCell::new());
         id
     }
     fn with_expr<T>(&self, id: ExprId, f: impl FnOnce(&Hir, &Self) -> T) -> T {
-        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 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.clone().raw();
             let self_mut = &mut *(self as *mut Self);
-            f(&mut self.hirs.get_unchecked_mut(idx).borrow_mut(), self_mut)
+            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) -> nixjit_lir::LookupResult {
+    fn lookup(&self, name: &str) -> LookupResult {
+        let mut arg_idx = 0;
         for scope in self.scopes.iter().rev() {
-            if let Some(val) = scope.get(name) {
-                return val.clone();
+            match scope {
+                Scope::Let(scope) => {
+                    if let Some(expr) = scope.get(name) {
+                        return LookupResult::Expr(unsafe { expr.clone() });
+                    }
+                }
+                Scope::Arg(ident) => {
+                    if ident.as_deref() == Some(name) {
+                        return LookupResult::Arg(unsafe { ArgIdx::from(arg_idx) });
+                    }
+                    arg_idx += 1;
+                }
+                Scope::With => return LookupResult::Unknown,
             }
         }
         LookupResult::NotFound
     }
 
-    fn new_dep(&mut self, expr: ExprId, dep: ExprId) {
-        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 new_dep(&mut self, expr: &ExprId, dep: ExprId) {
+        unsafe {
+            let expr = expr.clone().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<()> {
-        let idx = usize::from(expr);
-        if self.resolved[idx] {
-            return Ok(());
+    fn resolve(&mut self, expr: &ExprId) -> Result<()> {
+        unsafe {
+            let idx = expr.clone().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 hir = ir.unwrap_hir_unchecked();
+                    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,
+                            })),
+                        ),
+                    }
+                },
+            )?;
         }
-
-        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) {
-        self.funcs.insert(body, param);
+    fn new_func(&mut self, body: &ExprId, param: Param) {
+        self.funcs.insert(unsafe { body.clone() }, param);
     }
 
     fn with_let_env<'a, T>(
         &mut self,
-        bindings: impl IntoIterator<Item = (&'a String, &'a ExprId)>,
+        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(), LookupResult::Expr(*expr));
+            scope.insert(name.clone(), unsafe { expr.clone() });
         }
-        self.scopes.push(scope);
+        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(HashMap::new());
+        self.scopes.push(Scope::With);
         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 {
-        let mut scope = HashMap::new();
-        if let Some(ident) = ident {
-            scope.insert(ident.to_string(), LookupResult::Arg(ArgIdx::from(0)));
-        }
-        self.scopes.push(scope);
+    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<Self>> {
-        let idx = usize::from(expr);
-        let lir = unsafe { &*(self.lirs[idx].assume_init_ref() as *const Lir) };
+    fn eval(&mut self, expr: &ExprId) -> Result<nixjit_eval::Value> {
+        let idx = unsafe { expr.clone().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>> {
-        self.stack.pop().unwrap().values
+    fn pop_frame(&mut self) -> Vec<nixjit_eval::Value> {
+        self.stack.pop().unwrap()
     }
 
-    fn lookup_with<'a>(&'a self, ident: &str) -> Option<&'a nixjit_eval::Value<Self>> {
+    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);
@@ -169,13 +333,17 @@ impl EvalContext for Context {
         None
     }
 
-    fn lookup_arg<'a>(&'a self, offset: usize) -> Option<&'a Value<Self>> {
-        self.stack.last().and_then(|frame| frame.values.get(offset))
+    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<Self>>>,
+        namespace: std::rc::Rc<HashMap<String, nixjit_eval::Value>>,
         f: impl FnOnce(&mut Self) -> T,
     ) -> T {
         self.with_scopes.push(namespace);
@@ -186,53 +354,33 @@ impl EvalContext for Context {
 
     fn with_args_env<T>(
         &mut self,
-        args: Vec<nixjit_eval::Value<Self>>,
+        args: Vec<nixjit_eval::Value>,
         f: impl FnOnce(&mut Self) -> T,
-    ) -> (Vec<nixjit_eval::Value<Self>>, T) {
-        self.stack.push(Frame {
-            left: args.len(),
-            values: args,
-        });
+    ) -> (Vec<Value>, T) {
+        self.stack.push(args);
         let res = f(self);
-        (self.stack.pop().unwrap().values, res)
+        let frame = self.stack.pop().unwrap();
+        (frame, 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);
+    fn call_primop(&mut self, id: nixjit_ir::PrimOpId, args: Vec<Value>) -> Result<Value> {
+        unsafe {
+            (self.primops.get_unchecked(id.raw()))(self, args)
         }
-        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> {
-        &self.stack.last().unwrap().values[offset]
+    fn lookup_arg(&self, offset: usize) -> &nixjit_eval::Value {
+        let values = self.stack.last().unwrap();
+        &values[values.len() - offset - 1]
     }
 
-    fn lookup_stack(&self, offset: usize) -> &nixjit_eval::Value<Self> {
+    fn lookup_stack(&self, offset: usize) -> &nixjit_eval::Value {
         todo!()
     }
 
-    fn enter_with(&mut self, namespace: std::rc::Rc<HashMap<String, nixjit_eval::Value<Self>>>) {
+    fn enter_with(&mut self, namespace: std::rc::Rc<HashMap<String, nixjit_eval::Value>>) {
         self.with_scopes.push(namespace);
     }
 
@@ -240,3 +388,5 @@ impl JITContext for Context {
         self.with_scopes.pop();
     }
 }
+
+impl BuiltinsContext for Context {}