Merge pull request #30 from aNNiMON/takewhile

Add takeWhile and dropWhile operators
aNNiMON · Mar 5, 2016 · d571183 · d571183
2 parents 428ebe2 + 804bbc6
commit d571183
Show file tree

Hide file tree

Showing 4 changed files with 142 additions and 60 deletions.
diff --git a/build.gradle b/build.gradle
@@ -39,7 +39,7 @@ buildscript {
 
 cobertura {
     coverageFormats = ['html', 'xml']
-    coverageExcludes = ['.*com.annimon.stream.LsaIterator']
+    coverageExcludes = ['.*com.annimon.stream.LsaIterator', '.*com.annimon.stream.LsaExtIterator']
 }
 
 // maven signing

diff --git a/src/main/java/com/annimon/stream/LsaExtIterator.java b/src/main/java/com/annimon/stream/LsaExtIterator.java
@@ -0,0 +1,32 @@
+package com.annimon.stream;
+
+import java.util.Iterator;
+
+public abstract class LsaExtIterator<T> implements Iterator<T> {
+
+    protected T next;
+    protected boolean hasNext, isInit;
+
+    @Override
+    public boolean hasNext() {
+        if (!isInit) {
+            nextIteration();
+            isInit = true;
+        }
+        return hasNext;
+    }
+
+    @Override
+    public T next() {
+        final T result = next;
+        nextIteration();
+        return result;
+    }
+
+    protected abstract void nextIteration();
+
+    @Override
+    public void remove() {
+        throw new UnsupportedOperationException("remove not supported");
+    }
+}
diff --git a/src/main/java/com/annimon/stream/Stream.java b/src/main/java/com/annimon/stream/Stream.java
@@ -331,28 +331,10 @@ public T next() {
     public static <T> Stream<T> concat(Stream<? extends T> stream1, Stream<? extends T> stream2) {
         final Iterator<? extends T> it1 = stream1.iterator;
         final Iterator<? extends T> it2 = stream2.iterator;
-        return new Stream<T>(new LsaIterator<T>() {
-
-            private T next;
-            private boolean hasNext, isInit;
-
-            @Override
-            public boolean hasNext() {
-                if (!isInit) {
-                    nextIteration();
-                    isInit = true;
-                }
-                return hasNext;
-            }
+        return new Stream<T>(new LsaExtIterator<T>() {
 
             @Override
-            public T next() {
-                final T result = next;
-                nextIteration();
-                return result;
-            }
-
-            private void nextIteration() {
+            protected void nextIteration() {
                 if (it1.hasNext()) {
                     next = it1.next();
                     hasNext = true;
@@ -487,28 +469,10 @@ public <R> R custom(Function<Stream<T>, R> function) {
      * @return the new stream
      */
     public Stream<T> filter(final Predicate<? super T> predicate) {
-        return new Stream<T>(new LsaIterator<T>() {
-
-            private T next;
-            private boolean hasNext, isInit;
+        return new Stream<T>(new LsaExtIterator<T>() {
 
             @Override
-            public boolean hasNext() {
-                if (!isInit) {
-                    nextIteration();
-                    isInit = true;
-                }
-                return hasNext;
-            }
-
-            @Override
-            public T next() {
-                final T result = next;
-                nextIteration();
-                return result;
-            }
-
-            private void nextIteration() {
+            protected void nextIteration() {
                 while (iterator.hasNext()) {
                     next = iterator.next();
                     if (predicate.test(next)) {
@@ -555,29 +519,12 @@ public R next() {
      * @return the new stream
      */
     public <R> Stream<R> flatMap(final Function<? super T, ? extends Stream<? extends R>> mapper) {
-        return new Stream<R>(new LsaIterator<R>() {
+        return new Stream<R>(new LsaExtIterator<R>() {
 
-            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() {
+            protected void nextIteration() {
                 if ((inner != null) && inner.hasNext()) {
                     next = inner.next();
                     hasNext = true;
@@ -710,6 +657,55 @@ public T next() {
         });
     }
 
+    /**
+     * Takes elements while the predicate is true.
+     *
+     * <p>This is an intermediate operation.
+     *
+     * @param predicate  the predicate used to take elements
+     * @return the new stream
+     */
+    public Stream<T> takeWhile(final Predicate<? super T> predicate) {
+        return new Stream<T>(new LsaExtIterator<T>() {
+
+            @Override
+            protected void nextIteration() {
+                hasNext = iterator.hasNext() && predicate.test(next = iterator.next());
+            }
+        });
+    }
+
+    /**
+     * Drops elements while the predicate is true and returns the rest.
+     *
+     * <p>This is an intermediate operation.
+     *
+     * @param predicate  the predicate used to drop elements
+     * @return the new stream
+     */
+    public Stream<T> dropWhile(final Predicate<? super T> predicate) {
+        return new Stream<T>(new LsaExtIterator<T>() {
+
+            @Override
+            protected void nextIteration() {
+                if (!isInit) {
+                    // Skip first time
+                    while (hasNext = iterator.hasNext()) {
+                        next = iterator.next();
+                        if (!predicate.test(next)) {
+                            return;
+                        }
+                    }
+                }
+
+                hasNext = hasNext && iterator.hasNext();
+                if (!hasNext) return;
+
+                next = iterator.next();
+            }
+        });
+    }
+
     /**
      * Returns {@code Stream} with first {@code maxSize} elements.
      *

diff --git a/src/test/java/com/annimon/stream/StreamTest.java b/src/test/java/com/annimon/stream/StreamTest.java
@@ -473,6 +473,60 @@ public void testPeek() {
         assertEquals("01234", consumer.toString());
     }
 
+    @Test
+    public void testTakeWhile() {
+        final PrintConsumer<Integer> consumer = new PrintConsumer<Integer>();
+        long count = Stream.of(2, 4, 6, 7, 8, 10, 11)
+                .takeWhile(Functions.remainder(2))
+                .peek(consumer)
+                .count();
+        assertEquals(3, count);
+        assertEquals("246", consumer.toString());
+    }
+
+    @Test
+    public void testTakeWhileNonFirstMatch() {
+        long count = Stream.of(2, 4, 6, 7, 8, 10, 11)
+                .takeWhile(Functions.remainder(3))
+                .count();
+        assertEquals(0, count);
+    }
+
+    @Test
+    public void testTakeWhileAllMatch() {
+        long count = Stream.of(2, 4, 6, 7, 8, 10, 11)
+                .takeWhile(Functions.remainder(1))
+                .count();
+        assertEquals(7, count);
+    }
+
+    @Test
+    public void testDropWhile() {
+        final PrintConsumer<Integer> consumer = new PrintConsumer<Integer>();
+        long count = Stream.of(2, 4, 6, 7, 8, 10, 11)
+                .dropWhile(Functions.remainder(2))
+                .peek(consumer)
+                .count();
+        assertEquals(4, count);
+        assertEquals("781011", consumer.toString());
+    }
+
+    @Test
+    public void testDropNonFirstMatch() {
+        long count = Stream.of(2, 4, 6, 7, 8, 10, 11)
+                .dropWhile(Functions.remainder(3))
+                .count();
+        assertEquals(7, count);
+    }
+
+    @Test
+    public void testDropWhileAllMatch() {
+        long count = Stream.of(2, 4, 6, 7, 8, 10, 11)
+                .dropWhile(Functions.remainder(1))
+                .count();
+        assertEquals(0, count);
+    }
+
     @Test
     public void testLimit() {
         final PrintConsumer<Integer> consumer = new PrintConsumer<Integer>();