Merge pull request #5830 from pennae/zipAttrsWith

add zipAttrsWith primop
This commit is contained in:
Eelco Dolstra 2022-01-04 11:16:50 +01:00 committed by GitHub
commit 1ffacad8a5
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
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))