From 50807f3dd5241667dac0c0cc302042d648de4b42 Mon Sep 17 00:00:00 2001
From: Eelco Dolstra <eelco.dolstra@logicblox.com>
Date: Tue, 28 Jul 2015 17:27:32 +0200
Subject: [PATCH] Add primop genList
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This can be used to implement functions like ‘imap’ (or for that
matter, ‘map’) without the quadratic complexity incurred by calling
‘++’ repeatedly.
---
 doc/manual/expressions/builtins.xml | 18 ++++++++++++++++++
 src/libexpr/primops.cc              | 20 ++++++++++++++++++++
 tests/lang/lib.nix                  |  6 +++---
 3 files changed, 41 insertions(+), 3 deletions(-)

diff --git a/doc/manual/expressions/builtins.xml b/doc/manual/expressions/builtins.xml
index 099f90e78..446157056 100644
--- a/doc/manual/expressions/builtins.xml
+++ b/doc/manual/expressions/builtins.xml
@@ -364,6 +364,24 @@ builtins.fromJSON ''{"x": [1, 2, 3], "y": null}''
   </varlistentry>
 
 
+  <varlistentry><term><function>builtins.genList</function>
+  <replaceable>generator</replaceable> <replaceable>length</replaceable></term>
+
+    <listitem><para>Generate list of size
+    <replaceable>length</replaceable>, with each element
+    <replaceable>i></replaceable> equal to the value returned by
+    <replaceable>generator</replaceable> <literal>i</literal>. For
+    example,
+
+<programlisting>
+builtins.genList (x: x * x) 5
+</programlisting>
+
+    returns the list <literal>[ 0 1 4 9 16 ]</literal>.</para></listitem>
+
+  </varlistentry>
+
+
   <varlistentry><term><function>builtins.getAttr</function>
   <replaceable>s</replaceable> <replaceable>set</replaceable></term>
 
diff --git a/src/libexpr/primops.cc b/src/libexpr/primops.cc
index ece76975d..5da2f3463 100644
--- a/src/libexpr/primops.cc
+++ b/src/libexpr/primops.cc
@@ -1364,6 +1364,25 @@ static void prim_all(EvalState & state, const Pos & pos, Value * * args, Value &
 }
 
 
+/* Apply a function to every element of a list. */
+static void prim_genList(EvalState & state, const Pos & pos, Value * * args, Value & v)
+{
+    state.forceFunction(*args[0], pos);
+    auto len = state.forceInt(*args[1], pos);
+
+    if (len < 0)
+        throw EvalError(format("cannot create list of size %1%, at %2%") % len % pos);
+
+    state.mkList(v, len);
+
+    for (unsigned int n = 0; n < len; ++n) {
+        Value * arg = state.allocValue();
+        mkInt(*arg, n);
+        mkApp(*(v.listElems()[n] = state.allocValue()), *args[0], *arg);
+    }
+}
+
+
 /*************************************************************
  * Integer arithmetic
  *************************************************************/
@@ -1759,6 +1778,7 @@ void EvalState::createBaseEnv()
     addPrimOp("__foldl'", 3, prim_foldlStrict);
     addPrimOp("__any", 2, prim_any);
     addPrimOp("__all", 2, prim_all);
+    addPrimOp("__genList", 2, prim_genList);
 
     // Integer arithmetic
     addPrimOp("__add", 2, prim_add);
diff --git a/tests/lang/lib.nix b/tests/lang/lib.nix
index 262cdd7e8..028a53831 100644
--- a/tests/lang/lib.nix
+++ b/tests/lang/lib.nix
@@ -54,8 +54,8 @@ rec {
   const = x: y: x;
 
   range = first: last:
-    if builtins.lessThan last first
-    then []
-    else [first] ++ range (builtins.add first 1) last;
+    if first > last
+      then []
+      else genList (n: first + n) (last - first + 1);
 
 }