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:
parent
96d08fcd66
commit
00c993f48b
|
@ -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.
|
||||
|
|
|
@ -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
|
||||
|
|
1
tests/lang/eval-okay-zipAttrsWith.exp
Normal file
1
tests/lang/eval-okay-zipAttrsWith.exp
Normal 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 ]; }; }
|
9
tests/lang/eval-okay-zipAttrsWith.nix
Normal file
9
tests/lang/eval-okay-zipAttrsWith.nix
Normal 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))
|
Loading…
Reference in a new issue