RemoveImport should look through supertypes for Kotlin

rewrite-java/src/main/java/org/openrewrite/java/RemoveImport.java

@@ -18,7 +18,6 @@
 import com.fasterxml.jackson.annotation.JsonCreator;
 import lombok.EqualsAndHashCode;
 import org.jspecify.annotations.Nullable;
-import org.openrewrite.SourceFile;
 import org.openrewrite.internal.ListUtils;
 import org.openrewrite.java.internal.FormatFirstClassPrefix;
 import org.openrewrite.java.style.ImportLayoutStyle;
@@ -28,6 +27,7 @@
 import java.util.*;
 import java.util.concurrent.atomic.AtomicReference;
 
+import static java.util.Collections.singleton;
 import static org.openrewrite.Tree.randomId;
 import static org.openrewrite.java.style.ImportLayoutStyle.isPackageAlwaysFolded;
 
@@ -58,48 +58,85 @@ public RemoveImport(String type, boolean force) {
         J j = tree;
         if (tree instanceof JavaSourceFile) {
             JavaSourceFile cu = (JavaSourceFile) tree;
-            ImportLayoutStyle importLayoutStyle = Optional.ofNullable(((SourceFile) cu).getStyle(ImportLayoutStyle.class))
+            boolean isKotlin = !(cu instanceof J.CompilationUnit) && (cu.getSourcePath().toString().endsWith("kt") || cu.getSourcePath().toString().endsWith("kts")); // poor man's `cu instanceof K.CompilationUnit`
+            ImportLayoutStyle importLayoutStyle = Optional.ofNullable(cu.getStyle(ImportLayoutStyle.class))
                     .orElse(IntelliJ.importLayout());
 
             boolean typeUsed = false;
+            Set<String> types = new HashSet<>(singleton(type));
             Set<String> otherTypesInPackageUsed = new TreeSet<>();
 
             Set<String> methodsAndFieldsUsed = new HashSet<>();
             Set<String> otherMethodsAndFieldsInTypeUsed = new TreeSet<>();
             Set<String> originalImports = new HashSet<>();
             for (J.Import cuImport : cu.getImports()) {
-                if (cuImport.getQualid().getType() != null) {
-                    originalImports.add(((JavaType.FullyQualified) cuImport.getQualid().getType()).getFullyQualifiedName().replace("$", "."));
+                JavaType.FullyQualified fq = TypeUtils.asFullyQualified(cuImport.getQualid().getType());
+                if (fq != null) {
+                    String fqnType = TypeUtils.toFullyQualifiedName(fq.getFullyQualifiedName());
+                    originalImports.add(fqnType);
+                    if (isKotlin && type.equals(fqnType)) {
+                        // For Kotlin, the owning class interfaces with methods can be used without actually importing those interfaces directly...
+                        JavaType.Class owningClass = TypeUtils.asClass(fq.getOwningClass());
+                        if (owningClass != null) {
+                            Queue<JavaType.FullyQualified> toVisit = new LinkedList<>(owningClass.getInterfaces());
+                            Set<JavaType.FullyQualified> visited = new HashSet<>();
+                            while (!toVisit.isEmpty()) {
+                                JavaType.FullyQualified current = toVisit.poll();
+                                if (!visited.add(current)) {
+                                    continue;
+                                }
+                                toVisit.addAll(current.getInterfaces());
+                            }
+                            for (JavaType.FullyQualified current : visited) {
+                                types.add(TypeUtils.toFullyQualifiedName(current.getFullyQualifiedName()));
+                            }
+                        }
+                        // ... and there is the option to star imports references of Java sourced superclasses types
+                        while (fq.getSupertype() != null) {
+                            fq = fq.getSupertype();
+                            types.add(TypeUtils.toFullyQualifiedName(fq.getFullyQualifiedName()));
+                        }
+                    }
                 }
             }
 
             for (JavaType.Variable variable : cu.getTypesInUse().getVariables()) {
                 JavaType.FullyQualified fq = TypeUtils.asFullyQualified(variable.getOwner());
-                if (fq != null && (TypeUtils.fullyQualifiedNamesAreEqual(fq.getFullyQualifiedName(), type) ||
+                if (fq != null && (fullyQualifiedNamesAreEqual(fq.getFullyQualifiedName(), types) ||
                         TypeUtils.fullyQualifiedNamesAreEqual(fq.getFullyQualifiedName(), owner))) {
                     methodsAndFieldsUsed.add(variable.getName());
                 }
             }
 
             for (JavaType.Method method : cu.getTypesInUse().getUsedMethods()) {
-                if (method.hasFlags(Flag.Static)) {
-                    String declaringType = method.getDeclaringType().getFullyQualifiedName();
-                    if (TypeUtils.fullyQualifiedNamesAreEqual(declaringType, type)) {
+                if (method.hasFlags(Flag.Static) || isKotlin) {
+                    String declaringType = TypeUtils.toFullyQualifiedName(method.getDeclaringType().getFullyQualifiedName());
+                    if (fullyQualifiedNamesAreEqual(declaringType, types)) {
                         methodsAndFieldsUsed.add(method.getName());
                     } else if (declaringType.equals(owner)) {
                         if (method.getName().equals(type.substring(type.lastIndexOf('.') + 1))) {
                             methodsAndFieldsUsed.add(method.getName());
                         } else {
                             otherMethodsAndFieldsInTypeUsed.add(method.getName());
                         }
+                    } else if (declaringType.endsWith("Kt") || declaringType.endsWith("Kts")) { // Kotlin top level function
+                        for (JavaType.Method m : method.getDeclaringType().getMethods()) {
+                            if (m.getDeclaringType().getOwningClass() != null) {
+                                String declaringDeclaringType = m.getDeclaringType().getOwningClass().getFullyQualifiedName() + "." + m.getName();
+                                if (fullyQualifiedNamesAreEqual(declaringDeclaringType, types)) {
+                                    methodsAndFieldsUsed.add(method.getName());
+                                    break;
+                                }
+                            }
+                        }
                     }
                 }
             }
 
             for (JavaType javaType : cu.getTypesInUse().getTypesInUse()) {
                 if (javaType instanceof JavaType.FullyQualified) {
                     JavaType.FullyQualified fullyQualified = (JavaType.FullyQualified) javaType;
-                    if (TypeUtils.fullyQualifiedNamesAreEqual(fullyQualified.getFullyQualifiedName(), type)) {
+                    if (fullyQualifiedNamesAreEqual(fullyQualified.getFullyQualifiedName(), types)) {
                         typeUsed = true;
                     } else if (TypeUtils.fullyQualifiedNamesAreEqual(fullyQualified.getFullyQualifiedName(), owner) ||
                             TypeUtils.fullyQualifiedNamesAreEqual(fullyQualified.getPackageName(), owner)) {
@@ -113,7 +150,7 @@ public RemoveImport(String type, boolean force) {
             JavaSourceFile c = cu;
 
             boolean keepImport = !force && (typeUsed || !otherTypesInPackageUsed.isEmpty() && type.endsWith(".*"));
-            AtomicReference<Space> spaceForNextImport = new AtomicReference<>();
+            AtomicReference<@Nullable Space> spaceForNextImport = new AtomicReference<>();
             c = c.withImports(ListUtils.flatMap(c.getImports(), import_ -> {
                 if (spaceForNextImport.get() != null) {
                     Space removedPrefix = spaceForNextImport.get();
@@ -126,14 +163,14 @@ public RemoveImport(String type, boolean force) {
                 }
 
                 String typeName = import_.getTypeName();
-                if (import_.isStatic()) {
-                    String imported = import_.getQualid().getSimpleName();
-                    if (TypeUtils.fullyQualifiedNamesAreEqual(typeName + "." + imported, type) && (force || !methodsAndFieldsUsed.contains(imported))) {
+                String imported = import_.getQualid().getSimpleName();
+                if (import_.isStatic() || (isKotlin && !"*".equals(imported))) {
+                    if (fullyQualifiedNamesAreEqual(typeName + "." + imported, types) && (force || !methodsAndFieldsUsed.contains(imported))) {
                         // e.g. remove java.util.Collections.emptySet when type is java.util.Collections.emptySet
                         spaceForNextImport.set(import_.getPrefix());
                         return null;
-                    } else if ("*".equals(imported) && (TypeUtils.fullyQualifiedNamesAreEqual(typeName, type) ||
-                            TypeUtils.fullyQualifiedNamesAreEqual(typeName + type.substring(type.lastIndexOf('.')), type))) {
+                    } else if ("*".equals(imported) && (fullyQualifiedNamesAreEqual(typeName, types) ||
+                            fullyQualifiedNamesAreEqual(typeName + type.substring(type.lastIndexOf('.')), types))) {
                         if (methodsAndFieldsUsed.isEmpty() && otherMethodsAndFieldsInTypeUsed.isEmpty()) {
                             spaceForNextImport.set(import_.getPrefix());
                             return null;
@@ -142,12 +179,12 @@ public RemoveImport(String type, boolean force) {
                             methodsAndFieldsUsed.addAll(otherMethodsAndFieldsInTypeUsed);
                             return unfoldStarImport(import_, methodsAndFieldsUsed);
                         }
-                    } else if (TypeUtils.fullyQualifiedNamesAreEqual(typeName, type) && !methodsAndFieldsUsed.contains(imported)) {
+                    } else if (fullyQualifiedNamesAreEqual(typeName, types) && !methodsAndFieldsUsed.contains(imported)) {
                         // e.g. remove java.util.Collections.emptySet when type is java.util.Collections
                         spaceForNextImport.set(import_.getPrefix());
                         return null;
                     }
-                } else if (!keepImport && TypeUtils.fullyQualifiedNamesAreEqual(typeName, type)) {
+                } else if (!keepImport && fullyQualifiedNamesAreEqual(typeName, types)) {
                     spaceForNextImport.set(import_.getPrefix());
                     return null;
                 } else if (!keepImport && import_.getPackageName().equals(owner) &&
@@ -175,6 +212,15 @@ public RemoveImport(String type, boolean force) {
         return j;
     }
 
+    private boolean fullyQualifiedNamesAreEqual(String declaringType, Collection<String> types) {
+        for (String type : types) {
+            if (TypeUtils.fullyQualifiedNamesAreEqual(declaringType, type)) {
+                return true;
+            }
+        }
+        return false;
+    }
+
     private long countTrailingLinebreaks(Space space) {
         return space.getLastWhitespace().chars().filter(s -> s == '\n').count();
     }

rewrite-kotlin/src/test/java/org/openrewrite/kotlin/RemoveImportTest.java

@@ -15,6 +15,7 @@
  */
 package org.openrewrite.kotlin;
 
+import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.Test;
 import org.openrewrite.DocumentExample;
 import org.openrewrite.ExecutionContext;
@@ -23,6 +24,7 @@
 import org.openrewrite.test.RewriteTest;
 import org.openrewrite.test.TypeValidation;
 
+import static org.openrewrite.java.Assertions.java;
 import static org.openrewrite.kotlin.Assertions.kotlin;
 import static org.openrewrite.test.RewriteTest.toRecipe;
 
@@ -75,7 +77,7 @@ class A
     @Test
     void removeStarFoldPackage() {
         rewriteRun(
-          spec -> spec.recipe(removeTypeImportRecipe("java.io.OutputStream")).expectedCyclesThatMakeChanges(2),
+          spec -> spec.recipe(removeTypeImportRecipe("java.io.OutputStream")),
           kotlin(
             """
               import java.io.*
@@ -223,4 +225,106 @@ class A
           )
         );
     }
+
+    @Test
+    void keepWhenParentMembersAreUsed() {
+        rewriteRun(
+          spec -> spec.recipe(removeTypeImportRecipe("org.example.Child.Companion.one")),
+          kotlin(
+            """
+              package org.example
+              interface Shared {
+                  fun one() = "one"
+              }
+              open class Parent {
+                  companion object : Shared
+              }
+              class Child : Parent() {
+                  companion object : Shared {
+                      fun two() = "two"
+                  }
+              }
+              """
+          ),
+          kotlin(
+            """
+              import org.example.Child.Companion.one
+              import org.example.Child.Companion.two
+
+              class A {
+                  fun test() {
+                      one()
+                      two()
+                  }
+              }
+              """
+          )
+        );
+    }
+
+    @Test
+    void keepWhenUsingTopLevelFunctions() {
+        rewriteRun(
+          spec -> spec.recipe(removeTypeImportRecipe("org.example.one")),
+          kotlin(
+            """
+              package org.example
+
+              fun one() = "one"
+              fun two() = "two"
+              """
+          ),
+          kotlin(
+            """
+              package org.example2
+
+              import org.example.one
+              import org.example.two
+
+              class Aassss {
+                  fun test() {
+                      one()
+                      two()
+                  }
+              }
+              """
+          )
+        );
+    }
+
+    @Disabled("We cannot use Java sources as dependencies in Kotlin sources yet")
+    @Test
+    void keepStarFoldWhenUsingStaticChildAndParentMembersFromJavaClasses() {
+        rewriteRun(
+          // This kind of setup is only possible with a Java <> Kotlin mix, as you cannot use star imports for companion object members
+          java(
+            """
+              package org.example;
+              public class Parent {
+                  public static void a() {}
+                  public static void b() {}
+              }
+              public class Child extends Parent {
+                   public static void x() {}
+                   public static void y() {}
+              }
+              """
+          ),
+          kotlin(
+            """
+              import org.example.Child.*
+              import org.example.Child.a
+
+              class A {
+                  fun test() {
+                      a()
+                      b()
+                      x()
+                      y()
+                  }
+              }
+              """
+          )
+        );
+    }
 }