AnnotateNullableMethods: annotate methods that return nullable method… (#738)

* AnnotateNullableMethods: annotate methods that return nullable method invocations

* Update src/main/java/org/openrewrite/staticanalysis/AnnotateNullableMethods.java

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* Minor polish

* Clarify scope of search and use `reduce`

---------

Co-authored-by: stefanod <stefano.dallapalma@adyen.com>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: Tim te Beek <tim@moderne.io>

Author: 4 people

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

@@ -29,6 +29,7 @@
 import java.util.Arrays;
 import java.util.Comparator;
 import java.util.List;
+import java.util.Optional;
 import java.util.concurrent.atomic.AtomicBoolean;
 
 @EqualsAndHashCode(callSuper = false)
@@ -86,7 +87,7 @@ public J.MethodDeclaration visitMethodDeclaration(J.MethodDeclaration methodDecl
 
                 J.MethodDeclaration md = super.visitMethodDeclaration(methodDeclaration, ctx);
                 updateCursor(md);
-                if (FindNullableReturnStatements.find(md.getBody(), getCursor().getParentTreeCursor())) {
+                if (FindNullableReturnStatements.find(md.getBody(), getCursor().getParentTreeCursor(), nullableAnnotationClass)) {
                     J.MethodDeclaration annotatedMethod = JavaTemplate.builder("@" + fullyQualifiedName)
                             .javaParser(JavaParser.fromJavaVersion().dependsOn(
                                     String.format("package %s;public @interface %s {}", fullyQualifiedPackage, simpleName)))
@@ -145,8 +146,14 @@ private static class FindNullableReturnStatements extends JavaIsoVisitor<AtomicB
                 new MethodMatcher("java.util.Spliterator trySplit(..)")
         );
 
-        static boolean find(@Nullable J subtree, Cursor parentTreeCursor) {
-            return new FindNullableReturnStatements().reduce(subtree, new AtomicBoolean(), parentTreeCursor).get();
+        private final String nullableAnnotationClass;
+
+        private FindNullableReturnStatements(@Nullable String nullableAnnotationClass) {
+            this.nullableAnnotationClass = Optional.ofNullable(nullableAnnotationClass).orElse(DEFAULT_NULLABLE_ANN_CLASS);
+        }
+
+        static boolean find(@Nullable J subtree, Cursor parentTreeCursor, @Nullable String nullableAnnotationClass) {
+            return new FindNullableReturnStatements(nullableAnnotationClass).reduce(subtree, new AtomicBoolean(), parentTreeCursor).get();
         }
 
         @Override
@@ -176,7 +183,7 @@ private boolean maybeIsNull(@Nullable Expression returnExpression) {
                 return ((J.Literal) returnExpression).getValue() == null;
             }
             if (returnExpression instanceof J.MethodInvocation) {
-                return isKnowNullableMethod((J.MethodInvocation) returnExpression);
+                return isLocalNullableMethod((J.MethodInvocation) returnExpression) || isKnownNullableMethod((J.MethodInvocation) returnExpression);
             }
             if (returnExpression instanceof J.Ternary) {
                 J.Ternary ternary = (J.Ternary) returnExpression;
@@ -185,13 +192,42 @@ private boolean maybeIsNull(@Nullable Expression returnExpression) {
             return false;
         }
 
-        private boolean isKnowNullableMethod(J.MethodInvocation methodInvocation) {
+        private boolean isKnownNullableMethod(J.MethodInvocation methodInvocation) {
             for (MethodMatcher m : KNOWN_NULLABLE_METHODS) {
                 if (m.matches(methodInvocation)) {
                     return true;
                 }
             }
             return false;
         }
+
+        private boolean isLocalNullableMethod(J.MethodInvocation methodInvocation) {
+            JavaType.Method targetMethod = methodInvocation.getMethodType();
+            if (targetMethod == null) {
+                return false;
+            }
+
+            // Visit the entire compilation unit to find the method declaration
+            AnnotationMatcher annotationMatcher = new AnnotationMatcher("@" + nullableAnnotationClass);
+            J.CompilationUnit cu = getCursor().firstEnclosingOrThrow(J.CompilationUnit.class);
+            return new JavaIsoVisitor<AtomicBoolean>() {
+                @Override
+                public J.MethodDeclaration visitMethodDeclaration(J.MethodDeclaration method, AtomicBoolean p) {
+                    if (p.get()) {
+                        return method;
+                    }
+                    if (targetMethod.equals(method.getMethodType()) && method.getReturnTypeExpression() instanceof J.AnnotatedType) {
+                        for (J.Annotation annotation : ((J.AnnotatedType) method.getReturnTypeExpression()).getAnnotations()) {
+                            if (annotationMatcher.matches(annotation)) {
+                                p.set(true);
+                                break;
+                            }
+                        }
+                        return method;
+                    }
+                    return super.visitMethodDeclaration(method, p);
+                }
+            }.reduce(cu, new AtomicBoolean(false)).get();
+        }
     }
 }

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

@@ -373,4 +373,83 @@ public class Foo {
           )
         );
     }
+
+    @Test
+    void nullableMethodsInvocationsWithDefaultNullableClass() {
+        rewriteRun(
+          //language=java
+          java(
+            """
+              import org.jspecify.annotations.Nullable;
+
+              import java.util.Random;
+
+              public class Test {
+                  public @Nullable String maybeNullString() {
+                      return new Random().nextBoolean() ? "Not null" : null;
+                  }
+
+                  public String getString() {
+                      return maybeNullString();
+                  }
+              }
+              """,
+            """
+              import org.jspecify.annotations.Nullable;
+
+              import java.util.Random;
+
+              public class Test {
+                  public @Nullable String maybeNullString() {
+                      return new Random().nextBoolean() ? "Not null" : null;
+                  }
+
+                  public @Nullable String getString() {
+                      return maybeNullString();
+                  }
+              }
+              """
+          )
+        );
+    }
+
+    @Test
+    void nullableMethodsInvocationsWithCustomNullableClass() {
+        rewriteRun(
+          spec -> spec.recipe(new AnnotateNullableMethods("org.openrewrite.jgit.annotations.Nullable")),
+          //language=java
+          java(
+            """
+              import org.openrewrite.jgit.annotations.Nullable;
+
+              import java.util.Random;
+
+              public class Test {
+                  public @Nullable String maybeNullString() {
+                      return new Random().nextBoolean() ? "Not null" : null;
+                  }
+
+                  public String getString() {
+                      return maybeNullString();
+                  }
+              }
+              """,
+            """
+              import org.openrewrite.jgit.annotations.Nullable;
+
+              import java.util.Random;
+
+              public class Test {
+                  public @Nullable String maybeNullString() {
+                      return new Random().nextBoolean() ? "Not null" : null;
+                  }
+
+                  public @Nullable String getString() {
+                      return maybeNullString();
+                  }
+              }
+              """
+          )
+        );
+    }
 }