From 23a639326724c4a6f3207ac1975114c9edc278d5 Mon Sep 17 00:00:00 2001
From: Ivan Hetman <Ihromant@users.noreply.github.com>
Date: Wed, 21 Jun 2023 16:35:37 +0300
Subject: [PATCH] classlib: collectors partitioning by (#716)

---
 .../java/util/stream/TCollectors.java         | 21 +++++++++++++++++++
 .../java/util/stream/CollectorsTest.java      |  9 ++++++++
 2 files changed, 30 insertions(+)

diff --git a/classlib/src/main/java/org/teavm/classlib/java/util/stream/TCollectors.java b/classlib/src/main/java/org/teavm/classlib/java/util/stream/TCollectors.java
index f97c4230e..0fa30ef34 100644
--- a/classlib/src/main/java/org/teavm/classlib/java/util/stream/TCollectors.java
+++ b/classlib/src/main/java/org/teavm/classlib/java/util/stream/TCollectors.java
@@ -30,6 +30,7 @@ import java.util.function.BiConsumer;
 import java.util.function.BiFunction;
 import java.util.function.BinaryOperator;
 import java.util.function.Function;
+import java.util.function.Predicate;
 import java.util.function.Supplier;
 import java.util.function.ToDoubleFunction;
 import java.util.function.ToIntFunction;
@@ -357,4 +358,24 @@ public final class TCollectors {
             this.b = b;
         }
     }
+
+    public static <T, A, R> TCollector<T, ?, Map<Boolean, R>> partitioningBy(Predicate<? super T> predicate,
+            TCollector<? super T, A, R> downstream) {
+        BiConsumer<A, ? super T> acc = downstream.accumulator();
+        return teeing(TCollector.of(downstream.supplier(), (res, el) -> {
+                    if (!predicate.test(el)) {
+                        acc.accept(res, el);
+                    }
+                }, downstream.combiner(), downstream.finisher()),
+                TCollector.of(downstream.supplier(), (res1, el1) -> {
+                    if (predicate.test(el1)) {
+                        acc.accept(res1, el1);
+                    }
+                }, downstream.combiner(), downstream.finisher()),
+                (fls, tr) -> Map.of(false, fls, true, tr));
+    }
+
+    public static <T> TCollector<T, ?, Map<Boolean, List<T>>> partitioningBy(Predicate<? super T> predicate) {
+        return partitioningBy(predicate, toList());
+    }
 }
diff --git a/tests/src/test/java/org/teavm/classlib/java/util/stream/CollectorsTest.java b/tests/src/test/java/org/teavm/classlib/java/util/stream/CollectorsTest.java
index 6d5b9be2f..4efe438fe 100644
--- a/tests/src/test/java/org/teavm/classlib/java/util/stream/CollectorsTest.java
+++ b/tests/src/test/java/org/teavm/classlib/java/util/stream/CollectorsTest.java
@@ -25,6 +25,7 @@ import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.Set;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
@@ -178,4 +179,12 @@ public class CollectorsTest {
                 .collect(Collectors.teeing(Collectors.summingInt(String::length),
                 Collectors.averagingInt(String::length), (sum, avg) -> sum / avg)));
     }
+
+    @Test
+    public void partitioningBy() {
+        Map<Boolean, Set<Integer>> grouped = IntStream.range(0, 10).boxed()
+                .collect(Collectors.partitioningBy(i -> i % 2 == 0, Collectors.toSet()));
+        assertEquals(Set.of(1, 3, 5, 7, 9), grouped.get(false));
+        assertEquals(Set.of(0, 2, 4, 6, 8), grouped.get(true));
+    }
 }