From 00c993f48b1c752656d8cdd2c25e38f9ba572128 Mon Sep 17 00:00:00 2001 From: pennae Date: Sat, 25 Dec 2021 15:29:49 +0100 Subject: [PATCH] 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. --- doc/manual/src/release-notes/rl-next.md | 2 + src/libexpr/primops.cc | 84 +++++++++++++++++++++++++ tests/lang/eval-okay-zipAttrsWith.exp | 1 + tests/lang/eval-okay-zipAttrsWith.nix | 9 +++ 4 files changed, 96 insertions(+) create mode 100644 tests/lang/eval-okay-zipAttrsWith.exp create mode 100644 tests/lang/eval-okay-zipAttrsWith.nix 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))