ReplaceClassIsInstanceWithInstanceof: Fix more corner cases

openrewrite · Nov 4, 2024 · f384181 · f384181
1 parent f32b689
commit f384181
Show file tree

Hide file tree

Showing 2 changed files with 111 additions and 6 deletions.
diff --git a/src/main/java/org/openrewrite/staticanalysis/ReplaceClassIsInstanceWithInstanceof.java b/src/main/java/org/openrewrite/staticanalysis/ReplaceClassIsInstanceWithInstanceof.java
@@ -21,12 +21,10 @@
 import org.openrewrite.java.JavaTemplate;
 import org.openrewrite.java.JavaVisitor;
 import org.openrewrite.java.MethodMatcher;
-import org.openrewrite.java.tree.Expression;
-import org.openrewrite.java.tree.J;
+import org.openrewrite.java.tree.*;
 import org.openrewrite.java.tree.J.FieldAccess;
 import org.openrewrite.java.tree.J.Identifier;
 import org.openrewrite.java.tree.J.MethodInvocation;
-import org.openrewrite.java.tree.JavaType;
 
 import java.util.Collections;
 import java.util.Set;
@@ -62,10 +60,13 @@ public J visitMethodInvocation(MethodInvocation method, ExecutionContext ctx) {
                     // for code like "A.class.isInstance(a)", select is "String.class", name is "isInstance", argument is "a"
                     Expression objectExpression = method.getArguments().get(0);
                     FieldAccess fieldAccessPart = (FieldAccess) method.getSelect();
-                    String className = ((JavaType.Class) fieldAccessPart.getTarget().getType()).getClassName();
+                    String className = fieldAccessPart.getTarget().toString();
                     // upcast to type J, so J.MethodInvocation can be replaced by J.InstanceOf
-                    J.InstanceOf instanceOf = JavaTemplate.apply("#{any()} instanceof #{}",
-                            getCursor(), method.getCoordinates().replace(), objectExpression, className);
+                    JavaCoordinates coordinates = method.getCoordinates().replace();
+                    J.InstanceOf instanceOf = JavaTemplate.builder("#{any()} instanceof #{}")
+                            .build()
+                            .apply(getCursor(), coordinates, new Object[]{objectExpression, className});
+                    instanceOf = instanceOf.withClazz(fieldAccessPart.getTarget().withPrefix(instanceOf.getClazz().getPrefix()));
                     return maybeAutoFormat(method, instanceOf, ctx);
                 }
                 return super.visitMethodInvocation(method, ctx);

diff --git a/src/test/java/org/openrewrite/staticanalysis/ReplaceClassIsInstanceWithInstanceofTest.java b/src/test/java/org/openrewrite/staticanalysis/ReplaceClassIsInstanceWithInstanceofTest.java
@@ -161,4 +161,108 @@ boolean foo(Object obj) {
         );
     }
 
+    @Test
+    void imported() {
+        rewriteRun(
+          //language=java
+          java(
+            """
+              import java.util.Map;
+
+              class A {
+                  boolean foo(Object obj) {
+                      return Map.class.isInstance(obj);
+                  }
+              }
+              """,
+            """
+              import java.util.Map;
+
+              class A {
+                  boolean foo(Object obj) {
+                      return obj instanceof Map;
+                  }
+              }
+              """
+          )
+        );
+    }
+
+    @Test
+    void importedNested() {
+        rewriteRun(
+          //language=java
+          java(
+            """
+              import java.util.Map.Entry;
+
+              class A {
+                  boolean foo(Object obj) {
+                      return Entry.class.isInstance(obj);
+                  }
+              }
+              """,
+            """
+              import java.util.Map.Entry;
+
+              class A {
+                  boolean foo(Object obj) {
+                      return obj instanceof Entry;
+                  }
+              }
+              """
+          )
+        );
+    }
+
+    @Test
+    void typeFromSourcePath() {
+        rewriteRun(
+          //language=java
+          java(
+            """
+              package a;
+
+              class A {
+                  boolean foo(Object obj) {
+                      return A.class.isInstance(obj);
+                  }
+              }
+              """,
+            """
+              package a;
+
+              class A {
+                  boolean foo(Object obj) {
+                      return obj instanceof A;
+                  }
+              }
+              """
+          )
+        );
+    }
+
+    @Test
+    void defaultPackage() {
+        rewriteRun(
+          //language=java
+          java(
+            """
+              class A {
+                  boolean foo(Object obj) {
+                      return A.class.isInstance(obj);
+                  }
+              }
+              """,
+            """
+              class A {
+                  boolean foo(Object obj) {
+                      return obj instanceof A;
+                  }
+              }
+              """
+          )
+        );
+    }
+
 }