Improve type safety of CollectionUtils.toTreeSet

dubbo-common/src/main/java/org/apache/dubbo/common/utils/CollectionUtils.java

@@ -419,7 +419,7 @@ public static <T> T first(List<T> values) {
         return values.get(0);
     }
 
-    public static <T> Set<T> toTreeSet(Set<T> set) {
+    public static <T extends Comparable<? super T>> Set<T> toTreeSet(Set<T> set) {
         if (isEmpty(set)) {
             return set;
         }
@@ -429,6 +429,15 @@ public static <T> Set<T> toTreeSet(Set<T> set) {
         return set;
     }
 
+    public static <T> Set<T> toTreeSet(Set<T> set, Comparator<? super T> comparator) {
+        if (set == null) {
+            return set;
+        }
+        Set<T> treeSet = new TreeSet<>(comparator);
+        treeSet.addAll(set);
+        return treeSet;
+    }
+
     public static <T> Set<T> newHashSet(int expectedSize) {
         return new HashSet<>(capacity(expectedSize));
     }

dubbo-common/src/test/java/org/apache/dubbo/common/utils/CollectionUtilsTest.java

@@ -21,12 +21,14 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.Comparator;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.TreeSet;
 
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
@@ -255,6 +257,34 @@ void sortShouldSupportComparableSuperType() {
         assertEquals("b", sorted.get(1).getName());
     }
 
+    @Test
+    void toTreeSetShouldSupportExplicitComparator() {
+        Set<PriorityItem> items = new LinkedHashSet<>(Arrays.asList(new PriorityItem(3), new PriorityItem(1)));
+
+        Set<PriorityItem> sorted = CollectionUtils.toTreeSet(items, Comparator.comparingInt(PriorityItem::getPriority));
+
+        List<PriorityItem> sortedList = new ArrayList<>(sorted);
+        assertEquals(1, sortedList.get(0).getPriority());
+        assertEquals(3, sortedList.get(1).getPriority());
+    }
+
+    @Test
+    void toTreeSetShouldPreserveComparatorForEmptyInput() {
+        Comparator<PriorityItem> comparator = Comparator.comparingInt(PriorityItem::getPriority);
+
+        Set<PriorityItem> sorted = CollectionUtils.toTreeSet(new LinkedHashSet<>(), comparator);
+
+        assertTrue(sorted instanceof TreeSet);
+        Assertions.assertSame(comparator, ((TreeSet<PriorityItem>) sorted).comparator());
+    }
+
+    @Test
+    void toTreeSetShouldReturnNullForNullInputWithComparator() {
+        Comparator<PriorityItem> comparator = Comparator.comparingInt(PriorityItem::getPriority);
+
+        assertNull(CollectionUtils.toTreeSet(null, comparator));
+    }
+
     private static class Person implements Comparable<Person> {
 
         private final String name;
@@ -279,4 +309,17 @@ private Student(String name) {
             super(name);
         }
     }
+
+    private static class PriorityItem {
+
+        private final int priority;
+
+        private PriorityItem(int priority) {
+            this.priority = priority;
+        }
+
+        private int getPriority() {
+            return priority;
+        }
+    }
 }