libexpr: speculatively pre-force list elements on the worker pool

src/libexpr-tests/meson.build

@@ -57,6 +57,7 @@ sources = files(
   'nix_api_external.cc',
   'nix_api_value.cc',
   'nix_api_value_internal.cc',
+  'parallel-eval.cc',
   'primops.cc',
   'search-path.cc',
   'trivial.cc',

src/libexpr-tests/parallel-eval.cc

@@ -0,0 +1,63 @@
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+
+#include "nix/expr/tests/libexpr.hh"
+#include "nix/expr/parallel-eval.hh"
+
+namespace nix {
+
+class ParallelEvalTest : public LibExprTest
+{
+public:
+    ParallelEvalTest()
+        : LibExprTest(openStore("dummy://"), [](bool & readOnlyMode) {
+            EvalSettings settings{readOnlyMode};
+            settings.nixPath = {};
+            settings.evalCores = 4;
+            return settings;
+        })
+    {
+    }
+};
+
+TEST_F(ParallelEvalTest, executorEnabled)
+{
+    ASSERT_TRUE(state.executor->enabled);
+    ASSERT_EQ(state.executor->evalCores, 4u);
+}
+
+TEST_F(ParallelEvalTest, concatStringsSepThunks)
+{
+    auto v = eval("builtins.concatStringsSep \",\" (builtins.genList (i: builtins.toString (i + 1)) 32)");
+    ASSERT_THAT(
+        v, IsStringEq("1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32"));
+}
+
+TEST_F(ParallelEvalTest, concatStringsSepReportsFirstThrow)
+{
+    try {
+        eval("builtins.concatStringsSep \"\" (builtins.genList (i: throw \"e${builtins.toString i}\") 16)");
+        FAIL() << "expected ThrownError";
+    } catch (const ThrownError & e) {
+        ASSERT_THAT(e.what(), testing::HasSubstr("e0"));
+    }
+}
+
+TEST_F(ParallelEvalTest, toStringNestedListThunks)
+{
+    auto v = eval("builtins.toString (builtins.genList (i: builtins.genList (j: i * 4 + j + 1) 4) 4)");
+    ASSERT_THAT(v, IsStringEq("1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16"));
+}
+
+TEST_F(ParallelEvalTest, preForceListElementsDefaultForce)
+{
+    auto v = eval("builtins.genList (i: i + 1) 8", false);
+    ASSERT_EQ(v.type(), nList);
+    state.preForceListElements(v, noPos);
+    for (auto elem : v.listView()) {
+        state.forceValue(*elem, noPos);
+        ASSERT_EQ(elem->type(), nInt);
+    }
+}
+
+} // namespace nix

src/libexpr/eval.cc

@@ -2456,6 +2456,61 @@ EvalState::tryAttrsToString(const PosIdx pos, Value & v, NixStringContext & cont
     return {};
 }
 
+void EvalState::preForceListElements(Value & v, const PosIdx pos, std::function<void(Value &)> consume)
+{
+    if (!executor->enabled || Executor::amWorkerThread || !v.isList() || v.listSize() < 2)
+        return;
+
+    std::vector<Value *> pending;
+    pending.reserve(v.listSize());
+    for (auto v2 : v.listView()) {
+        if (!v2->isFinished())
+            pending.push_back(v2);
+        else if (consume) {
+            /* Finished WHNF still has work if consume walks into it. */
+            switch (v2->type()) {
+            case nAttrs:
+            case nList:
+            case nPath:
+                pending.push_back(v2);
+                break;
+            default:
+                break;
+            }
+        }
+    }
+    if (pending.size() < 2)
+        return;
+
+    auto sharedConsume = std::make_shared<std::function<void(Value &)>>(std::move(consume));
+    /* At most 4 work items per core so a long list of cheap elements
+       does not contend on the executor queue. */
+    size_t nChunks = std::min<size_t>(pending.size(), size_t(executor->evalCores) * 4);
+    size_t chunkSize = (pending.size() + nChunks - 1) / nChunks;
+
+    Executor::WorkItems work;
+    work.reserve(nChunks);
+    for (size_t i = 0; i < pending.size(); i += chunkSize) {
+        std::vector<RootValue> slice;
+        size_t n = std::min(chunkSize, pending.size() - i);
+        slice.reserve(n);
+        for (size_t j = 0; j < n; ++j)
+            slice.push_back(RootValue(pending[i + j]));
+        addWork(work, 0, [slice = std::move(slice), sharedConsume, this, pos]() {
+            for (auto & rv : slice) {
+                try {
+                    if (*sharedConsume)
+                        (*sharedConsume)(**rv);
+                    else
+                        forceValue(**rv, pos);
+                } catch (const Error &) {
+                }
+            }
+        });
+    }
+    executor->spawn(std::move(work));
+}
+
 BackedStringView EvalState::coerceToString(
     const PosIdx pos,
     Value & v,
@@ -2536,6 +2591,19 @@ BackedStringView EvalState::coerceToString(
         if (v.isList()) {
             std::string result;
             auto listView = v.listView();
+
+            preForceListElements(v, pos, [this, pos, coerceMore, canonicalizePath](Value & elem) {
+                NixStringContext scratch;
+                coerceToString(
+                    pos,
+                    elem,
+                    scratch,
+                    "while evaluating one element of the list",
+                    coerceMore,
+                    false,
+                    canonicalizePath);
+            });
+
             for (auto [n, v2] : enumerate(listView)) {
                 try {
                     result += *coerceToString(

src/libexpr/include/nix/expr/eval-settings.hh

@@ -503,13 +503,14 @@ public:
         1,
         "eval-cores",
         R"(
-          The number of threads used to evaluate Nix expressions. This currently affects the following commands:
+          The number of threads used to evaluate Nix expressions. This currently affects:
 
           * `nix search`
           * `nix flake check`
           * `nix flake show`
           * `nix eval --json`
-          * Any evaluation that uses `builtins.parallel`
+          * `builtins.parallel`
+          * `builtins.concatStringsSep` and string coercion of lists
 
           The value `0` causes Nix to use all available CPU cores in the system.
 

src/libexpr/include/nix/expr/eval.hh

@@ -728,6 +728,28 @@ public:
 
     std::string devirtualize(std::string_view s, const NixStringContext & context);
 
+    /**
+     * Run `consume` on the elements of list `v` on the worker pool.
+     * Returns without waiting; the caller then walks the list sequentially
+     * to build the result and to report errors.
+     *
+     * No-op unless the parallel evaluator is enabled, the caller is not
+     * already a worker thread, and at least two elements need work.
+     *
+     * `consume` may run concurrently with the caller and after this
+     * function returns. It must not capture pointers into the caller's
+     * stack. It may only force values and do other memoized, thread-safe
+     * work (for example `coerceToString` into a scratch context with
+     * `copyToStore = false`). It must not write the caller's string
+     * context.
+     *
+     * `Error` thrown by `consume` is discarded; the sequential walk
+     * rethrows it. `Interrupted` propagates and stops the worker pool.
+     *
+     * If `consume` is empty, elements are forced to weak-head normal form.
+     */
+    void preForceListElements(Value & v, const PosIdx pos, std::function<void(Value &)> consume = {});
+
     /**
      * String coercion.
      *

src/libexpr/primops.cc

@@ -5226,6 +5226,17 @@ static void prim_concatStringsSep(EvalState & state, const PosIdx pos, Value **
         pos,
         "while evaluating the second argument (the list of strings to concat) passed to builtins.concatStringsSep");
 
+    state.preForceListElements(*args[1], pos, [&state, pos](Value & elem) {
+        NixStringContext scratch;
+        state.coerceToString(
+            pos,
+            elem,
+            scratch,
+            "while evaluating one element of the list of strings to concat passed to builtins.concatStringsSep",
+            false,
+            false);
+    });
+
     std::string res;
     res.reserve((args[1]->listSize() + 32) * sep.size());
     bool first = true;