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