[functional] Add groupby and Stream.groupBy

Author: aNNiMON

modules/main/src/main/java/com/annimon/ownlang/modules/functional/StreamValue.java

@@ -31,6 +31,7 @@ private void init() {
         set("reduce", wrapTerminal(new functional_reduce()));
         set("forEach", wrapTerminal(new functional_forEach()));
         set("forEachIndexed", wrapTerminal(new functional_forEachIndexed()));
+        set("groupBy", wrapTerminal(new functional_groupBy()));
         set("toArray", args -> container);
         set("joining", container::joinToString);
         set("count", args -> NumberValue.of(container.size()));

modules/main/src/main/java/com/annimon/ownlang/modules/functional/functional.java

@@ -27,6 +27,7 @@ public Map<String, Function> functions() {
         result.put("sortby", new functional_sortBy());
         result.put("takewhile", new functional_takeWhile());
         result.put("dropwhile", new functional_dropWhile());
+        result.put("groupby", new functional_groupBy());
 
         result.put("chain", new functional_chain());
         result.put("stream", new functional_stream());

modules/main/src/main/java/com/annimon/ownlang/modules/functional/functional_groupBy.java

@@ -0,0 +1,67 @@
+package com.annimon.ownlang.modules.functional;
+
+import com.annimon.ownlang.exceptions.TypeException;
+import com.annimon.ownlang.lib.*;
+import java.util.ArrayList;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+
+public final class functional_groupBy implements Function {
+
+    @Override
+    public Value execute(Value[] args) {
+        Arguments.check(2, args.length);
+        
+        final Value container = args[0];
+        final Function classifier = ValueUtils.consumeFunction(args[1], 1);
+        return groupBy(container, classifier);
+    }
+
+    static Value groupBy(Value container, Function classifier) {
+        if (container.type() == Types.ARRAY) {
+            return groupByArray((ArrayValue) container, classifier);
+        }
+        if (container.type() == Types.MAP) {
+            return groupByMap((MapValue) container, classifier);
+        }
+        throw new TypeException("Invalid first argument. Array or map expected");
+    }
+
+    @SuppressWarnings("Java8MapApi")
+    static Value groupByArray(ArrayValue array, Function classifier) {
+        final var result = new LinkedHashMap<Value, List<Value>>();
+        for (Value element : array) {
+            final var key = classifier.execute(element);
+            var container = result.get(key);
+            if (container == null) {
+                container = new ArrayList<>();
+                result.put(key, container);
+            }
+            container.add(element);
+        }
+        return fromMapOfArrays(result);
+    }
+
+    static Value groupByMap(MapValue map, Function classifier) {
+        final var result = new LinkedHashMap<Value, Value>();
+        for (Map.Entry<Value, Value> element : map) {
+            final var k = element.getKey();
+            final var v = element.getValue();
+            final var key = classifier.execute(k, v);
+            var container = (MapValue) result.get(key);
+            if (container == null) {
+                container = new MapValue(10);
+                result.put(key, container);
+            }
+            container.set(k, v);
+        }
+        return new MapValue(result);
+    }
+
+    private static MapValue fromMapOfArrays(Map<Value, List<Value>> map) {
+        final var result = new LinkedHashMap<Value, Value>();
+        map.forEach((key, value) -> result.put(key, new ArrayValue(value)));
+        return new MapValue(result);
+    }
+}

ownlang-parser/src/test/resources/modules/functional/groupby.own

@@ -0,0 +1,15 @@
+use std, functional
+
+def testArraysGroupBy() {
+  arr = [1, 2, 3, 4, 1, 2, 3, 1, 2, 3]
+  result = groupby(arr, def(v) = v % 2 == 0)
+  assertEquals([2, 4, 2, 2], result[true])
+  assertEquals([1, 3, 1, 3, 1, 3], result[false])
+}
+
+def testMapsGroupBy() {
+  map = {"abc": 123, "test1": 234, "test2": 345, "test3": 456, "def": 567}
+  result = groupby(map, def(k, v) = k.startsWith("test"))
+  assertEquals({"test1": 234, "test2": 345, "test3": 456}, result[true])
+  assertEquals({"abc": 123, "def": 567}, result[false])
+}

ownlang-parser/src/test/resources/modules/functional/stream.own

@@ -53,6 +53,15 @@ def testCustom() {
   assertEquals([5,6,4,2], stream(data).custom(::reverse).toArray())
 }
 
+def reverse(container) {
+  size = length(container)
+  result = newarray(size)
+  for i : range(size) {
+    result[size - i - 1] = container[i]
+  }
+  return result
+}
+
 def testJoining() {
   data = [1,2,3,4]
   assertEquals("1234", stream(data).joining())
@@ -101,11 +110,18 @@ def testForEachMapIndexed() {
   assertEquals("a10b21", result)
 }
 
-def reverse(container) {
-  size = length(container)
-  result = newarray(size)
-  for i : range(size) {
-    result[size - i - 1] = container[i]
-  }
-  return result
-}
+def testArraysGroupBy() {
+  data = [1, 2, 3, 4, 1, 2, 3, 1, 2, 3]
+  result = stream(data)
+    .groupBy(def(v) = v % 2 == 0)
+  assertEquals([2, 4, 2, 2], result[true])
+  assertEquals([1, 3, 1, 3, 1, 3], result[false])
+}
+
+def testMapsGroupBy() {
+  data = {"abc": 123, "test1": 234, "test2": 345, "test3": 456, "def": 567}
+  result = stream(data)
+      .groupBy(def(entry) = entry[0].startsWith("test"))
+  assertEquals([["test1", 234], ["test2", 345], ["test3", 456]], sort(result[true]))
+  assertEquals([["abc", 123], ["def", 567]], sort(result[false]))
+}