diff --git a/dubbo-common/src/main/java/org/apache/dubbo/common/utils/CollectionUtils.java b/dubbo-common/src/main/java/org/apache/dubbo/common/utils/CollectionUtils.java index 9544dbeeec4..3817a1f9de9 100644 --- a/dubbo-common/src/main/java/org/apache/dubbo/common/utils/CollectionUtils.java +++ b/dubbo-common/src/main/java/org/apache/dubbo/common/utils/CollectionUtils.java @@ -419,7 +419,7 @@ public static T first(List values) { return values.get(0); } - public static Set toTreeSet(Set set) { + public static > Set toTreeSet(Set set) { if (isEmpty(set)) { return set; } @@ -429,6 +429,15 @@ public static Set toTreeSet(Set set) { return set; } + public static Set toTreeSet(Set set, Comparator comparator) { + if (set == null) { + return set; + } + Set treeSet = new TreeSet<>(comparator); + treeSet.addAll(set); + return treeSet; + } + public static Set newHashSet(int expectedSize) { return new HashSet<>(capacity(expectedSize)); } diff --git a/dubbo-common/src/test/java/org/apache/dubbo/common/utils/CollectionUtilsTest.java b/dubbo-common/src/test/java/org/apache/dubbo/common/utils/CollectionUtilsTest.java index be48654d2d6..20d801a3676 100644 --- a/dubbo-common/src/test/java/org/apache/dubbo/common/utils/CollectionUtilsTest.java +++ b/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 items = new LinkedHashSet<>(Arrays.asList(new PriorityItem(3), new PriorityItem(1))); + + Set sorted = CollectionUtils.toTreeSet(items, Comparator.comparingInt(PriorityItem::getPriority)); + + List sortedList = new ArrayList<>(sorted); + assertEquals(1, sortedList.get(0).getPriority()); + assertEquals(3, sortedList.get(1).getPriority()); + } + + @Test + void toTreeSetShouldPreserveComparatorForEmptyInput() { + Comparator comparator = Comparator.comparingInt(PriorityItem::getPriority); + + Set sorted = CollectionUtils.toTreeSet(new LinkedHashSet<>(), comparator); + + assertTrue(sorted instanceof TreeSet); + Assertions.assertSame(comparator, ((TreeSet) sorted).comparator()); + } + + @Test + void toTreeSetShouldReturnNullForNullInputWithComparator() { + Comparator comparator = Comparator.comparingInt(PriorityItem::getPriority); + + assertNull(CollectionUtils.toTreeSet(null, comparator)); + } + private static class Person implements Comparable { 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; + } + } }