From 90700736c7744b97ab69a4fadcd56fa242ec617f Mon Sep 17 00:00:00 2001
From: Silvan Mosberger <contact@infinisil.com>
Date: Thu, 2 Dec 2021 17:46:44 +0100
Subject: Introduce builtins.groupBy primop

This function is very useful in nixpkgs, but its implementation in Nix
itself is rather slow due to it requiring a lot of attribute set and
list appends.
---
 src/libexpr/primops.cc | 50 ++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 50 insertions(+)

(limited to 'src/libexpr/primops.cc')

diff --git a/src/libexpr/primops.cc b/src/libexpr/primops.cc
index d1f4d9009..66af373d7 100644
--- a/src/libexpr/primops.cc
+++ b/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);
-- 
cgit v1.2.3