brewery
notashelf /
bfbcef9c284c54772e80b5a7f33efe9f9bcf75fa

nixir

public

Import-resolving Nix IR plugin

Star0tarball
clone
ssh://git@git.notashelf.dev:33/notashelf/nixir.git

Commit bfbcef9c284c

tarball

NotAShelf <raf@notashelf.dev> · 2026-07-07 18:13 UTC

unverified2 files changed+191-150
irc/evaluator: speed it up

Signed-off-by: NotAShelf <raf@notashelf.dev>
Change-Id: If1ff039e333110772783dbd42b0d839d6a6a6964
diff --git a/src/irc/evaluator.cpp b/src/irc/evaluator.cppindex 1684d40..e56222e 100644--- a/src/irc/evaluator.cpp+++ b/src/irc/evaluator.cpp@@ -6,6 +6,7 @@ #include "nix/expr/eval.hh" #include "nix/expr/value.hh" #include "nix/util/url.hh"+#include <deque> #include <unordered_map>  namespace nix_irc {@@ -17,7 +18,10 @@ struct IREnvironment {   std::vector<Value*> bindings;   Value* with_attrs; -  explicit IREnvironment(IREnvironment* p = nullptr) : parent(p), with_attrs(nullptr) {}+  explicit IREnvironment(IREnvironment* p = nullptr, size_t binding_capacity = 0)+      : parent(p), with_attrs(nullptr) {+    bindings.reserve(binding_capacity);+  }    void bind(Value* val) { bindings.push_back(val); } @@ -63,9 +67,15 @@ struct Thunk { };  struct Evaluator::Impl {+  struct IRLambdaClosure {+    IREnvironment* env;+    std::shared_ptr<Node> body;+  };+   EvalState& state;-  std::unordered_map<Value*, std::unique_ptr<Thunk>> thunks;-  std::vector<std::unique_ptr<IREnvironment>> environments;+  std::unordered_map<Value*, Thunk> thunks;+  std::unordered_map<PrimOp*, IRLambdaClosure> ir_lambdas;+  std::deque<IREnvironment> environments;    explicit Impl(EvalState& s) : state(s) {} @@ -99,15 +109,14 @@ struct Evaluator::Impl {     return escaped;   } -  IREnvironment* make_env(IREnvironment* parent = nullptr) {-    auto env = new IREnvironment(parent);-    environments.push_back(std::unique_ptr<IREnvironment>(env));-    return env;+  IREnvironment* make_env(IREnvironment* parent = nullptr, size_t binding_capacity = 0) {+    environments.emplace_back(parent, binding_capacity);+    return &environments.back();   }    Value* make_thunk(const std::shared_ptr<Node>& expr, IREnvironment* env) {     Value* v = state.allocValue();-    thunks[v] = std::make_unique<Thunk>(expr, env);+    thunks.emplace(v, Thunk(expr, env));     return v;   } @@ -117,7 +126,7 @@ struct Evaluator::Impl {       return;     } -    Thunk* thunk = it->second.get();+    Thunk* thunk = &it->second;     if (thunk->blackholed) {       state.error<EvalError>("infinite recursion encountered").debugThrow();     }@@ -132,7 +141,9 @@ struct Evaluator::Impl {     if (!src)       return;     force(src);-    state.forceValue(*src, noPos);+    if (src->isThunk()) {+      state.forceValue(*src, noPos);+    }     switch (src->type()) {     case nInt:       dest.mkInt(src->integer());@@ -216,107 +227,116 @@ struct Evaluator::Impl {     } else if (auto* n = node->get_if<LambdaNode>()) {       auto lambda_env = env;       auto body = n->body;--      auto primOp = [this, lambda_env, body](EvalState& state, PosIdx pos, Value** args,-                                             Value& result) {-        auto call_env = make_env(lambda_env);-        call_env->bind(args[0]);-        eval_node(body, result, call_env);-      };--      v.mkPrimOp(new PrimOp{.name = n->param_name.value_or("lambda"),-                            .arity = static_cast<size_t>(n->arity),-                            .fun = primOp});+      auto* prim = new PrimOp{.name = n->param_name.value_or("lambda"),+                              .arity = static_cast<size_t>(n->arity),+                              .fun = [this, lambda_env, body](EvalState& state, PosIdx pos,+                                                              Value** args, Value& result) {+                                auto call_env = make_env(lambda_env, 1);+                                call_env->bind(args[0]);+                                eval_node(body, result, call_env);+                              }};++      ir_lambdas.emplace(prim, IRLambdaClosure{.env = lambda_env, .body = body});+      v.mkPrimOp(prim);     } else if (auto* n = node->get_if<AppNode>()) {-      Value* func_val = state.allocValue();-      eval_node(n->func, *func_val, env);-      force(func_val);+      Value func_val;+      eval_node(n->func, func_val, env);+      force(&func_val);        Value* arg_val = make_thunk(n->arg, env);+      if (func_val.isPrimOp()) {+        auto it = ir_lambdas.find(const_cast<PrimOp*>(func_val.primOp()));+        if (it != ir_lambdas.end()) {+          auto call_env = make_env(it->second.env, 1);+          call_env->bind(arg_val);+          eval_node(it->second.body, v, call_env);+          return;+        }+      } -      state.callFunction(*func_val, *arg_val, v, noPos);+      state.callFunction(func_val, *arg_val, v, noPos);     } else if (auto* n = node->get_if<BinaryOpNode>()) {-      Value* left = state.allocValue();-      Value* right = state.allocValue();+      Value left;+      Value right;        switch (n->op) {       case BinaryOp::AND:-        eval_node(n->left, *left, env);-        force(left);-        if (left->type() != nBool) {+        eval_node(n->left, left, env);+        force(&left);+        if (left.type() != nBool) {           state.error<EvalError>("type error in logical AND").debugThrow();         }-        if (!left->boolean()) {+        if (!left.boolean()) {           v.mkBool(false);         } else {-          eval_node(n->right, *right, env);-          force(right);-          if (right->type() != nBool) {+          eval_node(n->right, right, env);+          force(&right);+          if (right.type() != nBool) {             state.error<EvalError>("type error in logical AND").debugThrow();           }-          v.mkBool(right->boolean());+          v.mkBool(right.boolean());         }         break;       case BinaryOp::OR:-        eval_node(n->left, *left, env);-        force(left);-        if (left->type() != nBool) {+        eval_node(n->left, left, env);+        force(&left);+        if (left.type() != nBool) {           state.error<EvalError>("type error in logical OR").debugThrow();         }-        if (left->boolean()) {+        if (left.boolean()) {           v.mkBool(true);         } else {-          eval_node(n->right, *right, env);-          force(right);-          if (right->type() != nBool) {+          eval_node(n->right, right, env);+          force(&right);+          if (right.type() != nBool) {             state.error<EvalError>("type error in logical OR").debugThrow();           }-          v.mkBool(right->boolean());+          v.mkBool(right.boolean());         }         break;       case BinaryOp::IMPL:-        eval_node(n->left, *left, env);-        force(left);-        if (left->type() != nBool) {+        eval_node(n->left, left, env);+        force(&left);+        if (left.type() != nBool) {           state.error<EvalError>("type error in implication").debugThrow();         }-        if (!left->boolean()) {+        if (!left.boolean()) {           v.mkBool(true);         } else {-          eval_node(n->right, *right, env);-          force(right);-          if (right->type() != nBool) {+          eval_node(n->right, right, env);+          force(&right);+          if (right.type() != nBool) {             state.error<EvalError>("type error in implication").debugThrow();           }-          v.mkBool(right->boolean());+          v.mkBool(right.boolean());         }         break;       default:-        eval_node(n->left, *left, env);-        eval_node(n->right, *right, env);-        force(left);-        force(right);+        eval_node(n->left, left, env);+        eval_node(n->right, right, env);+        force(&left);+        force(&right);          switch (n->op) {         case BinaryOp::ADD:-          if (left->type() == nInt && right->type() == nInt) {-            v.mkInt((left->integer() + right->integer()).valueWrapping());-          } else if (left->type() == nString && right->type() == nString) {-            v.mkString(std::string(left->c_str()) + std::string(right->c_str()));-          } else if (left->type() == nPath && right->type() == nString) {+          if (left.type() == nInt && right.type() == nInt) {+            v.mkInt((left.integer() + right.integer()).valueWrapping());+          } else if (left.type() == nString && right.type() == nString) {+            v.mkString(std::string(left.c_str()) + std::string(right.c_str()));+          } else if (left.type() == nPath && right.type() == nString) {             // Path + string = path-            std::string leftPath = std::string(left->path().path.abs());-            std::string result = leftPath + std::string(right->c_str());+            std::string leftPath = std::string(left.path().path.abs());+            std::string result = leftPath + std::string(right.c_str());             v.mkPath(state.rootPath(CanonPath(result)));-          } else if (left->type() == nString && right->type() == nPath) {+          } else if (left.type() == nString && right.type() == nPath) {             // String + path = path-            std::string rightPath = std::string(right->path().path.abs());-            std::string result = std::string(left->c_str()) + rightPath;+            std::string rightPath = std::string(right.path().path.abs());+            std::string result = std::string(left.c_str()) + rightPath;             v.mkPath(state.rootPath(CanonPath(result)));-          } else if (left->type() == nPath && right->type() == nPath) {+          } else if (left.type() == nPath && right.type() == nPath) {             // Path + path = path-            std::string leftPath = std::string(left->path().path.abs());-            std::string rightPath = std::string(right->path().path.abs());+            std::string leftPath = std::string(left.path().path.abs());+            std::string rightPath = std::string(right.path().path.abs());             std::string result = leftPath + rightPath;             v.mkPath(state.rootPath(CanonPath(result)));           } else {@@ -324,84 +344,84 @@ struct Evaluator::Impl {           }           break;         case BinaryOp::SUB:-          if (left->type() == nInt && right->type() == nInt) {-            v.mkInt((left->integer() - right->integer()).valueWrapping());+          if (left.type() == nInt && right.type() == nInt) {+            v.mkInt((left.integer() - right.integer()).valueWrapping());           } else {             state.error<EvalError>("type error in subtraction").debugThrow();           }           break;         case BinaryOp::MUL:-          if (left->type() == nInt && right->type() == nInt) {-            v.mkInt((left->integer() * right->integer()).valueWrapping());+          if (left.type() == nInt && right.type() == nInt) {+            v.mkInt((left.integer() * right.integer()).valueWrapping());           } else {             state.error<EvalError>("type error in multiplication").debugThrow();           }           break;         case BinaryOp::DIV:-          if (left->type() == nInt && right->type() == nInt) {-            if (right->integer() == NixInt(0)) {+          if (left.type() == nInt && right.type() == nInt) {+            if (right.integer() == NixInt(0)) {               state.error<EvalError>("division by zero").debugThrow();             }-            v.mkInt((left->integer() / right->integer()).valueWrapping());+            v.mkInt((left.integer() / right.integer()).valueWrapping());           } else {             state.error<EvalError>("type error in division").debugThrow();           }           break;         case BinaryOp::EQ:-          v.mkBool(state.eqValues(*left, *right, noPos, ""));+          v.mkBool(state.eqValues(left, right, noPos, ""));           break;         case BinaryOp::NE:-          v.mkBool(!state.eqValues(*left, *right, noPos, ""));+          v.mkBool(!state.eqValues(left, right, noPos, ""));           break;         case BinaryOp::LT:-          if (left->type() == nInt && right->type() == nInt) {-            v.mkBool(left->integer() < right->integer());-          } else if (left->type() == nString && right->type() == nString) {-            v.mkBool(std::string(left->c_str()) < std::string(right->c_str()));+          if (left.type() == nInt && right.type() == nInt) {+            v.mkBool(left.integer() < right.integer());+          } else if (left.type() == nString && right.type() == nString) {+            v.mkBool(std::string(left.c_str()) < std::string(right.c_str()));           } else {             state.error<EvalError>("type error in comparison").debugThrow();           }           break;         case BinaryOp::GT:-          if (left->type() == nInt && right->type() == nInt) {-            v.mkBool(left->integer() > right->integer());-          } else if (left->type() == nString && right->type() == nString) {-            v.mkBool(std::string(left->c_str()) > std::string(right->c_str()));+          if (left.type() == nInt && right.type() == nInt) {+            v.mkBool(left.integer() > right.integer());+          } else if (left.type() == nString && right.type() == nString) {+            v.mkBool(std::string(left.c_str()) > std::string(right.c_str()));           } else {             state.error<EvalError>("type error in comparison").debugThrow();           }           break;         case BinaryOp::LE:-          if (left->type() == nInt && right->type() == nInt) {-            v.mkBool(left->integer() <= right->integer());-          } else if (left->type() == nString && right->type() == nString) {-            v.mkBool(std::string(left->c_str()) <= std::string(right->c_str()));+          if (left.type() == nInt && right.type() == nInt) {+            v.mkBool(left.integer() <= right.integer());+          } else if (left.type() == nString && right.type() == nString) {+            v.mkBool(std::string(left.c_str()) <= std::string(right.c_str()));           } else {             state.error<EvalError>("type error in comparison").debugThrow();           }           break;         case BinaryOp::GE:-          if (left->type() == nInt && right->type() == nInt) {-            v.mkBool(left->integer() >= right->integer());-          } else if (left->type() == nString && right->type() == nString) {-            v.mkBool(std::string(left->c_str()) >= std::string(right->c_str()));+          if (left.type() == nInt && right.type() == nInt) {+            v.mkBool(left.integer() >= right.integer());+          } else if (left.type() == nString && right.type() == nString) {+            v.mkBool(std::string(left.c_str()) >= std::string(right.c_str()));           } else {             state.error<EvalError>("type error in comparison").debugThrow();           }           break;         case BinaryOp::CONCAT: {           // List concatenation: left ++ right-          if (left->type() != nList || right->type() != nList) {+          if (left.type() != nList || right.type() != nList) {             state.error<EvalError>("list concatenation requires two lists").debugThrow();           } -          size_t left_size = left->listSize();-          size_t right_size = right->listSize();+          size_t left_size = left.listSize();+          size_t right_size = right.listSize();           size_t total_size = left_size + right_size;            auto builder = state.buildList(total_size);-          auto left_view = left->listView();-          auto right_view = right->listView();+          auto left_view = left.listView();+          auto right_view = right.listView();            // Copy elements from left list           size_t idx = 0;@@ -419,23 +439,23 @@ struct Evaluator::Impl {         }         case BinaryOp::MERGE: {           // // is attrset merge - right overrides left-          if (left->type() != nAttrs || right->type() != nAttrs) {+          if (left.type() != nAttrs || right.type() != nAttrs) {             state.error<EvalError>("attrset merge requires two attrsets").debugThrow();           }            // Build a map of right attrs first (these have priority)           std::unordered_map<Symbol, Value*> right_attrs;-          for (auto& attr : *right->attrs()) {+          for (auto& attr : *right.attrs()) {             right_attrs[attr.name] = attr.value;           }            // Copy right attrs to result-          auto builder = state.buildBindings(left->attrs()->size() + right->attrs()->size());-          for (auto& attr : *right->attrs()) {+          auto builder = state.buildBindings(left.attrs()->size() + right.attrs()->size());+          for (auto& attr : *right.attrs()) {             builder.insert(attr.name, attr.value);           }           // Add left attrs that don't exist in right-          for (auto& attr : *left->attrs()) {+          for (auto& attr : *left.attrs()) {             if (right_attrs.find(attr.name) == right_attrs.end()) {               builder.insert(attr.name, attr.value);             }@@ -449,21 +469,21 @@ struct Evaluator::Impl {         break;       }     } else if (auto* n = node->get_if<UnaryOpNode>()) {-      Value* operand = state.allocValue();-      eval_node(n->operand, *operand, env);-      force(operand);+      Value operand;+      eval_node(n->operand, operand, env);+      force(&operand);        switch (n->op) {       case UnaryOp::NEG:-        if (operand->type() == nInt) {-          v.mkInt((NixInt(0) - operand->integer()).valueWrapping());+        if (operand.type() == nInt) {+          v.mkInt((NixInt(0) - operand.integer()).valueWrapping());         } else {           state.error<EvalError>("type error in negation").debugThrow();         }         break;       case UnaryOp::NOT:-        if (operand->type() == nBool) {-          v.mkBool(!operand->boolean());+        if (operand.type() == nBool) {+          v.mkBool(!operand.boolean());         } else {           state.error<EvalError>("type error in logical NOT").debugThrow();         }@@ -472,21 +492,21 @@ struct Evaluator::Impl {         state.error<EvalError>("unknown unary operator").debugThrow();       }     } else if (auto* n = node->get_if<IfNode>()) {-      Value* cond = state.allocValue();-      eval_node(n->cond, *cond, env);-      force(cond);+      Value cond;+      eval_node(n->cond, cond, env);+      force(&cond); -      if (cond->type() != nBool) {+      if (cond.type() != nBool) {         state.error<EvalError>("condition must be a boolean").debugThrow();       } -      if (cond->boolean()) {+      if (cond.boolean()) {         eval_node(n->then_branch, v, env);       } else {         eval_node(n->else_branch, v, env);       }     } else if (auto* n = node->get_if<LetNode>()) {-      auto let_env = make_env(env);+      auto let_env = make_env(env, n->bindings.size());       // Nix's let is recursive: bind all names first, then evaluate       // We allocate Values immediately and evaluate into them       std::vector<Value*> values;@@ -502,7 +522,7 @@ struct Evaluator::Impl {       }       eval_node(n->body, v, let_env);     } else if (auto* n = node->get_if<LetRecNode>()) {-      auto letrec_env = make_env(env);+      auto letrec_env = make_env(env, n->bindings.size());       // Same as LetNode - both are recursive in Nix       std::vector<Value*> values;       for (const auto& [name, expr] : n->bindings) {@@ -522,7 +542,7 @@ struct Evaluator::Impl {       if (n->recursive) {         // For recursive attrsets, create environment where all bindings can         // see each other-        attr_env = make_env(env);+        attr_env = make_env(env, n->attrs.size());         for (const auto& binding : n->attrs) {           if (!binding.is_dynamic()) {             Value* thunk = make_thunk(binding.value, attr_env);@@ -557,24 +577,24 @@ struct Evaluator::Impl {        v.mkAttrs(bindings.finish());     } else if (auto* n = node->get_if<SelectNode>()) {-      Value* obj = state.allocValue();-      eval_node(n->expr, *obj, env);-      force(obj);+      Value obj;+      eval_node(n->expr, obj, env);+      force(&obj); -      Value* attr_val = state.allocValue();-      eval_node(n->attr, *attr_val, env);-      force(attr_val);+      Value attr_val;+      eval_node(n->attr, attr_val, env);+      force(&attr_val); -      if (obj->type() != nAttrs) {+      if (obj.type() != nAttrs) {         state.error<EvalError>("selection on non-attrset").debugThrow();       } -      if (attr_val->type() != nString) {+      if (attr_val.type() != nString) {         state.error<EvalError>("attribute name must be string").debugThrow();       } -      auto sym = state.symbols.create(attr_val->c_str());-      auto attr = obj->attrs()->get(sym);+      auto sym = state.symbols.create(attr_val.c_str());+      auto attr = obj.attrs()->get(sym);        if (attr) {         copy_value(v, attr->value);@@ -584,21 +604,21 @@ struct Evaluator::Impl {         state.error<EvalError>("attribute not found").debugThrow();       }     } else if (auto* n = node->get_if<HasAttrNode>()) {-      Value* obj = state.allocValue();-      eval_node(n->expr, *obj, env);-      force(obj);+      Value obj;+      eval_node(n->expr, obj, env);+      force(&obj); -      Value* attr_val = state.allocValue();-      eval_node(n->attr, *attr_val, env);-      force(attr_val);+      Value attr_val;+      eval_node(n->attr, attr_val, env);+      force(&attr_val); -      if (obj->type() != nAttrs) {+      if (obj.type() != nAttrs) {         v.mkBool(false);-      } else if (attr_val->type() != nString) {+      } else if (attr_val.type() != nString) {         state.error<EvalError>("attribute name must be string").debugThrow();       } else {-        auto sym = state.symbols.create(attr_val->c_str());-        auto attr = obj->attrs()->get(sym);+        auto sym = state.symbols.create(attr_val.c_str());+        auto attr = obj.attrs()->get(sym);         v.mkBool(attr != nullptr);       }     } else if (auto* n = node->get_if<WithNode>()) {@@ -614,30 +634,30 @@ struct Evaluator::Impl {       with_env->with_attrs = attrs;       eval_node(n->body, v, with_env);     } else if (auto* n = node->get_if<AssertNode>()) {-      Value* cond = state.allocValue();-      eval_node(n->cond, *cond, env);-      force(cond);+      Value cond;+      eval_node(n->cond, cond, env);+      force(&cond); -      if (cond->type() != nBool) {+      if (cond.type() != nBool) {         state.error<EvalError>("assertion must be boolean").debugThrow();       } -      if (!cond->boolean()) {+      if (!cond.boolean()) {         state.error<EvalError>("assertion failed").debugThrow();       }        eval_node(n->body, v, env);     } else if (auto* n = node->get_if<ImportNode>()) {       // Evaluate path expression to get the file path-      Value* path_val = state.allocValue();-      eval_node(n->path, *path_val, env);-      force(path_val);+      Value path_val;+      eval_node(n->path, path_val, env);+      force(&path_val);        // Path should be a string or path type, convert to SourcePath-      if (path_val->type() == nPath) {-        state.evalFile(path_val->path(), v);-      } else if (path_val->type() == nString) {-        auto path = state.rootPath(CanonPath(path_val->c_str()));+      if (path_val.type() == nPath) {+        state.evalFile(path_val.path(), v);+      } else if (path_val.type() == nString) {+        auto path = state.rootPath(CanonPath(path_val.c_str()));         state.evalFile(path, v);       } else {         state.error<EvalError>("import argument must be a path or string").debugThrow();diff --git a/src/irc/main.cpp b/src/irc/main.cppindex 3c00d70..828e91c 100644--- a/src/irc/main.cpp+++ b/src/irc/main.cpp@@ -2,6 +2,7 @@ #include "parser.h" #include "serializer.h" #include <cctype>+#include <chrono> #include <cstring> #include <filesystem> #include <iostream>@@ -12,6 +13,12 @@ namespace nix_irc { namespace fs = std::filesystem; +using Clock = std::chrono::steady_clock;++static long long elapsed_us(Clock::time_point start, Clock::time_point end) {+  return std::chrono::duration_cast<std::chrono::microseconds>(end - start).count();+}+ void print_usage(const char* prog) {   std::cout << "Usage: " << prog << " [options] <input.nix|flake#attr> [output.nixir]\n"             << "\nOptions:\n"@@ -204,11 +211,13 @@ int run_compile(int argc, char** argv) {   }    try {+    auto t_start = Clock::now();     Parser parser;     (void) search_paths;     (void) resolve_imports;      std::shared_ptr<Node> ast;+    long long resolve_us = 0;      if (is_flake_reference(input_file)) {       std::cout << "Compiling flake reference: " << input_file << "\n";@@ -218,6 +227,8 @@ int run_compile(int argc, char** argv) {       ast = parser.parse_file(input_file);     } +    auto t_parsed = Clock::now();+     if (!ast) {       std::cerr << "Error: Failed to parse input\n";       return 1;@@ -229,6 +240,7 @@ int run_compile(int argc, char** argv) {      std::cout << "Generating IR...\n";     auto ir = ir_gen.generate(ast);+    auto t_ir_generated = Clock::now();      IRModule module;     module.version = IR_VERSION;@@ -237,8 +249,17 @@ int run_compile(int argc, char** argv) {     std::cout << "Serializing to: " << output_file << "\n";     Serializer serializer;     serializer.serialize(module, output_file);+    auto t_serialized = Clock::now();      std::cout << "Done!\n";+    auto parse_us = elapsed_us(t_start, t_parsed);+    auto ir_gen_us = elapsed_us(t_parsed, t_ir_generated) - resolve_us;+    auto serialize_us = elapsed_us(t_ir_generated, t_serialized);+    auto total_us = elapsed_us(t_start, t_serialized);++    std::cerr << "nixirc timing: parse=" << parse_us << "us resolve=" << resolve_us+              << "us irgen=" << ir_gen_us << "us serialize=" << serialize_us+              << "us total=" << total_us << "us" << std::endl;     return 0;    } catch (const std::exception& e) {