Fix ReplaceLambdaWithMethodReference handling of nested class imports (#745)

* Fix ReplaceLambdaWithMethodReference handling of nested class imports

When transforming lambdas with instanceof checks or casts to method references,
preserve the original form of the class reference instead of always generating
a fully qualified name. This ensures that if a nested class like Map.Entry is
already imported as `import java.util.Map.Entry`, the transformed code uses
`Entry.class::isInstance` instead of `Map.Entry.class::isInstance`.

The fix directly uses the original expression from the instanceof/cast rather
than reconstructing it, allowing the ImportService to properly handle any
necessary import adjustments while preserving the user's preferred form.

Fixes #744

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-Authored-By: Claude <[email protected]>

* Minimize diff

---------

Co-authored-by: Claude <[email protected]>

Author: timtebeek

src/main/java/org/openrewrite/staticanalysis/JavaElementFactory.java

@@ -120,18 +120,13 @@ static J.MemberReference newInstanceMethodReference(Expression containing, Strin
         );
     }
 
-    static J.@Nullable FieldAccess newClassLiteral(@Nullable JavaType type, boolean qualified) {
-        JavaType.Class classType = getClassType(type);
-        if (classType == null) {
-            return null;
-        }
-
-        JavaType.Parameterized parameterized = new JavaType.Parameterized(null, classType, singletonList(type));
+    static J.FieldAccess newClassLiteral(JavaType.Class classType, JavaType originalType, J instanceOfClass) {
+        JavaType.Parameterized parameterized = new JavaType.Parameterized(null, classType, singletonList(originalType));
         return new J.FieldAccess(
                 randomId(),
                 Space.EMPTY,
                 Markers.EMPTY,
-                className(type, qualified),
+                instanceOfClass.withPrefix(Space.EMPTY), // Use the original expression directly
                 new JLeftPadded<>(
                         Space.EMPTY,
                         new J.Identifier(randomId(), Space.EMPTY, Markers.EMPTY, emptyList(), "class", parameterized, null),
@@ -141,7 +136,7 @@ static J.MemberReference newInstanceMethodReference(Expression containing, Strin
         );
     }
 
-    private static JavaType.@Nullable Class getClassType(@Nullable JavaType type) {
+    static JavaType.@Nullable Class getClassType(@Nullable JavaType type) {
         if (type instanceof JavaType.Class) {
             JavaType.Class classType = (JavaType.Class) type;
             if ("java.lang.Class".equals(classType.getFullyQualifiedName())) {

src/main/java/org/openrewrite/staticanalysis/ReplaceLambdaWithMethodReference.java

@@ -100,8 +100,11 @@ public J visitLambda(J.Lambda lambda, ExecutionContext ctx) {
                 J j = instanceOf.getClazz();
                 if ((j instanceof J.Identifier || j instanceof J.FieldAccess) &&
                         instanceOf.getExpression() instanceof J.Identifier) {
-                    J.FieldAccess classLiteral = newClassLiteral(((TypeTree) j).getType(), j instanceof J.FieldAccess);
-                    if (classLiteral != null) {
+                    // Create the class literal directly from the original expression
+                    JavaType originalType = ((TypeTree) j).getType();
+                    JavaType.Class classType = getClassType(originalType);
+                    if (classType != null) {
+                        J.FieldAccess classLiteral = newClassLiteral(classType, originalType, j);
                         //noinspection DataFlowIssue
                         JavaType.FullyQualified rawClassType = ((JavaType.Parameterized) classLiteral.getType()).getType();
                         Optional<JavaType.Method> isInstanceMethod = rawClassType.getMethods().stream().filter(m -> "isInstance".equals(m.getName())).findFirst();
@@ -123,11 +126,13 @@ public J visitLambda(J.Lambda lambda, ExecutionContext ctx) {
                     J tree = j.getTree();
                     if ((tree instanceof J.Identifier || tree instanceof J.FieldAccess) &&
                             !(j.getType() instanceof JavaType.GenericTypeVariable)) {
-                        J.FieldAccess classLiteral = newClassLiteral(((Expression) tree).getType(), tree instanceof J.FieldAccess);
-                        if (classLiteral != null) {
+                        // Create the class literal directly from the original expression
+                        JavaType.Class classType = getClassType(((Expression) tree).getType());
+                        if (classType != null) {
+                            J.FieldAccess classLiteral = newClassLiteral(classType, ((Expression) tree).getType(), tree);
                             //noinspection DataFlowIssue
-                            JavaType.FullyQualified classType = ((JavaType.Parameterized) classLiteral.getType()).getType();
-                            Optional<JavaType.Method> castMethod = classType.getMethods().stream().filter(m -> "cast".equals(m.getName())).findFirst();
+                            JavaType.FullyQualified fullClassType = ((JavaType.Parameterized) classLiteral.getType()).getType();
+                            Optional<JavaType.Method> castMethod = fullClassType.getMethods().stream().filter(m -> "cast".equals(m.getName())).findFirst();
                             if (castMethod.isPresent()) {
                                 J.MemberReference updated = newInstanceMethodReference(classLiteral, castMethod.get(), lambda.getType()).withPrefix(lambda.getPrefix());
                                 doAfterVisit(service(ImportService.class).shortenFullyQualifiedTypeReferencesIn(updated));
@@ -252,9 +257,7 @@ private boolean hasSelectWhoseReferenceMightChange(MethodCall method) {
                     JavaType.Variable fieldType = ((J.FieldAccess) select).getName().getFieldType();
                     return fieldType != null && fieldType.getOwner() instanceof JavaType.Class && !fieldType.hasFlags(Flag.Final);
                 }
-                if (select instanceof J.NewClass || select instanceof J.Parentheses) {
-                    return true;
-                }
+                return select instanceof J.NewClass || select instanceof J.Parentheses;
             }
             return false;
         }

src/test/java/org/openrewrite/staticanalysis/ReplaceLambdaWithMethodReferenceTest.java

@@ -1487,4 +1487,93 @@ public void run() {
           )
         );
     }
+
+    @Issue("https://github.com/openrewrite/rewrite-static-analysis/issues/744")
+    @Test
+    void methodReferenceNestedClassImportInvalid() {
+        rewriteRun(
+          //language=java
+          java(
+            """
+              import java.util.List;
+              import java.util.Map.Entry;
+
+              public class Foo {
+                  private void foo() {
+                      List.of().stream().filter(i -> i instanceof Entry);
+                  }
+              }
+              """,
+            """
+              import java.util.List;
+              import java.util.Map.Entry;
+
+              public class Foo {
+                  private void foo() {
+                      List.of().stream().filter(Entry.class::isInstance);
+                  }
+              }
+              """
+          )
+        );
+    }
+
+    @Issue("https://github.com/openrewrite/rewrite-static-analysis/issues/744")
+    @Test
+    void methodReferenceNestedClassFullyQualified() {
+        rewriteRun(
+          //language=java
+          java(
+            """
+              import java.util.List;
+
+              public class Foo {
+                  private void foo() {
+                      List.of().stream().filter(i -> i instanceof java.util.Map.Entry);
+                  }
+              }
+              """,
+            """
+              import java.util.List;
+              import java.util.Map;
+
+              public class Foo {
+                  private void foo() {
+                      List.of().stream().filter(Map.Entry.class::isInstance);
+                  }
+              }
+              """
+          )
+        );
+    }
+
+    @Issue("https://github.com/openrewrite/rewrite-static-analysis/issues/744")
+    @Test
+    void methodReferenceNestedClassCast() {
+        rewriteRun(
+          //language=java
+          java(
+            """
+              import java.util.List;
+              import java.util.Map.Entry;
+
+              public class Foo {
+                  private void foo() {
+                      List.of().stream().map(i -> (Entry) i);
+                  }
+              }
+              """,
+            """
+              import java.util.List;
+              import java.util.Map.Entry;
+
+              public class Foo {
+                  private void foo() {
+                      List.of().stream().map(Entry.class::cast);
+                  }
+              }
+              """
+          )
+        );
+    }
 }