feat: WIP

evaluator/nixjit_context/src/lib.rs

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