Merge pull request #5715 from Infinisil/groupBy-primop

Introduce `builtins.groupBy` primop
This commit is contained in:
Eelco Dolstra 2021-12-03 13:21:14 +01:00 committed by GitHub
commit 2e606e87c4
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 55 additions and 0 deletions

View file

@ -5,3 +5,6 @@
* `nix develop` now has a flag `--unpack` to run `unpackPhase`. * `nix develop` now has a flag `--unpack` to run `unpackPhase`.
* Lists can now be compared lexicographically using the `<` operator. * Lists can now be compared lexicographically using the `<` operator.
* New built-in function: `builtins.groupBy`, with the same functionality as
Nixpkgs' `lib.groupBy`, but faster.

View file

@ -2928,6 +2928,56 @@ static RegisterPrimOp primop_partition({
.fun = prim_partition, .fun = prim_partition,
}); });
static void prim_groupBy(EvalState & state, const Pos & pos, Value * * args, Value & v)
{
state.forceFunction(*args[0], pos);
state.forceList(*args[1], pos);
ValueVectorMap attrs;
for (auto vElem : args[1]->listItems()) {
Value res;
state.callFunction(*args[0], *vElem, res, pos);
string name = state.forceStringNoCtx(res, pos);
Symbol sym = state.symbols.create(name);
auto vector = attrs.try_emplace(sym, ValueVector()).first;
vector->second.push_back(vElem);
}
state.mkAttrs(v, attrs.size());
for (auto & i : attrs) {
Value * list = state.allocAttr(v, i.first);
auto size = i.second.size();
state.mkList(*list, size);
memcpy(list->listElems(), i.second.data(), sizeof(Value *) * size);
}
}
static RegisterPrimOp primop_groupBy({
.name = "__groupBy",
.args = {"f", "list"},
.doc = R"(
Groups elements of *list* together by the string returned from the
function *f* called on each element. It returns an attribute set
where each attribute value contains the elements of *list* that are
mapped to the same corresponding attribute name returned by *f*.
For example,
```nix
builtins.groupBy (builtins.substring 0 1) ["foo" "bar" "baz"]
```
evaluates to
```nix
{ b = [ "bar" "baz" ]; f = [ "foo" ]; }
```
)",
.fun = prim_groupBy,
});
static void prim_concatMap(EvalState & state, const Pos & pos, Value * * args, Value & v) static void prim_concatMap(EvalState & state, const Pos & pos, Value * * args, Value & v)
{ {
state.forceFunction(*args[0], pos); state.forceFunction(*args[0], pos);

View file

@ -425,9 +425,11 @@ void mkPath(Value & v, const char * s);
#if HAVE_BOEHMGC #if HAVE_BOEHMGC
typedef std::vector<Value *, traceable_allocator<Value *> > ValueVector; typedef std::vector<Value *, traceable_allocator<Value *> > ValueVector;
typedef std::map<Symbol, Value *, std::less<Symbol>, traceable_allocator<std::pair<const Symbol, Value *> > > ValueMap; typedef std::map<Symbol, Value *, std::less<Symbol>, traceable_allocator<std::pair<const Symbol, Value *> > > ValueMap;
typedef std::map<Symbol, ValueVector, std::less<Symbol>, traceable_allocator<std::pair<const Symbol, ValueVector> > > ValueVectorMap;
#else #else
typedef std::vector<Value *> ValueVector; typedef std::vector<Value *> ValueVector;
typedef std::map<Symbol, Value *> ValueMap; typedef std::map<Symbol, Value *> ValueMap;
typedef std::map<Symbol, ValueVector> ValueVectorMap;
#endif #endif