Add memoise primop

'builtins.memoise f x' is equal to 'f x', but uses a cache to speed up repeated invocations of 'f' with the same 'x'. A typical application is memoising evaluations of Nixpkgs in NixOps network specifications. For example, with the patch below, the time to evaluate a 10-machine NixOps network is reduced from 17.1s to 9.6s, while memory consumption goes from 4486 MiB to 3089 MiB. (This is with GC_INITIAL_HEAP_SIZE=16G.) Nixpkgs patch: diff --git a/pkgs/top-level/impure.nix b/pkgs/top-level/impure.nix index a9f21e45aed..f641067e022 100644 --- a/pkgs/top-level/impure.nix +++ b/pkgs/top-level/impure.nix @@ -79,7 +79,7 @@ in # not be passed. assert args ? localSystem -> !(args ? system || args ? platform); -import ./. (builtins.removeAttrs args [ "system" "platform" ] // { +builtins.memoise or (x: x) (import ./.) (builtins.removeAttrs args [ "system" "platform" ] // { inherit config overlays crossSystem; # Fallback: Assume we are building packages on the current (build, in GNU # Autotools parlance) system.
NixOS · Jan 11, 2018 · 0395b9b · 0395b9b · copumpkin · Feb 13, 2018
1 parent 435ccc7
commit 0395b9b
Show file tree

Hide file tree

Showing 3 changed files with 116 additions and 1 deletion.
diff --git a/src/libexpr/eval.cc b/src/libexpr/eval.cc
@@ -1683,6 +1683,8 @@ void EvalState::printStats()
     printMsg(v, format("  number of primop calls: %1%") % nrPrimOpCalls);
     printMsg(v, format("  number of function calls: %1%") % nrFunctionCalls);
     printMsg(v, format("  total allocations: %1% bytes") % (bEnvs + bLists + bValues + bAttrsets));
+    printMsg(v, format("  memoisation hits: %d") % nrMemoiseHits);
+    printMsg(v, format("  memoisation misses: %d") % nrMemoiseMisses);
 
 #if HAVE_BOEHMGC
     GC_word heapSize, totalBytes;

diff --git a/src/libexpr/eval.hh b/src/libexpr/eval.hh
@@ -89,7 +89,7 @@ private:
 
     /* A cache from path names to values. */
 #if HAVE_BOEHMGC
-    typedef std::map<Path, Value, std::less<Path>, traceable_allocator<std::pair<const Path, Value> > > FileEvalCache;
+    typedef std::map<Path, Value, std::less<Path>, traceable_allocator<std::pair<const Path, Value>>> FileEvalCache;
 #else
     typedef std::map<Path, Value> FileEvalCache;
 #endif
@@ -269,6 +269,8 @@ private:
     unsigned long nrListConcats = 0;
     unsigned long nrPrimOpCalls = 0;
     unsigned long nrFunctionCalls = 0;
+    unsigned long nrMemoiseHits = 0;
+    unsigned long nrMemoiseMisses = 0;
 
     bool countCalls;
 
@@ -287,6 +289,22 @@ private:
     friend struct ExprOpConcatLists;
     friend struct ExprSelect;
     friend void prim_getAttr(EvalState & state, const Pos & pos, Value * * args, Value & v);
+    friend void prim_memoise(EvalState & state, const Pos & pos, Value * * args, Value & v);
+
+    /* State for builtins.memoise. */
+    struct MemoArgComparator
+    {
+        EvalState & state;
+        MemoArgComparator(EvalState & state) : state(state) { }
+        bool operator()(Value * a, Value * b);
+    };
+
+    typedef std::map<Value *, Value, MemoArgComparator, traceable_allocator<std::pair<const Value *, Value>>> PerLambdaMemo;
+
+    typedef std::pair<Env *, ExprLambda *> LambdaKey;
+
+    // FIXME: use std::unordered_map
+    std::map<LambdaKey, PerLambdaMemo, std::less<LambdaKey>, traceable_allocator<std::pair<const LambdaKey, PerLambdaMemo>>> memos;
 };
 
 

diff --git a/src/libexpr/primops/memoise.cc b/src/libexpr/primops/memoise.cc
@@ -0,0 +1,95 @@
+#include "primops.hh"
+#include "eval-inline.hh"
+
+#include <cstring>
+
+namespace nix {
+
+bool EvalState::MemoArgComparator::operator()(Value * v1, Value * v2)
+{
+    if (v1 == v2) return false;
+
+    state.forceValue(*v1);
+    state.forceValue(*v2);
+
+    if (v1->type == v2->type) {
+        switch (v1->type) {
+
+        case tInt:
+            return v1->integer < v2->integer;
+
+        case tBool:
+            return v1->boolean < v2->boolean;
+
+        case tString:
+            return strcmp(v1->string.s, v2->string.s);
+
+        case tPath:
+            return strcmp(v1->path, v2->path);
+
+        case tNull:
+            return false;
+
+        case tList1:
+        case tList2:
+        case tListN:
+            unsigned int n;
+            for (n = 0; n < v1->listSize() && n < v2->listSize(); ++n) {
+                if ((*this)(v1->listElems()[n], v2->listElems()[n])) return true;
+                if ((*this)(v2->listElems()[n], v1->listElems()[n])) return false;
+            }
+
+            return n == v1->listSize() && n < v2->listSize();
+
+        case tAttrs:
+            Bindings::iterator i, j;
+            for (i = v1->attrs->begin(), j = v2->attrs->begin(); i != v1->attrs->end() && j != v2->attrs->end(); ++i, ++j) {
+                if (i->name < j->name) return true;
+                if (j->name < i->name) return false;
+                if ((*this)(i->value, j->value)) return true;
+                if ((*this)(j->value, i->value)) return false;
+            }
+            return i == v1->attrs->end() && j != v2->attrs->end();
+
+        case tLambda:
+            return std::make_pair(v1->lambda.env, v1->lambda.fun) < std::make_pair(v2->lambda.env, v2->lambda.fun);
+
+        case tFloat:
+            return v1->fpoint < v2->fpoint;
+
+        default:
+            break;
+        }
+    }
+
+    // As a fallback, use pointer equality.
+    return v1 < v2;
+}
+
+void prim_memoise(EvalState & state, const Pos & pos, Value * * args, Value & v)
+{
+    state.forceFunction(*args[0], pos);
+
+    EvalState::PerLambdaMemo foo(state);
+
+    auto & memo = state.memos.emplace(std::make_pair(args[0]->lambda.env, args[0]->lambda.fun), state).first->second;
+
+    auto result = memo.find(args[1]);
+
+    if (result != memo.end()) {
+        state.nrMemoiseHits++;
+        v = result->second;
+        return;
+    }
+
+    state.nrMemoiseMisses++;
+
+    state.callFunction(*args[0], *args[1], v, pos);
+
+    memo[args[1]] = v;
+
+}
+
+static RegisterPrimOp r("memoise", 2, prim_memoise);
+
+}