classlib: add mapMulti to Stream (#720)

classlib/src/main/java/org/teavm/classlib/java/util/stream/TDoubleStream.java

@@ -41,7 +41,8 @@ import org.teavm.classlib.java.util.stream.doubleimpl.TSingleDoubleStreamImpl;
 import org.teavm.classlib.java.util.stream.doubleimpl.TSpecializedConcatDoubleStream;
 
 public interface TDoubleStream extends TBaseStream<Double, TDoubleStream> {
-    interface Builder {
+    interface Builder extends DoubleConsumer {
+        @Override
         void accept(double t);
 
         default Builder add(double t) {
@@ -64,6 +65,14 @@ public interface TDoubleStream extends TBaseStream<Double, TDoubleStream> {
 
     TDoubleStream flatMap(DoubleFunction<? extends TDoubleStream> mapper);
 
+    default TDoubleStream mapMulti(DoubleMapMultiConsumer mapper) {
+        return flatMap(e -> {
+            Builder builder = builder();
+            mapper.accept(e, builder);
+            return builder.build();
+        });
+    }
+
     TDoubleStream distinct();
 
     TDoubleStream sorted();
@@ -147,4 +156,9 @@ public interface TDoubleStream extends TBaseStream<Double, TDoubleStream> {
             return new TGenericConcatDoubleStream(a, b);
         }
     }
+
+    @FunctionalInterface
+    interface DoubleMapMultiConsumer {
+        void accept(double value, DoubleConsumer dc);
+    }
 }

classlib/src/main/java/org/teavm/classlib/java/util/stream/TIntStream.java

@@ -43,7 +43,8 @@ import org.teavm.classlib.java.util.stream.intimpl.TSingleIntStreamImpl;
 import org.teavm.classlib.java.util.stream.intimpl.TSpecializedConcatIntStream;
 
 public interface TIntStream extends TBaseStream<Integer, TIntStream> {
-    interface Builder {
+    interface Builder extends IntConsumer {
+        @Override
         void accept(int t);
 
         default Builder add(int t) {
@@ -66,6 +67,14 @@ public interface TIntStream extends TBaseStream<Integer, TIntStream> {
 
     TIntStream flatMap(IntFunction<? extends TIntStream> mapper);
 
+    default TIntStream mapMulti(IntMapMultiConsumer mapper) {
+        return flatMap(e -> {
+            Builder builder = builder();
+            mapper.accept(e, builder);
+            return builder.build();
+        });
+    }
+
     TIntStream distinct();
 
     TIntStream sorted();
@@ -161,4 +170,9 @@ public interface TIntStream extends TBaseStream<Integer, TIntStream> {
             return new TGenericConcatIntStream(a, b);
         }
     }
+
+    @FunctionalInterface
+    interface IntMapMultiConsumer {
+        void accept(int value, IntConsumer ic);
+    }
 }

classlib/src/main/java/org/teavm/classlib/java/util/stream/TLongStream.java

@@ -43,7 +43,8 @@ import org.teavm.classlib.java.util.stream.longimpl.TSingleLongStreamImpl;
 import org.teavm.classlib.java.util.stream.longimpl.TSpecializedConcatLongStream;
 
 public interface TLongStream extends TBaseStream<Long, TLongStream> {
-    interface Builder {
+    interface Builder extends LongConsumer {
+        @Override
         void accept(long t);
 
         default Builder add(long t) {
@@ -66,6 +67,14 @@ public interface TLongStream extends TBaseStream<Long, TLongStream> {
 
     TLongStream flatMap(LongFunction<? extends TLongStream> mapper);
 
+    default TLongStream mapMulti(LongMapMultiConsumer mapper) {
+        return flatMap(e -> {
+            Builder builder = builder();
+            mapper.accept(e, builder);
+            return builder.build();
+        });
+    }
+
     TLongStream distinct();
 
     TLongStream sorted();
@@ -159,4 +168,9 @@ public interface TLongStream extends TBaseStream<Long, TLongStream> {
             return new TGenericConcatLongStream(a, b);
         }
     }
+
+    @FunctionalInterface
+    interface LongMapMultiConsumer {
+        void accept(long value, LongConsumer lc);
+    }
 }

classlib/src/main/java/org/teavm/classlib/java/util/stream/TStream.java

@@ -22,8 +22,11 @@ import java.util.function.BiConsumer;
 import java.util.function.BiFunction;
 import java.util.function.BinaryOperator;
 import java.util.function.Consumer;
+import java.util.function.DoubleConsumer;
 import java.util.function.Function;
+import java.util.function.IntConsumer;
 import java.util.function.IntFunction;
+import java.util.function.LongConsumer;
 import java.util.function.Predicate;
 import java.util.function.Supplier;
 import java.util.function.ToDoubleFunction;
@@ -42,7 +45,8 @@ import org.teavm.classlib.java.util.stream.impl.TSpecializedConcatStream;
 import org.teavm.classlib.java.util.stream.impl.TStreamBuilder;
 
 public interface TStream<T> extends TBaseStream<T, TStream<T>> {
-    interface Builder<T> {
+    interface Builder<T> extends Consumer<T> {
+        @Override
         void accept(T t);
 
         default Builder<T> add(T t) {
@@ -71,6 +75,38 @@ public interface TStream<T> extends TBaseStream<T, TStream<T>> {
 
     TDoubleStream flatMapToDouble(Function<? super T, ? extends TDoubleStream> mapper);
 
+    default <R> TStream<R> mapMulti(BiConsumer<? super T, ? super Consumer<R>> mapper) {
+        return flatMap(e -> {
+            TStream.Builder<R> builder = builder();
+            mapper.accept(e, builder);
+            return builder.build();
+        });
+    }
+
+    default TIntStream mapMultiToInt(BiConsumer<? super T, ? super IntConsumer> mapper) {
+        return flatMapToInt(e -> {
+            TIntStream.Builder builder = TIntStream.builder();
+            mapper.accept(e, builder);
+            return builder.build();
+        });
+    }
+
+    default TLongStream mapMultiToLong(BiConsumer<? super T, ? super LongConsumer> mapper) {
+        return flatMapToLong(e -> {
+            TLongStream.Builder builder = TLongStream.builder();
+            mapper.accept(e, builder);
+            return builder.build();
+        });
+    }
+
+    default TDoubleStream mapMultiToDouble(BiConsumer<? super T, ? super DoubleConsumer> mapper) {
+        return flatMapToDouble(e -> {
+            TDoubleStream.Builder builder = TDoubleStream.builder();
+            mapper.accept(e, builder);
+            return builder.build();
+        });
+    }
+
     TStream<T> distinct();
 
     TStream<T> sorted();

tests/src/test/java/org/teavm/classlib/java/util/stream/IntStreamTest.java

@@ -344,4 +344,14 @@ public class IntStreamTest {
         assertEquals(Integer.MIN_VALUE, empty.getMax());
         assertEquals(0L, empty.getSum());
     }
+
+    @Test
+    public void mapMultiWorks() {
+        int[] mapped = IntStream.rangeClosed(0, 3).mapMulti((cnt, cons) -> {
+            for (int i = 0; i < cnt; i++) {
+                cons.accept(cnt);
+            }
+        }).toArray();
+        assertArrayEquals(new int[] {1, 2, 2, 3, 3, 3}, mapped);
+    }
 }

tests/src/test/java/org/teavm/classlib/java/util/stream/StreamTest.java

@@ -440,4 +440,21 @@ public class StreamTest {
             // ok
         }
     }
+
+    @Test
+    public void mapMultiWorks() {
+        String[] mapped = Stream.of(" a ", "", "   bb       ").<String>mapMulti((s, cons) -> {
+            String trim = s.trim();
+            if (!trim.isEmpty()) {
+                cons.accept(trim);
+            }
+        }).toArray(String[]::new);
+        assertArrayEquals(new String[] {"a", "bb"}, mapped);
+        int[] mappedInt = Stream.of("a", "", "bb").mapMultiToInt((s, cons) -> {
+            if (s.length() % 2 == 0) {
+                cons.accept(s.length());
+            }
+        }).toArray();
+        assertArrayEquals(new int[] {0, 2}, mappedInt);
+    }
 }