Fix Stream.sorted and Stream.distinct laziness

aNNiMON · Mar 5, 2016 · 01d1a19 · 01d1a19
1 parent e821bc5
commit 01d1a19
Show file tree

Hide file tree

Showing 2 changed files with 63 additions and 9 deletions.
diff --git a/src/main/java/com/annimon/stream/Stream.java b/src/main/java/com/annimon/stream/Stream.java
@@ -557,11 +557,25 @@ protected void nextIteration() {
      * @return the new stream
      */
     public Stream<T> distinct() {
-        final Set<T> set = new HashSet<T>();
-        while (iterator.hasNext()) {
-            set.add(iterator.next());
-        }
-        return new Stream<T>(set);
+        return new Stream<T>(new LsaExtIterator<T>() {
+
+            private Iterator<T> distinctIterator;
+
+            @Override
+            protected void nextIteration() {
+                if (!isInit) {
+                    final Set<T> set = new HashSet<T>();
+                    while (iterator.hasNext()) {
+                        set.add(iterator.next());
+                    }
+                    distinctIterator = set.iterator();
+                }
+                hasNext = distinctIterator.hasNext();
+                if (hasNext) {
+                    next = distinctIterator.next();
+                }
+            }
+        });
     }
 
     /**
@@ -594,10 +608,24 @@ public int compare(T o1, T o2) {
      * @param comparator  the {@code Comparator} to compare elements
      * @return the new stream
      */
-    public Stream<T> sorted(Comparator<? super T> comparator) {
-        final List<T> list = collectToList();
-        Collections.sort(list, comparator);
-        return new Stream<T>(list);
+    public Stream<T> sorted(final Comparator<? super T> comparator) {
+        return new Stream<T>(new LsaExtIterator<T>() {
+
+            private Iterator<T> sortedIterator;
+
+            @Override
+            protected void nextIteration() {
+                if (!isInit) {
+                    final List<T> list = collectToList();
+                    Collections.sort(list, comparator);
+                    sortedIterator = list.iterator();
+                }
+                hasNext = sortedIterator.hasNext();
+                if (hasNext) {
+                    next = sortedIterator.next();
+                }
+            }
+        });
     }
 
     /**

diff --git a/src/test/java/com/annimon/stream/StreamTest.java b/src/test/java/com/annimon/stream/StreamTest.java
@@ -359,6 +359,19 @@ public void testDistinct() {
         assertThat(data, is(expected));
     }
 
+    @Test
+    public void testDistinctLazy() {
+        List<Integer> expected = Arrays.asList(-1, 1, 2, 3, 5);
+
+        List<Integer> input = new ArrayList<Integer>(10);
+        input.addAll(Arrays.asList(1, 1, 2, 3, 5));
+        Stream<Integer> stream = Stream.of(input).distinct().sorted();
+        input.addAll(Arrays.asList(3, 2, 1, 1, -1));
+
+        List<Integer> data = stream.collect(Collectors.<Integer>toList());
+        assertThat(data, is(expected));
+    }
+
     @Test
     public void testSorted() {
         List<Integer> expected = Arrays.asList(-7, 0, 3, 6, 9, 19);
@@ -368,6 +381,19 @@ public void testSorted() {
         assertThat(data, is(expected));
     }
 
+    @Test
+    public void testSortedLazy() {
+        List<Integer> expected = Arrays.asList(-7, 0, 3, 6, 9, 19);
+
+        List<Integer> input = new ArrayList<Integer>(6);
+        input.addAll(Arrays.asList(6, 3, 9));
+        Stream<Integer> stream = Stream.of(input).sorted();
+        input.addAll(Arrays.asList(0, -7, 19));
+
+        List<Integer> data = stream.collect(Collectors.<Integer>toList());
+        assertThat(data, is(expected));
+    }
+
     @Test
     public void testSortedWithComparator() {
         List<Integer> expected = Arrays.asList(19, 9, -7, 6, 3, 0);