[#54] StreamExHeadTailTest: flatMap test added (became possible after

#55)
amaembo · Jan 20, 2016 · 760976f · 760976f
1 parent 5b7735e
commit 760976f
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/src/test/java/one/util/streamex/StreamExHeadTailTest.java b/src/test/java/one/util/streamex/StreamExHeadTailTest.java
@@ -102,6 +102,11 @@ static <T> StreamEx<T> peek(StreamEx<T> input, Consumer<T> consumer) {
             return peek(tail, consumer).prepend(head);
         });
     }
+
+    // Stream.flatMap (TSO)
+    static <T, R> StreamEx<R> flatMap(StreamEx<T> input, Function<T, Stream<R>> mapper) {
+        return input.headTail((head, tail) -> flatMap(tail, mapper).prepend(mapper.apply(head)));
+    }
 
     // Stream.sorted
     static <T> StreamEx<T> sorted(StreamEx<T> input) {
@@ -402,6 +407,8 @@ public void testHeadTailTCO() {
         // 1+2+...+10000
         assertEquals(50005000, (int) limit(scanLeft(StreamEx.iterate(1, x -> x + 1), Integer::sum), 10000).reduce(
             (a, b) -> b).get());
+        assertEquals(50005000, (int) limit(flatMap(StreamEx.iterate(1, x -> x + 3), x -> StreamEx.of(x, x + 1, x + 2)),
+            10000).reduce(Integer::sum).get());
         assertEquals(asList(50005000), skip(
             appendReduction(IntStreamEx.rangeClosed(1, 10000).boxed(), 0, Integer::sum), 10000).toList());
         AtomicInteger sum = new AtomicInteger();