Add a specific case for SortedSet to preserve existing behavior

Author: lredor

src/main/java/org/openrewrite/java/migrate/guava/NoGuavaSetsFilter.java

@@ -19,8 +19,8 @@
 import org.openrewrite.Preconditions;
 import org.openrewrite.Recipe;
 import org.openrewrite.TreeVisitor;
-import org.openrewrite.java.JavaIsoVisitor;
 import org.openrewrite.java.JavaTemplate;
+import org.openrewrite.java.JavaVisitor;
 import org.openrewrite.java.MethodMatcher;
 import org.openrewrite.java.search.UsesMethod;
 import org.openrewrite.java.tree.J;
@@ -31,6 +31,7 @@
 
 public class NoGuavaSetsFilter extends Recipe {
     private static final MethodMatcher SETS_FILTER = new MethodMatcher("com.google.common.collect.Sets filter(java.util.Set, com.google.common.base.Predicate)");
+    private static final MethodMatcher SETS_FILTER_SORTED_SET = new MethodMatcher("com.google.common.collect.Sets filter(java.util.SortedSet, com.google.common.base.Predicate)");
 
     @Override
     public String getDisplayName() {
@@ -49,28 +50,44 @@ public Set<String> getTags() {
 
     @Override
     public TreeVisitor<?, ExecutionContext> getVisitor() {
-        return Preconditions.check(
-                new UsesMethod<>(SETS_FILTER),
-                new JavaIsoVisitor<ExecutionContext>() {
-                    @Override
-                    public J.MethodInvocation visitMethodInvocation(J.MethodInvocation method, ExecutionContext ctx) {
-                        if (SETS_FILTER.matches(method)) {
-                            maybeRemoveImport("com.google.common.base.Predicate");
-                            maybeRemoveImport("com.google.common.collect.Sets");
-                            maybeAddImport("java.util.function.Predicate");
-                            maybeAddImport("java.util.stream.Collectors");
+        TreeVisitor<?, ExecutionContext> precondition = Preconditions.or(new UsesMethod<>(SETS_FILTER), new UsesMethod<>(SETS_FILTER_SORTED_SET));
+        return Preconditions.check(precondition, new JavaVisitor<ExecutionContext>() {
+                @Override
+                public J visitMethodInvocation(J.MethodInvocation method, ExecutionContext ctx) {
+                    if (SETS_FILTER_SORTED_SET.matches(method)) {
+                        maybeRemoveImport("com.google.common.base.Predicate");
+                        maybeRemoveImport("com.google.common.collect.Sets");
+                        maybeAddImport("java.util.TreeSet");
+                        maybeAddImport("java.util.function.Predicate");
+                        maybeAddImport("java.util.stream.Collectors");
 
-                            return JavaTemplate.builder("#{any(java.util.Collection)}.stream().filter(#{any(java.util.function.Predicate)}).collect(Collectors.toSet())")
-                                    .imports("java.util.stream.Collectors")
-                                    .build()
-                                    .apply(getCursor(),
-                                            method.getCoordinates().replace(),
-                                            method.getArguments().get(0),
-                                            method.getArguments().get(1));
-                        }
-                        return super.visitMethodInvocation(method, ctx);
+                        return JavaTemplate.builder("#{any(java.util.Collection)}.stream().filter(#{any(java.util.function.Predicate)}).collect(Collectors.toCollection(TreeSet::new))")
+                                .imports("java.util.TreeSet")
+                                .imports("java.util.stream.Collectors")
+                                .build()
+                                .apply(getCursor(),
+                                        method.getCoordinates().replace(),
+                                        method.getArguments().get(0),
+                                        method.getArguments().get(1));
                     }
+                    if (SETS_FILTER.matches(method)) {
+                        maybeRemoveImport("com.google.common.base.Predicate");
+                        maybeRemoveImport("com.google.common.collect.Sets");
+                        maybeAddImport("java.util.TreeSet");
+                        maybeAddImport("java.util.function.Predicate");
+                        maybeAddImport("java.util.stream.Collectors");
+
+                        return JavaTemplate.builder("#{any(java.util.Collection)}.stream().filter(#{any(java.util.function.Predicate)}).collect(Collectors.toSet())")
+                                .imports("java.util.stream.Collectors")
+                                .build()
+                                .apply(getCursor(),
+                                        method.getCoordinates().replace(),
+                                        method.getArguments().get(0),
+                                        method.getArguments().get(1));
+                    }
+                    return super.visitMethodInvocation(method, ctx);
                 }
+            }
         );
     }
 }

src/test/java/org/openrewrite/java/migrate/guava/NoGuavaSetsFilterTest.java

@@ -66,4 +66,40 @@ public static Set<Object> test(Set<Object> set, Predicate<Object> isNotNull) {
           )
         );
     }
+
+    @Test
+    void replaceSetsFilterUsingSortedSet() {
+        //language=java
+        rewriteRun(
+          java(
+            """
+              import java.util.Set;
+              import java.util.SortedSet;
+
+              import com.google.common.base.Predicate;
+              import com.google.common.collect.Sets;
+
+              class Test {
+                  public static Set<Object> test(SortedSet<Object> set, Predicate<Object> isNotNull) {
+                      return Sets.filter(set, isNotNull);
+                  }
+              }
+              """,
+            """
+              import java.util.Set;
+              import java.util.SortedSet;
+              import java.util.TreeSet;
+              import java.util.stream.Collectors;
+
+              import com.google.common.base.Predicate;
+
+              class Test {
+                  public static Set<Object> test(SortedSet<Object> set, Predicate<Object> isNotNull) {
+                      return set.stream().filter(isNotNull).collect(Collectors.toCollection(TreeSet::new));
+                  }
+              }
+              """
+          )
+        );
+    }
 }