feat(ir): use petgraph

Cargo.lock

@@ -342,6 +342,12 @@ dependencies = [
  "windows-sys 0.59.0",
 ]
 
+[[package]]
+name = "fixedbitset"
+version = "0.5.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1d674e81391d1e1ab681a28d99df07927c6d4aa5b027d7da16ba32d1d21ecd99"
+
 [[package]]
 name = "foldhash"
 version = "0.1.5"
@@ -495,6 +501,7 @@ dependencies = [
  "hashbrown 0.15.3",
  "itertools",
  "lru",
+ "petgraph",
  "priority-queue",
  "regex",
  "replace_with",
@@ -503,6 +510,18 @@ dependencies = [
  "thiserror",
 ]
 
+[[package]]
+name = "petgraph"
+version = "0.8.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "54acf3a685220b533e437e264e4d932cfbdc4cc7ec0cd232ed73c08d03b8a7ca"
+dependencies = [
+ "fixedbitset",
+ "hashbrown 0.15.3",
+ "indexmap",
+ "serde",
+]
+
 [[package]]
 name = "priority-queue"
 version = "2.5.0"

Cargo.toml

@@ -10,10 +10,6 @@ repl = ["dep:rustyline"]
 name = "repl"
 required-features = ["repl"]
 
-[[bin]]
-name = "scc"
-required-features = ["repl"]
-
 [profile.perf]
 debug = 2
 strip = false
@@ -30,6 +26,7 @@ derive_more = { version = "2.0", features = ["full"] }
 ecow = "0.2"
 regex = "1.11"
 hashbrown = "0.15"
+petgraph = "0.8"
 priority-queue = "2.5"
 lru = "0.14"
 replace_with = "0.1"

flake.nix

@@ -14,7 +14,7 @@
         {
           default = pkgs.mkShell {
             packages = with pkgs; [
-              (fenix.packages.${system}.complete.withComponents [
+              /* (fenix.packages.${system}.complete.withComponents [
                 "cargo"
                 "clippy"
                 "rust-src"
@@ -22,7 +22,7 @@
                 "rustfmt"
                 "rust-analyzer"
                 "miri"
-              ])
+              ]) */
               gdb
               valgrind
               gemini-cli

src/bin/repl.rs

@@ -40,6 +40,7 @@ fn main() -> Result<()> {
                 }
                 let expr = root.tree().expr().unwrap();
                 let downgraded = unwrap!(downgrade(expr));
+                println!("{downgraded:?}");
                 println!("{}", unwrap!(eval(downgraded)));
             }
             Err(ReadlineError::Interrupted) => {

src/bin/scc.rs

@@ -1,59 +0,0 @@
-use itertools::Itertools;
-use rustyline::error::ReadlineError;
-use rustyline::{DefaultEditor, Result};
-
-use nixjit::error::Error;
-use nixjit::ir::downgrade;
-
-macro_rules! unwrap {
-    ($e:expr) => {
-        match $e {
-            Ok(ok) => ok,
-            Err(err) => {
-                println!("{err}");
-                continue;
-            }
-        }
-    };
-}
-
-fn main() -> Result<()> {
-    let mut rl = DefaultEditor::new()?;
-    loop {
-        let readline = rl.readline("nixjit-scc-analyzer> ");
-        match readline {
-            Ok(expr) => {
-                if expr.trim().is_empty() {
-                    continue;
-                }
-                rl.add_history_entry(expr.as_str())?;
-                let root = rnix::Root::parse(&expr);
-                if !root.errors().is_empty() {
-                    println!(
-                        "{}",
-                        Error::ParseError(
-                            root.errors().iter().map(|err| err.to_string()).join(";")
-                        )
-                    );
-                    continue;
-                }
-                let expr = root.tree().expr().unwrap();
-                let downgraded = unwrap!(downgrade(expr));
-                println!("{:?}", downgraded);
-            }
-            Err(ReadlineError::Interrupted) => {
-                println!("CTRL-C");
-                break;
-            }
-            Err(ReadlineError::Eof) => {
-                println!("CTRL-D");
-                break;
-            }
-            Err(err) => {
-                println!("Error: {err:?}");
-                break;
-            }
-        }
-    }
-    Ok(())
-}

src/engine/mod.rs

@@ -3,13 +3,14 @@ use std::cell::OnceCell;
 use std::rc::Rc;
 
 use hashbrown::{HashMap, HashSet};
+use itertools::Itertools;
 use priority_queue::PriorityQueue;
 
 use crate::env::Env;
 use crate::error::Result;
 use crate::eval::Evaluate;
 use crate::eval::jit::{JITCompiler, JITFunc};
-use crate::ir::{Dep, Downgraded, Ir, SccNode};
+use crate::ir::{Downgraded, Ir};
 use crate::ty::internal as i;
 use crate::ty::public::Value;
 
@@ -22,7 +23,8 @@ type EnvIdx = usize;
 pub struct Engine {
     pub thunks: Box<[Ir]>,
     pub funcs: Box<[Ir]>,
-    pub func_deps: Vec<HashMap<Dep, usize>>,
+    pub func_deps: Vec<HashMap<usize, usize>>,
+    pub func_arg_deps: Vec<Vec<usize>>,
     jit: JITCompiler,
     compiled: Box<[OnceCell<JITFunc>]>,
     tasks: PriorityQueue<CompileTask, usize>,
@@ -33,6 +35,7 @@ pub fn eval(downgraded: Downgraded) -> Result<Value> {
         downgraded.thunks,
         downgraded.funcs,
         downgraded.func_deps,
+        downgraded.func_arg_deps,
         JITCompiler::new(),
     );
     engine.eval(downgraded.graph)
@@ -42,7 +45,8 @@ impl Engine {
     pub fn new(
         thunks: Box<[Ir]>,
         funcs: Box<[Ir]>,
-        func_deps: Vec<HashMap<Dep, usize>>,
+        func_deps: Vec<HashMap<usize, usize>>,
+        func_arg_deps: Vec<Vec<usize>>,
         jit: JITCompiler,
     ) -> Self {
         Self {
@@ -53,34 +57,29 @@ impl Engine {
             thunks,
             funcs,
             func_deps,
+            func_arg_deps,
             jit,
         }
     }
 
-    pub fn eval(&mut self, graph: Vec<SccNode>) -> Result<Value> {
+    pub fn eval(&mut self, graph: Vec<Vec<usize>>) -> Result<Value> {
         let mut env = Env::new();
-        let last = graph.last().unwrap().members[0];
-        for SccNode { members, .. } in graph.into_iter() {
-            if members.len() == 1 {
-                for member in members.into_iter() {
         let engine = unsafe { &mut *(self as *mut Self) };
-                    let mut val = self.thunks[member].eval(engine, &mut env)?;
+        env.new_frame(0);
-                    val.force(engine, &mut env)?;
+        env.new_frame(0);
-                    env.insert_cache(member, val);
+        for members in graph.into_iter() {
-                }
+            if members.len() == 1 {
+                let val = self.thunks[members[0]].eval(engine, &mut env)?;
+                env.push_stack(val);
             } else {
                 todo!();
                 for member in members.into_iter() {
-                    let val = self.thunks[member].clone().eval(self, &mut env)?;
+                    let val = self.thunks[member].eval(engine, &mut env)?;
-                    env.insert_cache(member, val);
+                    env.push_stack(val);
                 }
             }
         }
-        env.lookup_cache(last, |_| unreachable!()).map(|mut val| {
+        Ok(env.pop_frame().last_mut().unwrap().force(self, &mut env)?.to_public(self, &mut HashSet::new()))
-            Ok(val
-                .force(self, &mut env)?
-                .to_public(self, &mut HashSet::new()))
-        })?
     }
 
     pub fn eval_thunk(&mut self, idx: usize, env: &mut Env) -> Result<i::Value> {
@@ -105,27 +104,15 @@ impl Engine {
     }
 
     pub fn eval_func_deps(&mut self, idx: usize, env: &mut Env) -> Result<()> {
-        for (&dep, _) in unsafe { &*(&self.func_deps[idx] as *const HashMap<Dep, usize>) }.iter() {
-            match dep {
-                Dep::Arg(idx) => {
-                    if let i::Value::Thunk(idx) = env.lookup_arg(idx) {
-                        env.insert_cache_lazy(idx, |env| {
         let engine = unsafe { &mut *(self as *mut Self) };
-                            let mut val = self.thunks[idx].eval(engine, env)?;
+        let engine_mut = unsafe { &mut *(self as *mut Self) };
-                            val.force(engine, env)?;
+        env.new_frame(self.func_deps[idx].len());
-                            val.ok()
+        for (&dep, _) in engine.func_deps[idx].iter().sorted_by_key(|&(_, &k)| k) {
-                        })?;
+            let val = self.thunks[dep].eval(engine_mut, env)?;
-                    }
+            env.push_stack(val);
-                }
-                Dep::Thunk(idx) => {
-                    env.insert_cache_lazy(idx, |env| {
-                        let engine = unsafe { &mut *(self as *mut Self) };
-                        let mut val = self.thunks[idx].eval(engine, env)?;
-                        val.force(engine, env)?;
-                        val.ok()
-                    })?;
-                }
         }
+        for &arg in engine.func_arg_deps[idx].iter() {
+            env.force_arg(arg);
         }
         Ok(())
     }

src/engine/test.rs

@@ -1,11 +1,8 @@
 #![allow(unused_macros)]
 
-extern crate test;
 
 use std::collections::BTreeMap;
 
-use test::{Bencher, black_box};
-
 use crate::ir::downgrade;
 use crate::ty::common::Const;
 use crate::ty::public::*;
@@ -210,14 +207,11 @@ fn test_func() {
     test_expr("let fix = f: let x = f x; in x; in (fix (self: { x = 1; y = self.x + 1; })).y", int!(2));
 }
 
-#[bench]
+#[test]
 #[ignore]
-fn bench_fib(b: &mut Bencher) {
+fn test_fib() {
-    b.iter(|| {
     test_expr(
         "let fib = n: if n == 1 || n == 2 then 1 else (fib (n - 1)) + (fib (n - 2)); in fib 30",
         int!(832040),
     );
-        black_box(())
-    })
 }

src/env.rs

@@ -3,78 +3,77 @@ use std::rc::Rc;
 
 use hashbrown::HashMap;
 
-use crate::error::Result;
 use crate::ty::internal::Value;
 
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, Default)]
 pub struct Env {
     stack: Vec<Value>,
+    args: Vec<Value>,
+    frame_offsets: Vec<usize>,
     with: Vec<Rc<HashMap<String, Value>>>,
 }
 
-impl Default for Env {
-    fn default() -> Self {
-        Self::new()
-    }
-}
-
 impl Env {
     pub fn new() -> Self {
-        Self {
+        Self::default()
-            stack: Vec::new(),
+    }
-            with: Vec::new(),
+
+    pub fn reserve_args(&mut self, size: usize) {
+        self.args.reserve(size);
+    }
+
+    pub fn push_arg(&mut self, value: Value) {
+        self.args.push(value);
+    }
+
+    pub fn push_args(&mut self, values: impl IntoIterator<Item = Value>) {
+        self.args.extend(values);
+    }
+
+    pub fn lookup_arg(&self, offset: usize) -> &Value {
+        &self.args[self.args.len() - offset - 1]
+    }
+
+    pub fn pop_args(&mut self, len: usize) -> Vec<Value> {
+        self.args.split_off(self.args.len() - len)
+    }
+
+    pub fn drop_args(&mut self, len: usize) {
+        self.args.truncate(self.args.len() - len);
+    }
+
+    pub fn force_arg(&mut self, offset: usize) {
+        let env = unsafe { &mut *(self as *mut Env) };
+        if let thunk @ Value::Thunk(_) = &mut self.args[offset] {
+            let &mut Value::Thunk(offset) = thunk else { unreachable!() };
+            *thunk = env.lookup_stack(offset).clone()
         }
     }
 
-    pub fn with_new_cache<T>(
+    pub fn new_frame(&mut self, size: usize) {
-        &mut self,
+        self.stack.reserve(size);
-        len: usize,
+        self.frame_offsets.push(self.stack.len());
-        f: impl FnOnce(&mut Self) -> T,
-    ) -> (T, Vec<Value>) {
-        self.stack.reserve(len);
-        let ret = f(self);
-        (ret, self.stack.split_off(self.stack.len() - len))
     }
 
-    pub fn with_cache<T>(
+    pub fn push_stack(&mut self, value: Value) {
-        &mut self,
+        self.stack.push(value);
-        cache: Vec<Value>,
-        len: usize,
-        f: impl FnOnce(&mut Self) -> T,
-    ) -> (T, Vec<Value>) {
-        self.stack.extend(cache);
-        let ret = f(self);
-        (ret, self.stack.split_off(self.stack.len() - len))
     }
 
-    /* pub fn insert_cache(&mut self, idx: usize, val: Value) {
+    pub fn resume_stack(&mut self, iter: impl IntoIterator<Item = Value>) {
-        self.cache.last_mut().unwrap().insert(idx, val);
+        self.stack.extend(iter);
-    } */
-
-    /* pub fn insert_cache_lazy(
-        &mut self,
-        idx: usize,
-        f: impl FnOnce(&mut Self) -> Result<Value>,
-    ) -> Result<()> {
-        if self.cache.last().unwrap().get(&idx).is_none() {
-            let val = f(self)?;
-            self.cache.last_mut().unwrap().insert(idx, val);
-        }
-        Ok(())
-    } */
-
-    pub fn insert_stack(&mut self, iter: impl IntoIterator<Item = fn(&mut Self) -> Result<Value>>) -> Result<()> {
-        let iter = iter.into_iter();
-        self.stack.reserve(iter.size_hint().0);
-        for f in iter {
-            let val = f(self)?;
-            self.stack.push(val);
-        }
-        Ok(())
     }
 
-    pub fn lookup_stack(&mut self, idx: usize) -> Value {
+    pub fn lookup_stack(&self, offset: usize) -> &Value {
-        self.stack.get(self.stack.len() - idx - 1).unwrap().clone()
+        dbg!(self.frame_offsets[self.frame_offsets.len() - 2], offset);
+        &self.stack[self.frame_offsets[self.frame_offsets.len() - 2] + offset]
+    }
+
+    pub fn pop_frame(&mut self) -> Vec<Value> {
+        self.stack.split_off(self.frame_offsets.pop().unwrap())
+    }
+
+    pub fn drop_frame(&mut self) {
+        self.stack.truncate(self.frame_offsets.pop().unwrap());
     }
 
     pub fn lookup_with(&self, symbol: &str) -> Option<Value> {
@@ -86,11 +85,11 @@ impl Env {
         None
     }
 
-    pub fn exit_with(&mut self) {
-        self.with.pop();
-    }
-
     pub fn enter_with(&mut self, map: Rc<HashMap<String, Value>>) {
         self.with.push(map)
     }
+
+    pub fn exit_with(&mut self) {
+        self.with.pop();
+    }
 }

src/eval/jit/compile.rs

@@ -530,19 +530,13 @@ impl JITCompile for Var {
 
 impl JITCompile for Arg {
     fn compile(&self, ctx: &mut JITContext, engine: ir::Value, env: ir::Value) -> StackSlot {
-        ctx.lookup_arg(env, self.level)
+        ctx.lookup_arg(env, self.offset)
     }
 }
 
-impl JITCompile for LetVar {
+impl JITCompile for StackVar {
     fn compile(&self, ctx: &mut JITContext, engine: ir::Value, env: ir::Value) -> StackSlot {
-        unreachable!()
+        ctx.lookup_stack(env, self.offset)
-    }
-}
-
-impl JITCompile for Cache {
-    fn compile(&self, ctx: &mut JITContext, engine: ir::Value, env: ir::Value) -> StackSlot {
-        todo!()
     }
 }
 

src/eval/jit/helpers.rs

@@ -25,8 +25,12 @@ pub extern "C" fn helper_call(
         .unwrap();
 }
 
-pub extern "C" fn helper_lookup_arg(env: &Env, level: usize, ret: &mut MaybeUninit<Value>) {
+pub extern "C" fn helper_lookup_stack(env: &Env, offset: usize, ret: &mut MaybeUninit<Value>) {
-    ret.write(env.lookup_arg(level));
+    ret.write(env.lookup_stack(offset).clone());
+}
+
+pub extern "C" fn helper_lookup_arg(env: &Env, offset: usize, ret: &mut MaybeUninit<Value>) {
+    ret.write(env.lookup_arg(offset).clone());
 }
 
 pub extern "C" fn helper_lookup(

src/eval/jit/mod.rs

@@ -250,6 +250,24 @@ impl<'comp, 'ctx> JITContext<'comp, 'ctx> {
         slot
     }
 
+    fn lookup_stack(&mut self, env: ir::Value, idx: usize) -> StackSlot {
+        let slot = self.alloca();
+        let lookup_stack = self
+            .compiler
+            .module
+            .declare_func_in_func(self.compiler.lookup_stack, self.builder.func);
+        let idx = self
+            .builder
+            .ins()
+            .iconst(self.compiler.ptr_type, idx as i64);
+        let ptr = self
+            .builder
+            .ins()
+            .stack_addr(self.compiler.ptr_type, slot, 0);
+        self.builder.ins().call(lookup_stack, &[env, idx, ptr]);
+        slot
+    }
+
     fn lookup_arg(&mut self, env: ir::Value, idx: usize) -> StackSlot {
         let slot = self.alloca();
         let lookup_arg = self
@@ -374,6 +392,7 @@ pub struct JITCompiler {
     func_sig: Signature,
 
     call: FuncId,
+    lookup_stack: FuncId,
     lookup_arg: FuncId,
     lookup: FuncId,
     select: FuncId,
@@ -413,6 +432,7 @@ impl JITCompiler {
         let mut builder = JITBuilder::with_isa(isa, cranelift_module::default_libcall_names());
 
         builder.symbol("helper_call", helper_call as _);
+        builder.symbol("helper_lookup_stack", helper_lookup_stack as _);
         builder.symbol("helper_lookup_arg", helper_lookup_arg as _);
         builder.symbol("helper_lookup", helper_lookup as _);
         builder.symbol("helper_select", helper_select as _);
@@ -466,6 +486,18 @@ impl JITCompiler {
             .declare_function("helper_call", Linkage::Import, &call_sig)
             .unwrap();
 
+        let mut lookup_stack_sig = module.make_signature();
+        lookup_stack_sig.params.extend(
+            [AbiParam {
+                value_type: ptr_type,
+                purpose: ArgumentPurpose::Normal,
+                extension: ArgumentExtension::None,
+            }; 3],
+        );
+        let lookup_stack = module
+            .declare_function("helper_lookup_stack", Linkage::Import, &lookup_stack_sig)
+            .unwrap();
+
         // fn(env: &Env, level: usize, ret: &mut MaybeUninit<Value>)
         let mut lookup_arg_sig = module.make_signature();
         lookup_arg_sig.params.extend(
@@ -669,6 +701,7 @@ impl JITCompiler {
             func_sig,
 
             call,
+            lookup_stack,
             lookup_arg,
             lookup,
             select,

src/eval/mod.rs

@@ -294,26 +294,20 @@ impl Evaluate for ir::Var {
 }
 
 impl Evaluate for ir::Arg {
-    fn eval(&self, _: &mut Engine, env: &mut Env) -> Result<Value> {
-        env.lookup_arg(self.level).ok()
-    }
-}
-
-impl Evaluate for ir::LetVar {
-    fn eval(&self, _: &mut Engine, env: &mut Env) -> Result<Value> {
-        unreachable!()
-    }
-}
-
-impl Evaluate for ir::Cache {
     fn eval(&self, engine: &mut Engine, env: &mut Env) -> Result<Value> {
-        todo!()
+        env.lookup_arg(self.offset).clone().ok()
+    }
+}
+
+impl Evaluate for ir::StackVar {
+    fn eval(&self, engine: &mut Engine, env: &mut Env) -> Result<Value> {
+        env.lookup_stack(self.offset).clone().ok()
     }
 }
 
 impl Evaluate for ir::Thunk {
     fn eval(&self, engine: &mut Engine, env: &mut Env) -> Result<Value> {
-        Value::Thunk(self.idx).ok()
+        unreachable!()
     }
 }
 

src/ir/ctx.rs

@@ -1,10 +1,10 @@
 use derive_more::Unwrap;
 use hashbrown::HashMap;
 
-use crate::error::Result;
+use crate::{error::Result, ir::sort_dependencies};
 use crate::ty::common::Const;
 
-use super::{Dep, Func, Ir, LoadFunc, MaybeThunk, SccAnalyzer, SccNode, Thunk};
+use super::{Func, Ir, LoadFunc, MaybeThunk, Thunk};
 
 #[derive(Clone, Copy, Unwrap, Debug)]
 pub enum Index {
@@ -15,7 +15,8 @@ pub enum Index {
 pub struct DowngradeContext {
     pub thunks: Vec<(Ir, bool)>,
     pub thunk_deps: Vec<HashMap<usize, usize>>,
-    pub func_deps: Vec<HashMap<Dep, usize>>,
+    pub func_deps: Vec<HashMap<usize, usize>>,
+    pub func_arg_deps: Vec<Vec<usize>>,
     pub funcs: Vec<Func>,
 }
 
@@ -35,7 +36,6 @@ enum EnvNode<'a> {
 pub enum LookupResult {
     Builtin(Ir),
     MaybeThunk(MaybeThunk),
-    Let { level: usize, idx: usize },
     SingleArg { idx: usize },
     MultiArg { idx: usize, default: Option<Ir> },
     With,
@@ -147,6 +147,7 @@ impl DowngradeContext {
             thunks: Vec::new(),
             thunk_deps: Vec::new(),
             func_deps: Vec::new(),
+            func_arg_deps: Vec::new(),
             funcs: Vec::new(),
         }
     }
@@ -169,17 +170,29 @@ impl DowngradeContext {
         }
     }
 
-    pub fn new_dep(&mut self, this: Index, dep: Dep) -> usize {
+    pub fn new_thunk_dep(&mut self, this: Index, thunk: usize) -> usize {
         match this {
             Index::Thunk(idx) => {
-                let len = self.thunk_deps.len();
+                let len = self.thunk_deps[idx].len();
                 *self.thunk_deps[idx]
-                    .entry(dep.unwrap_thunk())
+                    .entry(thunk)
                     .or_insert(len)
             }
             Index::Func(idx) => {
-                let len = self.thunk_deps.len();
+                let len = self.func_deps[idx].len();
-                *self.func_deps[idx].entry(dep).or_insert(len)
+                *self.func_deps[idx].entry(thunk).or_insert(len)
+            }
+        }
+    }
+
+    pub fn new_arg_dep(&mut self, this: Index, arg: usize) {
+        match this {
+            Index::Func(idx) => {
+                self.func_arg_deps[idx].push(arg);
+            }
+            Index::Thunk(idx) => {
+                return;
+                panic!("unexpected dependency from Thunk({idx}) to Arg({arg})")
             }
         }
     }
@@ -188,6 +201,7 @@ impl DowngradeContext {
         let idx = self.funcs.len();
         self.funcs.push(func);
         self.func_deps.push(HashMap::new());
+        self.func_arg_deps.push(Vec::new());
         LoadFunc { idx }
     }
 
@@ -246,15 +260,17 @@ impl DowngradeContext {
 pub struct Downgraded {
     pub thunks: Box<[Ir]>,
     pub funcs: Box<[Ir]>,
-    pub func_deps: Vec<HashMap<Dep, usize>>,
+    pub func_deps: Vec<HashMap<usize, usize>>,
-    pub graph: Vec<SccNode>,
+    pub func_arg_deps: Vec<Vec<usize>>,
+    pub graph: Vec<Vec<usize>>,
 }
 
 impl Downgraded {
     pub fn new(ctx: DowngradeContext) -> Self {
         Self {
-            graph: SccAnalyzer::new(&ctx).analyze(),
+            graph: sort_dependencies(&ctx),
             func_deps: ctx.func_deps,
+            func_arg_deps: ctx.func_arg_deps,
             thunks: ctx.thunks.into_iter().map(|(ir, _)| ir).collect(),
             funcs: ctx
                 .funcs

src/ir/graph.rs

@@ -0,0 +1,37 @@
+use hashbrown::HashMap;
+use petgraph::graph::{DiGraph, NodeIndex};
+use petgraph::algo::{condensation, toposort};
+
+use super::*;
+
+struct GraphBuilder<'ctx> {
+    ctx: &'ctx DowngradeContext,
+    graph: DiGraph<usize, ()>,
+    nodes: HashMap<usize, NodeIndex>
+}
+
+impl GraphBuilder<'_> {
+    fn connect(&mut self, idx: usize) {
+        let node = self.graph.add_node(idx);
+        self.nodes.insert(idx, node);
+        let deps = self.ctx.thunk_deps.get(idx).unwrap();
+        for (&refee, _) in deps {
+            if !self.nodes.contains_key(&refee) {
+                self.connect(refee);
+            }
+            self.graph.add_edge(node, *self.nodes.get(&refee).unwrap(), ());
+        }
+    }
+}
+
+pub fn sort_dependencies(ctx: &DowngradeContext) -> Vec<Vec<usize>> {
+    let mut builder = GraphBuilder {
+        ctx,
+        graph: DiGraph::new(),
+        nodes: HashMap::new(),
+    };
+    builder.connect(ctx.thunks.len() - 1);
+    let mut graph = condensation(builder.graph, false);
+    let sorted = toposort(&graph, None).unwrap();
+    sorted.into_iter().map(|idx| graph.remove_node(idx).unwrap()).collect_vec()
+}

src/ir/mod.rs

@@ -143,11 +143,9 @@ ir! {
     Str => { val: String },
     Var => { sym: String },
     #[derive(Copy)]
-    Arg => { level: usize },
+    Arg => { offset: usize },
     #[derive(Copy)]
-    LetVar => { level: usize, idx: usize },
+    StackVar => { offset: usize },
-    #[derive(Copy)]
-    Cache => { idx: usize },
     #[derive(Copy)]
     Thunk => { idx: usize },
     Path => { expr: Box<Ir> },
@@ -364,13 +362,24 @@ impl Thunk {
         ctx: &mut DowngradeContext,
         env: &Env<'a, 'env>,
     ) -> Result<Ir> {
-        let idx = ctx.new_dep(self_idx, Dep::Thunk(self.idx));
+        let idx = ctx.new_thunk_dep(self_idx, self.idx);
         ctx.resolve_thunk(self.idx, env)?;
-        Cache { idx }.ir().ok()
+        StackVar { offset: idx }.ir().ok()
     }
 }
 
-impl Cache {
+impl Arg {
+    fn resolve<'a, 'env>(
+        self,
+        self_idx: Index,
+        ctx: &mut DowngradeContext,
+        env: &Env<'a, 'env>,
+    ) -> Result<Ir> {
+        unreachable!()
+    }
+}
+
+impl StackVar {
     fn resolve<'a, 'env>(
         self,
         self_idx: Index,
@@ -705,17 +714,22 @@ impl Var {
         };
         match res {
             Builtin(ir) => ir,
-            Let { level, idx } => LetVar { level, idx }.ir(),
+            SingleArg { idx } => {
-            SingleArg { idx: level } => Arg { level }.ir(),
+                ctx.new_arg_dep(self_idx, idx);
+                Arg { offset: idx }.ir()
+            },
             MultiArg {
-                idx: level,
+                idx,
                 default,
-            } => Select {
+            } => {
-                expr: Arg { level }.ir().boxed(),
+                ctx.new_thunk_dep(self_idx, idx);
+                Select {
+                    expr: Arg { offset: idx }.ir().boxed(),
                     attrpath: vec![Attr::Str(self.sym)],
                     default: default.map(Box::new),
                 }
-            .ir(),
+                .ir()
+            }
             MaybeThunk(thunk) => thunk.resolve(self_idx, ctx, env)?,
             With => self.ir(),
         }
@@ -723,28 +737,6 @@ impl Var {
     }
 }
 
-impl Arg {
-    fn resolve<'a, 'env>(
-        self,
-        _: Index,
-        _: &mut DowngradeContext,
-        _: &Env<'a, 'env>,
-    ) -> Result<Ir> {
-        self.ir().ok()
-    }
-}
-
-impl LetVar {
-    fn resolve<'a, 'env>(
-        self,
-        _: Index,
-        _: &mut DowngradeContext,
-        _: &Env<'a, 'env>,
-    ) -> Result<Ir> {
-        self.ir().ok()
-    }
-}
-
 impl Downgrade for ast::AttrSet {
     fn downgrade(self, ctx: &mut DowngradeContext) -> Result<Ir> {
         let rec = self.rec_token().is_some();

src/ir/scc.rs

@@ -1,203 +0,0 @@
-use std::collections::VecDeque;
-
-use hashbrown::{HashMap, HashSet};
-
-use super::*;
-
-#[derive(Debug, Clone)]
-pub struct SccNode {
-    pub id: usize,
-    pub members: Vec<usize>,
-    pub deps: HashSet<usize>,
-}
-
-#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Unwrap)]
-pub enum Dep {
-    Thunk(usize),
-    Arg(usize),
-}
-
-#[derive(Default, Debug)]
-pub struct SccGraph {
-    nodes: HashMap<usize, SccNode>,
-    root: usize,
-}
-
-impl SccGraph {
-    fn new(ctx: &DowngradeContext, sccs: Vec<Vec<usize>>) -> Self {
-        let mut graph = SccGraph::default();
-        let mut thunk_to_scc = HashMap::new();
-
-        for (id, members) in sccs.iter().enumerate() {
-            for &thunk_id in members {
-                thunk_to_scc.insert(thunk_id, id);
-            }
-            graph.nodes.insert(
-                id,
-                SccNode {
-                    id,
-                    members: members.clone(),
-                    deps: HashSet::new(),
-                },
-            );
-        }
-        graph.root = thunk_to_scc[&(ctx.thunks.len() - 1)];
-
-        for (from_node_id, from_deps) in ctx.thunk_deps.iter().enumerate() {
-            let from_scc_id = thunk_to_scc[&from_node_id];
-            for (&to_node_id, _) in from_deps {
-                let to_scc_id = thunk_to_scc[&to_node_id];
-                if from_scc_id != to_scc_id {
-                    graph
-                        .nodes
-                        .get_mut(&from_scc_id)
-                        .unwrap()
-                        .deps
-                        .insert(to_scc_id);
-                }
-            }
-        }
-
-        graph
-    }
-
-    fn sorted(self) -> Vec<SccNode> {
-        let mut reachable = HashSet::new();
-        let mut stack = vec![self.root];
-        reachable.insert(self.root);
-
-        while let Some(id) = stack.pop() {
-            if let Some(node) = self.nodes.get(&id) {
-                for &dep_id in &node.deps {
-                    if reachable.insert(dep_id) {
-                        stack.push(dep_id);
-                    }
-                }
-            }
-        }
-
-        let mut in_degrees: HashMap<usize, usize> = HashMap::new();
-        let mut reverse_adj: HashMap<usize, Vec<usize>> = HashMap::new();
-
-        for &id in &reachable {
-            in_degrees.insert(id, 0);
-        }
-
-        for &id in &reachable {
-            if let Some(node) = self.nodes.get(&id) {
-                for &dep_id in &node.deps {
-                    if reachable.contains(&dep_id) {
-                        reverse_adj.entry(dep_id).or_default().push(id);
-                        *in_degrees.get_mut(&id).unwrap() += 1;
-                    }
-                }
-            }
-        }
-
-        let mut queue = in_degrees
-            .iter()
-            .filter(|&(_, &deg)| deg == 0)
-            .map(|(&id, _)| id)
-            .collect::<VecDeque<_>>();
-
-        queue.make_contiguous().sort();
-
-        let mut sorted_order = Vec::new();
-        while let Some(u) = queue.pop_front() {
-            if let Some(node) = self.nodes.get(&u) {
-                sorted_order.push(node.clone());
-            }
-
-            if let Some(dependents) = reverse_adj.get(&u) {
-                for &v in dependents {
-                    if let Some(degree) = in_degrees.get_mut(&v) {
-                        *degree -= 1;
-                        if *degree == 0 {
-                            queue.push_back(v);
-                        }
-                    }
-                }
-            }
-        }
-
-        if sorted_order.len() != reachable.len() {
-            unreachable!("Cycle detected in the reachable part of SCC graph!");
-        }
-
-        sorted_order
-    }
-}
-
-pub struct SccAnalyzer<'ctx> {
-    ctx: &'ctx DowngradeContext,
-    index: usize,
-    stack: Vec<usize>,
-    on_stack: HashSet<usize>,
-    indices: HashMap<usize, usize>,
-    low_links: HashMap<usize, usize>,
-    sccs: Vec<Vec<usize>>,
-}
-
-impl<'ctx> SccAnalyzer<'ctx> {
-    pub fn new(ctx: &'ctx DowngradeContext) -> Self {
-        Self {
-            ctx,
-            index: 0,
-            stack: Vec::new(),
-            on_stack: HashSet::new(),
-            indices: HashMap::new(),
-            low_links: HashMap::new(),
-            sccs: Vec::new(),
-        }
-    }
-
-    pub fn analyze(mut self) -> Vec<SccNode> {
-        for idx in 0..self.ctx.thunks.len() {
-            if !self.indices.contains_key(&idx) {
-                self.strong_connect(idx);
-            }
-        }
-        SccGraph::new(self.ctx, self.sccs).sorted()
-    }
-
-    fn strong_connect(&mut self, v_id: usize) {
-        self.indices.insert(v_id, self.index);
-        self.low_links.insert(v_id, self.index);
-        self.index += 1;
-        self.stack.push(v_id);
-        self.on_stack.insert(v_id);
-
-        let Some(deps) = self.ctx.thunk_deps.get(v_id) else {
-            unreachable!()
-        };
-        for (&w_id, _) in deps {
-            if !self.indices.contains_key(&w_id) {
-                self.strong_connect(w_id);
-                let v_low_link = *self.low_links.get(&v_id).unwrap();
-                let w_low_link = *self.low_links.get(&w_id).unwrap();
-                if w_low_link < v_low_link {
-                    self.low_links.insert(v_id, w_low_link);
-                }
-            } else if self.on_stack.contains(&w_id) {
-                let v_low_link = *self.low_links.get(&v_id).unwrap();
-                let w_index = *self.indices.get(&w_id).unwrap();
-                if w_index < v_low_link {
-                    self.low_links.insert(v_id, w_index);
-                }
-            }
-        }
-
-        if self.low_links[&v_id] == self.indices[&v_id] {
-            let mut scc = Vec::new();
-            loop {
-                let w_id = self.stack.pop().unwrap();
-                self.on_stack.remove(&w_id);
-                scc.push(w_id);
-                if w_id == v_id {
-                    break;
-                }
-            }
-            self.sccs.push(scc);
-        }
-    }
-}

src/lib.rs

@@ -1,4 +1,3 @@
-#![cfg_attr(test, feature(test))]
 #![allow(dead_code)]
 
 mod builtins;

src/ty/internal/func.rs

@@ -7,5 +7,5 @@ use super::Value;
 pub struct PartialFunc {
     pub idx: usize,
     pub args: Vec<Value>,
-    pub cache: HashMap<usize, Value>,
+    pub frame: Vec<Value>,
 }

src/ty/internal/mod.rs

@@ -183,21 +183,21 @@ impl Value {
                 let self::PartialFunc {
                     idx,
                     args: old_args,
-                    cache,
+                    frame,
                 } = Rc::make_mut(func);
                 let idx = *idx;
                 let len = args.len() + old_args.len();
                 env.reserve_args(len);
-                env.enter_args(std::mem::take(old_args));
+                env.push_args(core::mem::take(old_args));
                 let mut args = args.into_iter().peekable();
-                env.enter_arg(args.next().unwrap());
+                env.push_arg(args.next().unwrap());
-                let (mut ret, cache) = env.with_cache(std::mem::take(cache), |env| {
+                env.resume_stack(core::mem::take(frame));
                 engine.eval_func_deps(idx, env)?;
                 let mut ret = engine.call_func(idx, env)?;
                 while args.peek().is_some() {
                     match ret {
                         Value::Func(func) => {
-                                env.enter_arg(args.next().unwrap());
+                            env.push_arg(args.next().unwrap());
                             engine.eval_func_deps(idx, env)?;
                             ret = engine.call_func(func, env)?;
                         }
@@ -215,31 +215,32 @@ impl Value {
                         _ => todo!(),
                     }
                 }
-                    ret.ok()
+                let frame = env.pop_frame();
-                });
+                if let Value::Func(idx) = ret {
                     let args = env.pop_args(len);
-                if let Ok(Value::Func(idx)) = ret {
+                    ret = PartialFunc(self::PartialFunc::new(idx, args, frame).into());
-                    ret = PartialFunc(self::PartialFunc::new(idx, args, cache).into()).ok();
+                } else if let Value::PartialFunc(func) = &mut ret {
-                } else if let Ok(Value::PartialFunc(func)) = &mut ret {
+                    let args = env.pop_args(len);
-                    Rc::make_mut(func).args.extend(args);
+                    let func = Rc::make_mut(func);
+                    func.args.extend(args);
+                } else {
+                    env.drop_args(len);
                 }
-                ret
+                ret.ok()
             }
             &mut Func(idx) => {
                 let len = args.len();
-                let mut args = args.into_iter().peekable();
                 env.reserve_args(len);
-                env.enter_arg(args.next().unwrap());
+                let mut args = args.into_iter().peekable();
-                let (mut ret, cache) = env.with_new_cache(|env| {
+                env.push_arg(args.next().unwrap());
                 engine.eval_func_deps(idx, env)?;
                 let mut ret = engine.call_func(idx, env)?;
-                    ret.force(engine, env)?;
                 while args.peek().is_some() {
                     match ret {
-                            Value::Func(idx) => {
+                        Value::Func(func) => {
-                                env.enter_arg(args.next().unwrap());
+                            env.push_arg(args.next().unwrap());
                             engine.eval_func_deps(idx, env)?;
-                                ret = engine.call_func(idx, env)?;
+                            ret = engine.call_func(func, env)?;
                         }
                         Value::PartialFunc(_) => {
                             todo!()
@@ -255,15 +256,18 @@ impl Value {
                         _ => todo!(),
                     }
                 }
-                    ret.ok()
+                let frame = env.pop_frame();
-                });
+                if let Value::Func(idx) = ret {
                     let args = env.pop_args(len);
-                if let Ok(Value::Func(idx)) = ret {
+                    ret = PartialFunc(self::PartialFunc::new(idx, args, frame).into());
-                    ret = Value::PartialFunc(self::PartialFunc::new(idx, args, cache).into()).ok()
+                } else if let Value::PartialFunc(func) = &mut ret {
-                } else if let Ok(Value::PartialFunc(func)) = &mut ret {
+                    let args = env.pop_args(len);
-                    Rc::make_mut(func).args.extend(args);
+                    let func = Rc::make_mut(func);
+                    func.args.extend(args);
+                } else {
+                    env.drop_args(len);
                 }
-                ret
+                ret.ok()
             }
             Catchable(_) => return Ok(()),
             other => todo!("{}", other.typename()),
@@ -517,7 +521,7 @@ impl Value {
 
     pub fn force(&mut self, engine: &mut Engine, env: &mut Env) -> Result<&mut Self> {
         if let &mut Value::Thunk(idx) = self {
-            *self = env.lookup_stack(idx, |_| unreachable!())?
+            *self = env.lookup_stack(idx).clone();
         }
         Ok(self)
     }