Fix Stream.skip laziness

aNNiMON · Mar 3, 2016 · 680d602 · 680d602
1 parent 42be36d
commit 680d602
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 5 deletions.
diff --git a/src/main/java/com/annimon/stream/Stream.java b/src/main/java/com/annimon/stream/Stream.java
@@ -716,11 +716,29 @@ public T next() {
      * @param n  the number of elements to skip
      * @return the new stream
      */
-    public Stream<T> skip(long n) {
-        for (long i = 0; i < n && iterator.hasNext(); i++) {
-            iterator.next();
-        }
-        return this;
+    public Stream<T> skip(final long n) {
+        return new Stream<T>(new LsaIterator<T>() {
+
+            private long skippedCount;
+
+            @Override
+            public boolean hasNext() {
+                if (skippedCount < n) {
+                    while (skippedCount < n) {
+                        if (!iterator.hasNext()) return false;
+                        iterator.next();
+                        skippedCount++;
+                    }
+                    return true;
+                }
+                return iterator.hasNext();
+            }
+
+            @Override
+            public T next() {
+                return iterator.next();
+            }
+        });
     }
 
     /**

diff --git a/src/test/java/com/annimon/stream/StreamTest.java b/src/test/java/com/annimon/stream/StreamTest.java
@@ -461,6 +461,20 @@ public void testSkipMoreThanCount() {
         assertEquals(0, count);
     }
 
+    @Test
+    public void testSkipLazy() {
+        final List<Integer> data = new ArrayList<Integer>(10);
+        data.add(0);
+
+        final PrintConsumer consumer = new PrintConsumer();
+        Stream stream = Stream.of(data).skip(3);
+        data.addAll(Arrays.asList(1, 2, 3, 4, 5));
+        stream = stream.peek(consumer);
+
+        assertEquals(3, stream.count());
+        assertEquals("345", consumer.toString());
+    }
+
     @Test
     public void testSkipAndLimit() {
         final PrintConsumer consumer = new PrintConsumer();