fix TreeSet.ofAll with comparator and empty values (#2255)

* test for TreeSet.removeAll case that removes all elements from the set * fix TreeSet.ofAll() discarding the comparator when creating empty set * removed unnecessary @SuppressWarnings * test for TreeSet.ofAll with explicit comparator and empty initial values
vavr-io · Jan 6, 2019 · 6179333 · 6179333
1 parent 00784b9
commit 6179333
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 3 deletions.
diff --git a/vavr/src/main/java/io/vavr/collection/TreeSet.java b/vavr/src/main/java/io/vavr/collection/TreeSet.java
@@ -192,7 +192,6 @@ public static <T extends Comparable<? super T>> TreeSet<T> fill(int n, Supplier<
         return fill(Comparators.naturalComparator(), n, s);
     }
 
-    @SuppressWarnings("unchecked")
     public static <T extends Comparable<? super T>> TreeSet<T> ofAll(Iterable<? extends T> values) {
         return ofAll(Comparators.naturalComparator(), values);
     }
@@ -201,10 +200,10 @@ public static <T extends Comparable<? super T>> TreeSet<T> ofAll(Iterable<? exte
     public static <T> TreeSet<T> ofAll(Comparator<? super T> comparator, Iterable<? extends T> values) {
         Objects.requireNonNull(comparator, "comparator is null");
         Objects.requireNonNull(values, "values is null");
-        if (values instanceof TreeSet && ((TreeSet) values).comparator() == comparator) {
+        if (values instanceof TreeSet && ((TreeSet<?>) values).comparator() == comparator) {
             return (TreeSet<T>) values;
         } else {
-            return values.iterator().hasNext() ? new TreeSet<>(RedBlackTree.ofAll(comparator, values)) : (TreeSet<T>) empty();
+            return values.iterator().hasNext() ? new TreeSet<>(RedBlackTree.ofAll(comparator, values)) : empty(comparator);
         }
     }
 

diff --git a/vavr/src/test/java/io/vavr/collection/TreeSetTest.java b/vavr/src/test/java/io/vavr/collection/TreeSetTest.java
@@ -236,6 +236,14 @@ public void shouldCollectNonEmpty() {
 
     // -- construct
 
+    @Test
+    public void shouldConstructEmptySetWithExplicitComparator() {
+        final TreeSet<Integer> ts = TreeSet.<Integer> of(Comparators.naturalComparator()
+            .reversed())
+            .addAll(Array.ofAll(1, 2, 3));
+        assertThat(ts.toArray()).isEqualTo(Array.of(3, 2, 1));
+    }
+
     @Test
     public void shouldConstructStreamFromEmptyJavaStream() {
         final TreeSet<Integer> actual = ofJavaStream(java.util.stream.Stream.<Integer>empty());
@@ -282,6 +290,17 @@ public void shouldKeepComparator() {
         final List<Integer> expected = List.of(3, 2, 1);
         assertThat(actual).isEqualTo(expected);
     }
+
+    // -- removeAll
+
+    @Test
+    public void shouldKeepComparatorOnRemoveAll() {
+        final TreeSet<Integer> ts = TreeSet.of(Comparators.naturalComparator()
+            .reversed(), 1, 2, 3)
+            .removeAll(Array.ofAll(1, 2, 3))
+            .addAll(Array.ofAll(4, 5, 6));
+        assertThat(ts.toArray()).isEqualTo(Array.of(6, 5, 4));
+    }
 
     // -- diff