Merge pull request #5830 from pennae/zipAttrsWith

add zipAttrsWith primop

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))