Add builtin function "partition"

The implementation of "partition" in Nixpkgs is O(n^2) (because of the
use of ++), and for some reason was causing stack overflows in
multi-threaded evaluation (not sure why).

This reduces "nix-env -qa --drv-path" runtime by 0.197s and memory
usage by 298 MiB (in non-Boehm mode).

7 files changed, 54 insertions, 10 deletions

diff --git a/src/libexpr/eval.cc b/src/libexpr/eval.cc
index ace0517fbf07..22396865f095 100644
--- a/src/libexpr/eval.cc
+++ b/src/libexpr/eval.cc
@@ -293,6 +293,8 @@ EvalState::EvalState(const Strings & _searchPath, ref<Store> store)
     , sColumn(symbols.create("column"))
     , sFunctor(symbols.create("__functor"))
     , sToString(symbols.create("__toString"))
+    , sRight(symbols.create("right"))
+    , sWrong(symbols.create("wrong"))
     , store(store)
     , baseEnv(allocEnv(128))
     , staticBaseEnv(false, 0)
diff --git a/src/libexpr/eval.hh b/src/libexpr/eval.hh
index 47e4d99bf6ed..6f1befa2a23e 100644
--- a/src/libexpr/eval.hh
+++ b/src/libexpr/eval.hh
@@ -71,7 +71,8 @@ public:
 
     const Symbol sWith, sOutPath, sDrvPath, sType, sMeta, sName, sValue,
         sSystem, sOverrides, sOutputs, sOutputName, sIgnoreNulls,
-        sFile, sLine, sColumn, sFunctor, sToString;
+        sFile, sLine, sColumn, sFunctor, sToString,
+        sRight, sWrong;
     Symbol sDerivationNix;
 
     /* If set, force copying files to the Nix store even if they
diff --git a/src/libexpr/json-to-value.cc b/src/libexpr/json-to-value.cc
index 1daf84600dca..f671802bcc24 100644
--- a/src/libexpr/json-to-value.cc
+++ b/src/libexpr/json-to-value.cc
@@ -12,15 +12,6 @@ static void skipWhitespace(const char * & s)
 }
 
 
-#if HAVE_BOEHMGC
-typedef std::vector<Value *, gc_allocator<Value *> > ValueVector;
-typedef std::map<Symbol, Value *, std::less<Symbol>, gc_allocator<Value *> > ValueMap;
-#else
-typedef std::vector<Value *> ValueVector;
-typedef std::map<Symbol, Value *> ValueMap;
-#endif
-
-
 static string parseJSONString(const char * & s)
 {
     string res;
diff --git a/src/libexpr/primops.cc b/src/libexpr/primops.cc
index 51ee16353ef6..5c0b77c15e78 100644
--- a/src/libexpr/primops.cc
+++ b/src/libexpr/primops.cc
@@ -1443,6 +1443,40 @@ static void prim_sort(EvalState & state, const Pos & pos, Value * * args, Value
 }
 
 
+static void prim_partition(EvalState & state, const Pos & pos, Value * * args, Value & v)
+{
+    state.forceFunction(*args[0], pos);
+    state.forceList(*args[1], pos);
+
+    auto len = args[1]->listSize();
+
+    ValueVector right, wrong;
+
+    for (unsigned int n = 0; n < len; ++n) {
+        auto vElem = args[1]->listElems()[n];
+        state.forceValue(*vElem);
+        Value res;
+        state.callFunction(*args[0], *vElem, res, pos);
+        if (state.forceBool(res))
+            right.push_back(vElem);
+        else
+            wrong.push_back(vElem);
+    }
+
+    state.mkAttrs(v, 2);
+
+    Value * vRight = state.allocAttr(v, state.sRight);
+    state.mkList(*vRight, right.size());
+    memcpy(vRight->listElems(), right.data(), sizeof(Value *) * right.size());
+
+    Value * vWrong = state.allocAttr(v, state.sWrong);
+    state.mkList(*vWrong, wrong.size());
+    memcpy(vWrong->listElems(), wrong.data(), sizeof(Value *) * wrong.size());
+
+    v.attrs->sort();
+}
+
+
 /*************************************************************
  * Integer arithmetic
  *************************************************************/
@@ -1881,6 +1915,7 @@ void EvalState::createBaseEnv()
     addPrimOp("__all", 2, prim_all);
     addPrimOp("__genList", 2, prim_genList);
     addPrimOp("__sort", 2, prim_sort);
+    addPrimOp("__partition", 2, prim_partition);
 
     // Integer arithmetic
     addPrimOp("__add", 2, prim_add);
diff --git a/src/libexpr/value.hh b/src/libexpr/value.hh
index f5e485748812..048522016c62 100644
--- a/src/libexpr/value.hh
+++ b/src/libexpr/value.hh
@@ -250,4 +250,13 @@ void mkPath(Value & v, const char * s);
 size_t valueSize(Value & v);
 
 
+#if HAVE_BOEHMGC
+typedef std::vector<Value *, gc_allocator<Value *> > ValueVector;
+typedef std::map<Symbol, Value *, std::less<Symbol>, gc_allocator<Value *> > ValueMap;
+#else
+typedef std::vector<Value *> ValueVector;
+typedef std::map<Symbol, Value *> ValueMap;
+#endif
+
+
 }
diff --git a/tests/lang/eval-okay-partition.exp b/tests/lang/eval-okay-partition.exp
new file mode 100644
index 000000000000..cd8b8b020c05
--- /dev/null
+++ b/tests/lang/eval-okay-partition.exp
@@ -0,0 +1 @@
+{ right = [ 0 2 4 6 8 10 100 102 104 106 108 110 ]; wrong = [ 1 3 5 7 9 101 103 105 107 109 ]; }
diff --git a/tests/lang/eval-okay-partition.nix b/tests/lang/eval-okay-partition.nix
new file mode 100644
index 000000000000..846d2ce49486
--- /dev/null
+++ b/tests/lang/eval-okay-partition.nix
@@ -0,0 +1,5 @@
+with import ./lib.nix;
+
+builtins.partition
+  (x: x / 2 * 2 == x)
+  (builtins.concatLists [ (range 0 10) (range 100 110) ])