Fix issue #23

aNNiMON · Feb 2, 2016 · b97ebb8 · b97ebb8
1 parent 892721f
commit b97ebb8
Show file tree

Hide file tree

Showing 2 changed files with 64 additions and 15 deletions.
diff --git a/src/main/java/com/annimon/stream/Stream.java b/src/main/java/com/annimon/stream/Stream.java
@@ -381,21 +381,33 @@ public Stream<T> filter(final Predicate<? super T> predicate) {
         return new Stream<T>(new LsaIterator<T>() {
 
             private T next;
+            private boolean hasNext, isInit;
 
             @Override
             public boolean hasNext() {
-                while (iterator.hasNext()) {
-                    next = iterator.next();
-                    if (predicate.test(next)) {
-                        return true;
-                    }
+                if (!isInit) {
+                    nextIteration();
+                    isInit = true;
                 }
-                return false;
+                return hasNext;
             }
 
             @Override
             public T next() {
-                return next;
+                final T result = next;
+                nextIteration();
+                return result;
+            }
+
+            private void nextIteration() {
+                while (iterator.hasNext()) {
+                    next = iterator.next();
+                    if (predicate.test(next)) {
+                        hasNext = true;
+                        return;
+                    }
+                }
+                hasNext = false;
             }
         });
     }
@@ -438,12 +450,29 @@ public <R> Stream<R> flatMap(final Function<? super T, ? extends Stream<? extend
 
             private R next;
             private Iterator<? extends R> inner;
+            private boolean hasNext, isInit;
 
             @Override
             public boolean hasNext() {
+                if (!isInit) {
+                    nextIteration();
+                    isInit = true;
+                }
+                return hasNext;
+            }
+
+            @Override
+            public R next() {
+                final R result = next;
+                nextIteration();
+                return result;
+            }
+
+            private void nextIteration() {
                 if ((inner != null) && inner.hasNext()) {
                     next = inner.next();
-                    return true;
+                    hasNext = true;
+                    return;
                 }
                 while (iterator.hasNext()) {
                     if (inner == null || !inner.hasNext()) {
@@ -455,15 +484,11 @@ public boolean hasNext() {
                     }
                     if ((inner != null) && inner.hasNext()) {
                         next = inner.next();
-                        return true;
+                        hasNext = true;
+                        return;
                     }
                 }
-                return false;
-            }
-
-            @Override
-            public R next() {
-                return next;
+                hasNext = false;
             }
         });
     }

diff --git a/src/test/java/com/annimon/stream/StreamTest.java b/src/test/java/com/annimon/stream/StreamTest.java
@@ -154,6 +154,30 @@ public void testConcat() {
         assertEquals("abcdefgh", consumer.toString());
     }
 
+    @Test
+    public void testConcatOfFilter() {
+        final PrintConsumer consumer = new PrintConsumer();
+        Stream<Integer> stream1 = Stream.ofRange(0, 5).filter(Functions.remainder(1));
+        Stream<Integer> stream2 = Stream.ofRange(5, 10).filter(Functions.remainder(1));
+        Stream.concat(stream1, stream2).forEach(consumer);
+        assertEquals("0123456789", consumer.toString());
+    }
+
+    @Test
+    public void testConcatOfFlatMap() {
+        final PrintConsumer consumer = new PrintConsumer();
+        final Function<Integer, Stream<Integer>> flatmapFunc = new Function<Integer, Stream<Integer>>() {
+            @Override
+            public Stream<Integer> apply(Integer value) {
+                return Stream.of(value, value);
+            }
+        };
+        Stream<Integer> stream1 = Stream.ofRange(1, 3).flatMap(flatmapFunc); // 1122
+        Stream<Integer> stream2 = Stream.ofRange(3, 5).flatMap(flatmapFunc); // 3344
+        Stream.concat(stream1, stream2).forEach(consumer);
+        assertEquals("11223344", consumer.toString());
+    }
+
 
     @Test
     public void testFilter() {