lib.concatMap and lib.mapAttrs to be builtins

5 files changed, 52 insertions, 0 deletions

diff --git a/src/libexpr/primops.cc b/src/libexpr/primops.cc
index 3a6c4035b8..7ff57f0b01 100644
--- a/src/libexpr/primops.cc
+++ b/src/libexpr/primops.cc
@@ -1356,6 +1356,24 @@ static void prim_functionArgs(EvalState & state, const Pos & pos, Value * * args
 }
 
 
+/* Apply a function to every element of a list. */
+static void prim_mapAttrs(EvalState & state, const Pos & pos, Value * * args, Value & v)
+{
+    state.forceFunction(*args[0], pos);
+    state.forceAttrs(*args[1], pos);
+
+    state.mkAttrs(v, args[1]->attrs->size());
+
+    for (auto & i : *args[1]->attrs) {
+        Value vName, vFun2;
+        mkString(vName, i.name);
+        state.callFunction(*args[0], vName, vFun2, pos);
+        state.callFunction(vFun2, *i.value, *state.allocAttr(v, i.name), pos);
+    }
+}
+
+
+
 /*************************************************************
  * Lists
  *************************************************************/
@@ -1627,6 +1645,28 @@ static void prim_partition(EvalState & state, const Pos & pos, Value * * args, V
 }
 
 
+/* concatMap = f: list: concatLists (map f list); */
+/* C++-version is to avoid allocating `mkApp', call `f' eagerly */
+static void prim_concatMap(EvalState & state, const Pos & pos, Value * * args, Value & v)
+{
+    state.forceFunction(*args[0], pos);
+    state.forceList(*args[1], pos);
+    auto len = args[1]->listSize();
+
+    Value vList;
+    state.mkList(vList, len);
+
+    for (unsigned int n = 0; n < len; ++n) {
+        Value * vElem = args[1]->listElems()[n];
+        state.forceValue(*vElem);
+        state.callFunction(*args[0], *vElem, *(vList.listElems()[n] = state.allocValue()), pos);
+    }
+
+    state.concatLists(v, len, vList.listElems(), pos);
+}
+
+
+
 /*************************************************************
  * Integer arithmetic
  *************************************************************/
@@ -2212,6 +2252,7 @@ void EvalState::createBaseEnv()
     addPrimOp("__intersectAttrs", 2, prim_intersectAttrs);
     addPrimOp("__catAttrs", 2, prim_catAttrs);
     addPrimOp("__functionArgs", 1, prim_functionArgs);
+    addPrimOp("__mapAttrs", 2, prim_mapAttrs);
 
     // Lists
     addPrimOp("__isList", 1, prim_isList);
@@ -2229,6 +2270,7 @@ void EvalState::createBaseEnv()
     addPrimOp("__genList", 2, prim_genList);
     addPrimOp("__sort", 2, prim_sort);
     addPrimOp("__partition", 2, prim_partition);
+    addPrimOp("__concatMap", 2, prim_concatMap);
 
     // Integer arithmetic
     addPrimOp("__add", 2, prim_add);
diff --git a/tests/lang/eval-okay-concatmap.exp b/tests/lang/eval-okay-concatmap.exp
new file mode 100644
index 0000000000..3b8be7739d
--- /dev/null
+++ b/tests/lang/eval-okay-concatmap.exp
@@ -0,0 +1 @@
+[ [ 1 3 5 7 9 ] [ "a" "z" "b" "z" ] ]
diff --git a/tests/lang/eval-okay-concatmap.nix b/tests/lang/eval-okay-concatmap.nix
new file mode 100644
index 0000000000..97da5d37a4
--- /dev/null
+++ b/tests/lang/eval-okay-concatmap.nix
@@ -0,0 +1,5 @@
+with import ./lib.nix;
+
+[ (builtins.concatMap (x: if x / 2 * 2 == x then [] else [ x ]) (range 0 10))
+  (builtins.concatMap (x: [x] ++ ["z"]) ["a" "b"])
+]
diff --git a/tests/lang/eval-okay-mapattrs.exp b/tests/lang/eval-okay-mapattrs.exp
new file mode 100644
index 0000000000..3f113f17ba
--- /dev/null
+++ b/tests/lang/eval-okay-mapattrs.exp
@@ -0,0 +1 @@
+{ x = "x-foo"; y = "y-bar"; }
diff --git a/tests/lang/eval-okay-mapattrs.nix b/tests/lang/eval-okay-mapattrs.nix
new file mode 100644
index 0000000000..f075b6275e
--- /dev/null
+++ b/tests/lang/eval-okay-mapattrs.nix
@@ -0,0 +1,3 @@
+with import ./lib.nix;
+
+builtins.mapAttrs (name: value: name + "-" + value) { x = "foo"; y = "bar"; }