diff --git a/src/main/java/com/annimon/stream/Stream.java b/src/main/java/com/annimon/stream/Stream.java index 3b90784c..3c4f96d4 100644 --- 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 distinct() { - final Set set = new HashSet(); - while (iterator.hasNext()) { - set.add(iterator.next()); - } - return new Stream(set); + return new Stream(new LsaExtIterator() { + + private Iterator distinctIterator; + + @Override + protected void nextIteration() { + if (!isInit) { + final Set set = new HashSet(); + 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 sorted(Comparator comparator) { - final List list = collectToList(); - Collections.sort(list, comparator); - return new Stream(list); + public Stream sorted(final Comparator comparator) { + return new Stream(new LsaExtIterator() { + + private Iterator sortedIterator; + + @Override + protected void nextIteration() { + if (!isInit) { + final List 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 index e772499f..9d48f884 100644 --- 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 expected = Arrays.asList(-1, 1, 2, 3, 5); + + List input = new ArrayList(10); + input.addAll(Arrays.asList(1, 1, 2, 3, 5)); + Stream stream = Stream.of(input).distinct().sorted(); + input.addAll(Arrays.asList(3, 2, 1, 1, -1)); + + List data = stream.collect(Collectors.toList()); + assertThat(data, is(expected)); + } + @Test public void testSorted() { List 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 expected = Arrays.asList(-7, 0, 3, 6, 9, 19); + + List input = new ArrayList(6); + input.addAll(Arrays.asList(6, 3, 9)); + Stream stream = Stream.of(input).sorted(); + input.addAll(Arrays.asList(0, -7, 19)); + + List data = stream.collect(Collectors.toList()); + assertThat(data, is(expected)); + } + @Test public void testSortedWithComparator() { List expected = Arrays.asList(19, 9, -7, 6, 3, 0);