From 658ef711ab80aebf223538dde50ccdb1875106bf Mon Sep 17 00:00:00 2001
From: Ivan Hetman <Ihromant@users.noreply.github.com>
Date: Mon, 16 Oct 2023 15:09:56 +0300
Subject: [PATCH] classlib: additional stream collectors (#822)

---
 .../java/util/stream/TCollectors.java         | 49 +++++++++++++++++++
 .../java/util/stream/CollectorsTest.java      | 11 +++++
 2 files changed, 60 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 0fa30ef34..33957ce7d 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
@@ -63,6 +63,10 @@ public final class TCollectors {
         return toCollection(HashSet::new);
     }
 
+    public static <T> TCollector<T, ?, Set<T>> toUnmodifiableSet() {
+        return collectingAndThen(toSet(), Collections::unmodifiableSet);
+    }
+
     public static TCollector<CharSequence, ?, String> joining() {
         return TCollector.of(StringBuilder::new, StringBuilder::append, StringBuilder::append,
                 StringBuilder::toString);
@@ -90,6 +94,41 @@ public final class TCollectors {
                 sb -> sb.insert(0, prefix).append(suffix).toString());
     }
 
+    public static <T, A, R, K> TCollector<T, ?, K> mapping(Function<? super T, ? extends A> mapper,
+            TCollector<? super A, R, K> downstream) {
+        BiConsumer<R, ? super A> downstreamAccumulator = downstream.accumulator();
+        return TCollector.of(downstream.supplier(), (R r, T t) -> downstreamAccumulator.accept(r, mapper.apply(t)),
+                downstream.combiner(), downstream.finisher(),
+                downstream.characteristics().toArray(TCollector.Characteristics[]::new));
+    }
+
+    public static <T, A, R, K> TCollector<T, ?, K> flatMapping(
+            Function<? super T, ? extends TStream<? extends A>> mapper, TCollector<? super A, R, K> downstream) {
+        BiConsumer<R, ? super A> downstreamAccumulator = downstream.accumulator();
+        return TCollector.of(downstream.supplier(),
+                (R r, T t) -> {
+                    TStream<? extends A> result = mapper.apply(t);
+                    if (result != null) {
+                        result.forEach(a -> downstreamAccumulator.accept(r, a));
+                    }
+                },
+                downstream.combiner(), downstream.finisher(),
+                downstream.characteristics().toArray(TCollector.Characteristics[]::new));
+    }
+
+    public static <T, A, R> TCollector<T, ?, R> filtering(Predicate<? super T> predicate,
+            TCollector<? super T, A, R> downstream) {
+        BiConsumer<A, ? super T> downstreamAccumulator = downstream.accumulator();
+        return TCollector.of(downstream.supplier(),
+                (r, t) -> {
+                    if (predicate.test(t)) {
+                        downstreamAccumulator.accept(r, t);
+                    }
+                },
+                downstream.combiner(), downstream.finisher(),
+                downstream.characteristics().toArray(TCollector.Characteristics[]::new));
+    }
+
     public static <E, K, V> TCollector<E, ?, Map<K, V>> toMap(Function<? super E, ? extends K> keyMapper,
             Function<? super E, ? extends V> valueMapper) {
         return TCollector.of(HashMap::new,
@@ -116,11 +155,21 @@ public final class TCollectors {
                 TCollector.Characteristics.IDENTITY_FINISH);
     }
 
+    public static <E, K, V> TCollector<E, ?, Map<K, V>> toUnmodifiableMap(Function<? super E, ? extends K> keyMapper,
+            Function<? super E, ? extends V> valueMapper) {
+        return collectingAndThen(toMap(keyMapper, valueMapper), Collections::unmodifiableMap);
+    }
+
     public static <E, K, V> TCollector<E, ?, Map<K, V>> toMap(Function<? super E, ? extends K> keyMapper,
             Function<? super E, ? extends V> valueMapper, BinaryOperator<V> mergeFunction) {
         return toMap(keyMapper, valueMapper, mergeFunction, HashMap::new);
     }
 
+    public static <E, K, V> TCollector<E, ?, Map<K, V>> toUnmodifiableMap(Function<? super E, ? extends K> keyMapper,
+            Function<? super E, ? extends V> valueMapper, BinaryOperator<V> mergeFunction) {
+        return collectingAndThen(toMap(keyMapper, valueMapper, mergeFunction), Collections::unmodifiableMap);
+    }
+
     public static <E, K, V, M extends Map<K, V>> TCollector<E, ?, M> toMap(Function<? super E, ? extends K> keyMapper,
             Function<? super E, ? extends V> valueMapper, BinaryOperator<V> mergeFunction, Supplier<M> mapFactory) {
         return TCollector.of(mapFactory,
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 36666f8c3..14ec4c552 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
@@ -185,4 +185,15 @@ public class CollectorsTest {
         assertEquals(Set.of(1, 3, 5, 7, 9), grouped.get(false));
         assertEquals(Set.of(0, 2, 4, 6, 8), grouped.get(true));
     }
+
+    @Test
+    public void simpleCollectors() {
+        List<String> l = List.of("a", "b", "c", "d");
+        assertEquals("aAbBcCdD", l.stream().collect(
+                Collectors.mapping(s -> s + s.toUpperCase(), Collectors.joining())));
+        assertEquals("acd", l.stream().collect(
+                Collectors.filtering(s -> s.indexOf('b') < 0, Collectors.joining())));
+        assertEquals("aaabbbcccddd", l.stream().collect(
+                Collectors.flatMapping(s -> Stream.of(s, s, s), Collectors.joining())));
+    }
 }