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.
This commit is contained in:
pennae 2021-12-25 15:29:49 +01:00
parent 96d08fcd66
commit 00c993f48b
4 changed files with 96 additions and 0 deletions

View file

@ -4,3 +4,5 @@
more compliant one](https://github.com/ToruNiina/toml11). more compliant one](https://github.com/ToruNiina/toml11).
* Added `:st`/`:show-trace` commands to nix repl, which are used to * Added `:st`/`:show-trace` commands to nix repl, which are used to
set or toggle display of error traces. set or toggle display of error traces.
* New builtin function `builtins.zipAttrsWith` with same functionality
as `lib.zipAttrsWith` from nixpkgs, but much more efficient.

View file

@ -2503,6 +2503,90 @@ static RegisterPrimOp primop_mapAttrs({
.fun = prim_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 * Lists

View file

@ -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 ]; }; }

View file

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