Update ChangeType recipe to support kotlin alias import (#3558)

* Update ChangeType recipe to support kotlin alias import
openrewrite · Sep 21, 2023 · f27cd1b · f27cd1b
1 parent f9ad4f8
commit f27cd1b
Showing 1 changed file with 53 additions and 0 deletions.
diff --git a/rewrite-java/src/main/java/org/openrewrite/java/ChangeType.java b/rewrite-java/src/main/java/org/openrewrite/java/ChangeType.java
@@ -15,10 +15,12 @@
  */
 package org.openrewrite.java;
 
+import lombok.AllArgsConstructor;
 import lombok.EqualsAndHashCode;
 import lombok.Value;
 import org.openrewrite.*;
 import org.openrewrite.internal.ListUtils;
+import org.openrewrite.internal.lang.NonNull;
 import org.openrewrite.internal.lang.Nullable;
 import org.openrewrite.java.search.UsesType;
 import org.openrewrite.java.tree.*;
@@ -85,6 +87,8 @@ public J visit(@Nullable Tree tree, ExecutionContext ctx) {
     private static class ChangeTypeVisitor extends JavaVisitor<ExecutionContext> {
         private final JavaType.Class originalType;
         private final JavaType targetType;
+        @Nullable
+        private J.Identifier importAlias;
 
         @Nullable
         private final Boolean ignoreDefinition;
@@ -96,6 +100,7 @@ private ChangeTypeVisitor(String oldFullyQualifiedTypeName, String newFullyQuali
             this.originalType = JavaType.ShallowClass.build(oldFullyQualifiedTypeName);
             this.targetType = JavaType.buildType(newFullyQualifiedTypeName);
             this.ignoreDefinition = ignoreDefinition;
+            importAlias = null;
         }
 
         @Override
@@ -125,6 +130,14 @@ public J visitClassDeclaration(J.ClassDeclaration classDecl, ExecutionContext ct
 
         @Override
         public J visitImport(J.Import import_, ExecutionContext ctx) {
+            // Just Collect kotlin alias for original type here
+            if (import_.getAlias() != null) {
+                if (hasSameFQN(import_, originalType)) {
+                    // collect Kotlin alias
+                    importAlias = import_.getAlias();
+                }
+            }
+
             // visitCompilationUnit() handles changing the imports.
             // If we call super.visitImport() then visitFieldAccess() will change the imports before AddImport/RemoveImport see them.
             // visitFieldAccess() doesn't have the import-specific formatting logic that AddImport/RemoveImport do.
@@ -180,6 +193,7 @@ public J visitImport(J.Import import_, ExecutionContext ctx) {
                         if (!"java.lang".equals(fullyQualifiedTarget.getPackageName())) {
                             maybeAddImport(fullyQualifiedTarget.getPackageName(), fullyQualifiedTarget.getClassName(), null, true);
                         }
+                        maybeUpdateAlias();
                     }
                 }
 
@@ -279,6 +293,7 @@ public J visitMethodInvocation(J.MethodInvocation method, ExecutionContext ctx)
                             if (fqn != null && TypeUtils.isOfClassType(fqn, originalType.getFullyQualifiedName()) &&
                                 method.getSimpleName().equals(anImport.getQualid().getSimpleName())) {
                                 maybeAddImport(((JavaType.FullyQualified) targetType).getFullyQualifiedName(), method.getName().getSimpleName());
+                                maybeUpdateAlias();
                                 break;
                             }
                         }
@@ -288,6 +303,15 @@ public J visitMethodInvocation(J.MethodInvocation method, ExecutionContext ctx)
             return super.visitMethodInvocation(method, ctx);
         }
 
+        private void maybeUpdateAlias() {
+            if (importAlias != null) {
+                UpdateImportAlias visitor = new UpdateImportAlias(targetType, importAlias);
+                if (!getAfterVisit().contains(visitor)) {
+                    doAfterVisit(visitor);
+                }
+            }
+        }
+
         private Expression updateOuterClassTypes(Expression typeTree) {
             if (typeTree instanceof J.FieldAccess) {
                 JavaType.FullyQualified type = (JavaType.FullyQualified) targetType;
@@ -437,6 +461,25 @@ private boolean isTargetFullyQualifiedType(@Nullable JavaType.FullyQualified fq)
         }
     }
 
+
+
+    // Visitor to backfill the missing kotlin alias import
+    @AllArgsConstructor
+    private static class UpdateImportAlias extends JavaIsoVisitor<ExecutionContext> {
+        @NonNull
+        private final JavaType targetType;
+        @NonNull
+        private J.Identifier importAlias;
+
+        @Override
+        public J.Import visitImport(J.Import _import, ExecutionContext executionContext) {
+            if (hasSameFQN(_import, targetType)) {
+                return _import.withAlias(importAlias);
+            }
+            return _import;
+        }
+    }
+
     private static class ChangeClassDefinition extends JavaIsoVisitor<ExecutionContext> {
         private final JavaType.Class originalType;
         private final JavaType.Class targetType;
@@ -590,4 +633,14 @@ public static JavaType.FullyQualified getTopLevelClassName(JavaType.FullyQualifi
         }
         return getTopLevelClassName(classType.getOwningClass());
     }
+
+    public static boolean hasSameFQN(J.Import import_, JavaType targetType) {
+        JavaType.FullyQualified type = TypeUtils.asFullyQualified(targetType);
+        String fqn = type != null ? type.getFullyQualifiedName() : null;
+
+        JavaType.FullyQualified curType = TypeUtils.asFullyQualified(Optional.ofNullable(import_.getQualid()).map(J.FieldAccess::getType).orElse(null));
+        String curFqn = curType != null ? curType.getFullyQualifiedName() : null;
+
+        return fqn != null && fqn.equals(curFqn);
+    }
 }