add zipAttrsWith primop

nixpkgs can save a good bit of eval memory with this primop. zipAttrsWith is
used quite a bit around nixpkgs (eg in the form of recursiveUpdate), but the
most costly application for this primop is in the module system. it improves
the implementation of zipAttrsWith from nixpkgs by not checking an attribute
multiple times if it occurs more than once in the input list, allocates less
values and set elements, and just avoids many a temporary object in general.

nixpkgs has a more generic version of this operation, zipAttrsWithNames, but
this version is only used once so isn't suitable for being the base of a new
primop. if it were to be used more we should add a second primop instead.

4 files changed, 96 insertions, 0 deletions

diff --git a/doc/manual/src/release-notes/rl-next.md b/doc/manual/src/release-notes/rl-next.md
index 2a67a39b1..a516887b1 100644
--- a/doc/manual/src/release-notes/rl-next.md
+++ b/doc/manual/src/release-notes/rl-next.md
@@ -4,3 +4,5 @@
   more compliant one](https://github.com/ToruNiina/toml11).
 * Added `:st`/`:show-trace` commands to nix repl, which are used to
   set or toggle display of error traces.
+* New builtin function `builtins.zipAttrsWith` with same functionality
+  as `lib.zipAttrsWith` from nixpkgs, but much more efficient.
diff --git a/src/libexpr/primops.cc b/src/libexpr/primops.cc
index 7c964bd0d..ab99962f7 100644
--- a/src/libexpr/primops.cc
+++ b/src/libexpr/primops.cc
@@ -2503,6 +2503,90 @@ static RegisterPrimOp primop_mapAttrs({
     .fun = prim_mapAttrs,
 });
 
+static void prim_zipAttrsWith(EvalState & state, const Pos & pos, Value * * args, Value & v)
+{
+    // we will first count how many values are present for each given key.
+    // we then allocate a single attrset and pre-populate it with lists of
+    // appropriate sizes, stash the pointers to the list elements of each,
+    // and populate the lists. after that we replace the list in the every
+    // attribute with the merge function application. this way we need not
+    // use (slightly slower) temporary storage the GC does not know about.
+
+    std::map<Symbol, std::pair<size_t, Value * *>> attrsSeen;
+
+    state.forceFunction(*args[0], pos);
+    state.forceList(*args[1], pos);
+    const auto listSize = args[1]->listSize();
+    const auto listElems = args[1]->listElems();
+
+    for (unsigned int n = 0; n < listSize; ++n) {
+        Value * vElem = listElems[n];
+        try {
+            state.forceAttrs(*vElem);
+            for (auto & attr : *vElem->attrs)
+                attrsSeen[attr.name].first++;
+        } catch (TypeError & e) {
+            e.addTrace(pos, hintfmt("while invoking '%s'", "zipAttrsWith"));
+            throw;
+        }
+    }
+
+    state.mkAttrs(v, attrsSeen.size());
+    for (auto & [sym, elem] : attrsSeen) {
+        Value * list = state.allocAttr(v, sym);
+        state.mkList(*list, elem.first);
+        elem.second = list->listElems();
+    }
+
+    for (unsigned int n = 0; n < listSize; ++n) {
+        Value * vElem = listElems[n];
+        for (auto & attr : *vElem->attrs)
+            *attrsSeen[attr.name].second++ = attr.value;
+    }
+
+    for (auto & attr : *v.attrs) {
+        Value * name = state.allocValue();
+        mkString(*name, attr.name);
+        Value * call1 = state.allocValue();
+        mkApp(*call1, *args[0], *name);
+        Value * call2 = state.allocValue();
+        mkApp(*call2, *call1, *attr.value);
+        attr.value = call2;
+    }
+}
+
+static RegisterPrimOp primop_zipAttrsWith({
+    .name = "__zipAttrsWith",
+    .args = {"f", "list"},
+    .doc = R"(
+      Transpose a list of attribute sets into an attribute set of lists,
+      then apply `mapAttrs`.
+
+      `f` receives two arguments: the attribute name and a non-empty
+      list of all values encountered for that attribute name.
+
+      The result is an attribute set where the attribute names are the
+      union of the attribute names in each element of `list`. The attribute
+      values are the return values of `f`.
+
+      ```nix
+      builtins.zipAttrsWith
+        (name: values: { inherit name values; })
+        [ { a = "x"; } { a = "y"; b = "z"; } ]
+      ```
+
+      evaluates to
+
+      ```
+      {
+        a = { name = "a"; values = [ "x" "y" ]; };
+        b = { name = "b"; values = [ "z" ]; };
+      }
+      ```
+    )",
+    .fun = prim_zipAttrsWith,
+});
+
 
 /*************************************************************
  * Lists
diff --git a/tests/lang/eval-okay-zipAttrsWith.exp b/tests/lang/eval-okay-zipAttrsWith.exp
new file mode 100644
index 000000000..9c0b15d22
--- /dev/null
+++ b/tests/lang/eval-okay-zipAttrsWith.exp
@@ -0,0 +1 @@
+{ "0" = { n = "0"; v = [ 5 23 29 ]; }; "1" = { n = "1"; v = [ 7 30 ]; }; "2" = { n = "2"; v = [ 18 ]; }; "4" = { n = "4"; v = [ 10 ]; }; "5" = { n = "5"; v = [ 15 25 26 31 ]; }; "6" = { n = "6"; v = [ 3 14 ]; }; "7" = { n = "7"; v = [ 12 ]; }; "8" = { n = "8"; v = [ 2 6 8 9 ]; }; "9" = { n = "9"; v = [ 0 16 ]; }; a = { n = "a"; v = [ 17 21 22 27 ]; }; c = { n = "c"; v = [ 11 24 ]; }; d = { n = "d"; v = [ 4 13 28 ]; }; e = { n = "e"; v = [ 20 ]; }; f = { n = "f"; v = [ 1 19 ]; }; }
diff --git a/tests/lang/eval-okay-zipAttrsWith.nix b/tests/lang/eval-okay-zipAttrsWith.nix
new file mode 100644
index 000000000..877d4e5fa
--- /dev/null
+++ b/tests/lang/eval-okay-zipAttrsWith.nix
@@ -0,0 +1,9 @@
+with import ./lib.nix;
+
+let
+  str = builtins.hashString "sha256" "test";
+in
+builtins.zipAttrsWith
+  (n: v: { inherit n v; })
+  (map (n: { ${builtins.substring n 1 str} = n; })
+    (range 0 31))