Fixes DistictIterator to not eat null values (#2426)

* Fixes #2425 * Releasing last consumed 'next' element.
vavr-io · Jul 22, 2019 · 54a3d31 · 54a3d31
1 parent 3295b85
commit 54a3d31
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 3 deletions.
diff --git a/vavr/src/main/java/io/vavr/collection/Iterator.java b/vavr/src/main/java/io/vavr/collection/Iterator.java
@@ -2125,7 +2125,8 @@ final class DistinctIterator<T, U> extends AbstractIterator<T> {
         private final Iterator<? extends T> that;
         private io.vavr.collection.Set<U> known;
         private final Function<? super T, ? extends U> keyExtractor;
-        private T next = null;
+        private boolean nextDefined = false;
+        private T next;
 
         DistinctIterator(Iterator<? extends T> that, Set<U> set, Function<? super T, ? extends U> keyExtractor) {
             this.that = that;
@@ -2135,20 +2136,27 @@ final class DistinctIterator<T, U> extends AbstractIterator<T> {
 
         @Override
         public boolean hasNext() {
-            while (next == null && that.hasNext()) {
+            return nextDefined || searchNext();
+        }
+
+        private boolean searchNext() {
+            while (that.hasNext()) {
                 final T elem = that.next();
                 final U key = keyExtractor.apply(elem);
                 if (!known.contains(key)) {
                     known = known.add(key);
+                    nextDefined = true;
                     next = elem;
+                    return true;
                 }
             }
-            return next != null;
+            return false;
         }
 
         @Override
         public T getNext() {
             final T result = next;
+            nextDefined = false;
             next = null;
             return result;
         }

diff --git a/vavr/src/test/java/io/vavr/collection/IteratorTest.java b/vavr/src/test/java/io/vavr/collection/IteratorTest.java
@@ -467,6 +467,28 @@ public void shouldGenerateFiniteStreamBasedOnOptionSupplier() {
         assertThat(Iterator.iterate(new OptionSupplier(1, 4)).take(50000).reduce((i, j) -> i + j)).isEqualTo(6);
     }
 
+    // -- distinct
+
+    @Test
+    public void shouldStayEmptyOnDistinct() {
+        assertThat(empty().distinct().toList()).isSameAs(List.empty());
+    }
+
+    @Test(/* #2425 */)
+    public void shouldNotEatNullOnDistinct() {
+        assertThat(of((String) null).distinct().toList()).isEqualTo(List.of((String) null));
+    }
+
+    @Test
+    public void shouldKeepDistinctElementsOnDistinct() {
+        assertThat(of(1, 2, 3).distinct().toList()).isEqualTo(List.of(1, 2, 3));
+    }
+
+    @Test
+    public void shouldRemoveDuplicatesOnDistinct() {
+        assertThat(of(1, 2, 1, 3, 3).distinct().toList()).isEqualTo(List.of(1, 2, 3));
+    }
+
     // -- groupBy
 
     @Override