Merge pull request #5715 from Infinisil/groupBy-primop

Introduce `builtins.groupBy` primop

doc/manual/src/release-notes/rl-next.md

@@ -5,3 +5,6 @@
 * `nix develop` now has a flag `--unpack` to run `unpackPhase`.
 
 * 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.

src/libexpr/primops.cc

@@ -2928,6 +2928,56 @@ static RegisterPrimOp primop_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)
 {
     state.forceFunction(*args[0], pos);

src/libexpr/value.hh

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