Add AnnotateNullableMethods recipe

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

@@ -0,0 +1,183 @@
+/*
+ * Copyright 2024 the original author or authors.
+ * <p>
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ * <p>
+ * https://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.openrewrite.staticanalysis;
+
+import lombok.AllArgsConstructor;
+import org.openrewrite.Cursor;
+import org.openrewrite.ExecutionContext;
+import org.openrewrite.Recipe;
+import org.openrewrite.TreeVisitor;
+import org.openrewrite.java.*;
+import org.openrewrite.java.service.AnnotationService;
+import org.openrewrite.java.tree.Expression;
+import org.openrewrite.java.tree.J;
+import org.openrewrite.java.tree.JavaType;
+
+import java.util.Arrays;
+import java.util.Comparator;
+import java.util.List;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+public class AnnotateNullableMethods extends Recipe {
+
+    private static final String NULLABLE_ANN_CLASS = "org.jspecify.annotations.Nullable";
+    private static final AnnotationMatcher NULLABLE_ANNOTATION_MATCHER =
+            new AnnotationMatcher("@" + NULLABLE_ANN_CLASS);
+
+    @Override
+    public String getDisplayName() {
+        return "Annotate methods which may return null with @Nullable";
+    }
+
+    @Override
+    public String getDescription() {
+        return "Automatically adds the @org.jspecify.annotation.Nullable to non-private methods" +
+                "that may return null. This recipe scans for methods that do not already have a @Nullable" +
+                "annotation and checks their return statements for potential null values. It also" +
+                "identifies known methods from standard libraries that may return null, such as methods" +
+                "from Map, Queue, Deque, NavigableSet, and Spliterator. The return of streams, or lambdas" +
+                " are not taken into account.";
+    }
+
+    @Override
+    public TreeVisitor<?, ExecutionContext> getVisitor() {
+        return new AnnotateNullableMethodsVisitor();
+    }
+
+    private static class AnnotateNullableMethodsVisitor extends JavaIsoVisitor<ExecutionContext> {
+        AtomicBoolean annotatedNullable = new AtomicBoolean(false);
+
+        @Override
+        public J.MethodDeclaration visitMethodDeclaration(J.MethodDeclaration md, ExecutionContext ctx) {
+            if (!md.hasModifier(J.Modifier.Type.Public)) {
+                return md;
+            }
+
+            if (md.getMethodType() != null && md.getMethodType().getReturnType() instanceof JavaType.Primitive) {
+                return md;
+            }
+
+            if (service(AnnotationService.class).matches(getCursor(), NULLABLE_ANNOTATION_MATCHER)) {
+                return md;
+            }
+
+            md = super.visitMethodDeclaration(md, ctx);
+            updateCursor(md);
+
+            if (FindNullableReturnStatements.find(md).get()) {
+                maybeAddImport(NULLABLE_ANN_CLASS);
+                if (!annotatedNullable.getAndSet(true)) {
+                    doAfterVisit(new NullableOnMethodReturnType().getVisitor());
+                }
+                return JavaTemplate.builder("@Nullable")
+                        .imports(NULLABLE_ANN_CLASS)
+                        .javaParser(JavaParser.fromJavaVersion().classpath("jspecify"))
+                        .build()
+                        .apply(getCursor(), md.getCoordinates().addAnnotation(Comparator.comparing(J.Annotation::getSimpleName)));
+            }
+            return md;
+        }
+    }
+
+    @AllArgsConstructor
+    private static class FindNullableReturnStatements extends JavaIsoVisitor<AtomicBoolean> {
+        private static final List<MethodMatcher> KNOWN_NULLABLE_METHODS = getMatchersKnownNullableMethods();
+
+        static AtomicBoolean find(J subtree) {
+            return new FindNullableReturnStatements().reduce(subtree, new AtomicBoolean());
+        }
+
+        private static List<MethodMatcher> getMatchersKnownNullableMethods() {
+            return Arrays.asList(
+                    new MethodMatcher("java.util.Map computeIfAbsent(..)"),
+                    new MethodMatcher("java.util.Map computeIfPresent(..)"),
+                    new MethodMatcher("java.util.Map get(..)"),
+                    new MethodMatcher("java.util.Map merge(..)"),
+                    new MethodMatcher("java.util.Map put(..)"),
+                    new MethodMatcher("java.util.Map putIfAbsent(..)"),
+
+                    new MethodMatcher("java.util.Queue poll(..)"),
+                    new MethodMatcher("java.util.Queue peek(..)"),
+
+                    new MethodMatcher("java.util.Deque peekFirst(..)"),
+                    new MethodMatcher("java.util.Deque pollFirst(..)"),
+                    new MethodMatcher("java.util.Deque peekLast(..)"),
+
+                    new MethodMatcher("java.util.NavigableSet lower(..)"),
+                    new MethodMatcher("java.util.NavigableSet floor(..)"),
+                    new MethodMatcher("java.util.NavigableSet ceiling(..)"),
+                    new MethodMatcher("java.util.NavigableSet higher(..)"),
+                    new MethodMatcher("java.util.NavigableSet pollFirst(..)"),
+                    new MethodMatcher("java.util.NavigableSet pollLast(..)"),
+
+                    new MethodMatcher("java.util.NavigableMap lowerEntry(..)"),
+                    new MethodMatcher("java.util.NavigableMap floorEntry(..)"),
+                    new MethodMatcher("java.util.NavigableMap ceilingEntry(..)"),
+                    new MethodMatcher("java.util.NavigableMap higherEntry(..)"),
+                    new MethodMatcher("java.util.NavigableMap lowerKey(..)"),
+                    new MethodMatcher("java.util.NavigableMap floorKey(..)"),
+                    new MethodMatcher("java.util.NavigableMap ceilingKey(..)"),
+                    new MethodMatcher("java.util.NavigableMap higherKey(..)"),
+                    new MethodMatcher("java.util.NavigableMap firstEntry(..)"),
+                    new MethodMatcher("java.util.NavigableMap lastEntry(..)"),
+                    new MethodMatcher("java.util.NavigableMap pollFirstEntry(..)"),
+                    new MethodMatcher("java.util.NavigableMap pollLastEntry(..)"),
+
+                    new MethodMatcher("java.util.Spliterator trySplit(..)")
+            );
+        }
+
+        @Override
+        public J.Return visitReturn(J.Return retrn, AtomicBoolean containsNullableReturn) {
+            if (containsNullableReturn.get()) {
+                return retrn;
+            }
+
+            J.Return r = super.visitReturn(retrn, containsNullableReturn);
+            updateCursor(r);
+
+            // If the returns is contained within a lambda statement, we don't consider it.
+            Cursor ex = getCursor().dropParentUntil(e -> e instanceof J.MethodDeclaration || e instanceof J.Lambda);
+            if (!(ex.getValue() instanceof J.MethodDeclaration)) {
+                return r;
+            }
+
+            if (r.getExpression() != null && maybeIsNull(r.getExpression())) {
+                containsNullableReturn.set(true);
+            }
+
+            return r;
+        }
+
+        private boolean maybeIsNull(Expression returnExpression) {
+            if (returnExpression instanceof J.Literal && ((J.Literal) returnExpression).getValue() == null) {
+                return true;
+            } else if (returnExpression instanceof J.MethodInvocation) {
+                return isKnowNullableMethod((J.MethodInvocation) returnExpression);
+            }
+            return false;
+        }
+
+        private boolean isKnowNullableMethod(J.MethodInvocation methodInvocation) {
+            for (MethodMatcher m : KNOWN_NULLABLE_METHODS) {
+                if (m.matches(methodInvocation)) {
+                    return true;
+                }
+            }
+            return false;
+        }
+    }
+}

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

@@ -0,0 +1,197 @@
+/*
+ * Copyright 2024 the original author or authors.
+ * <p>
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ * <p>
+ * https://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.openrewrite.staticanalysis;
+
+import org.junit.jupiter.api.Test;
+import org.openrewrite.DocumentExample;
+import org.openrewrite.java.JavaParser;
+import org.openrewrite.test.RecipeSpec;
+import org.openrewrite.test.RewriteTest;
+
+import static org.openrewrite.java.Assertions.java;
+
+class AnnotateNullableMethodsTest implements RewriteTest {
+    @Override
+    public void defaults(RecipeSpec spec) {
+        spec.recipe(new AnnotateNullableMethods()).parser(JavaParser.fromJavaVersion().classpath("jspecify"));
+    }
+
+    @DocumentExample
+    @Test
+    void methodReturnsNullLiteral() {
+        rewriteRun(
+          spec -> spec.recipe(new AnnotateNullableMethods()),
+          //language=java
+          java(
+            """
+              public class Test {
+
+                  public String getString() {
+                      return null;
+                  }
+
+                  public String getStringWithMultipleReturn() {
+                      if (System.currentTimeMillis() % 2 == 0) {
+                          return "Not null";
+                      }
+                      return null;
+                  }
+              }
+              """,
+            """
+              import org.jspecify.annotations.Nullable;
+
+              public class Test {      
+
+                  public @Nullable String getString() {
+                      return null;
+                  }
+
+                  public @Nullable String getStringWithMultipleReturn() {
+                      if (System.currentTimeMillis() % 2 == 0) {
+                          return "Not null";
+                      }
+                      return null;
+                  }
+              }
+              """
+          )
+        );
+    }
+
+    @Test
+    void methodReturnNullButIsAlreadyAnnotated() {
+        rewriteRun(
+          spec -> spec.recipe(new AnnotateNullableMethods()),
+          //language=java
+          java(
+            """
+              import org.jspecify.annotations.Nullable;
+
+              public class Test {
+                  public @Nullable String getString() {
+                      return null;
+                  }
+
+                  public @Nullable String getStringWithMultipleReturn() {
+                      if (System.currentTimeMillis() % 2 == 0) {
+                          return "Not null";
+                      }
+                      return null;
+                  }
+              }
+              """
+          )
+        );
+    }
+
+    @Test
+    void methodDoesNotReturnNull() {
+        rewriteRun(
+          spec -> spec.recipe(new AnnotateNullableMethods()),
+          //language=java
+          java(
+            """
+              package org.example;
+
+              public class Test {
+                  public String getString() {
+                      return "Hello";
+                  }
+              }
+              """
+          )
+        );
+    }
+
+    @Test
+    void methodReturnsDelegateKnowNullableMethod() {
+        rewriteRun(
+          spec -> spec.recipe(new AnnotateNullableMethods()),
+          //language=java
+          java(
+            """
+              import java.util.HashMap;
+              import java.util.Map;
+
+              public class Test {
+
+                  public String getString() {
+                      Map<String, String> map = new HashMap<>();
+                      return map.get("key");
+                  }
+              }
+              """,
+            """
+              import org.jspecify.annotations.Nullable;
+
+              import java.util.HashMap;
+              import java.util.Map;
+
+              public class Test {
+
+                  public @Nullable String getString() {
+                      Map<String, String> map = new HashMap<>();
+                      return map.get("key");
+                  }
+              }
+              """
+          )
+        );
+    }
+
+    @Test
+    void methodWithLambdaShouldNotBeAnnotated() {
+        rewriteRun(
+          //language=java
+          java(
+            """
+                import java.util.stream.Stream;
+              class A {
+                  public Runnable getRunnable() {
+                      return () -> null;
+                  }
+
+                  public Integer someStream(){
+                      // Stream with lambda class.
+                        return Stream.of(1, 2, 3)
+                            .map(i -> {if (i == 2) return null; else return i;})
+                            .reduce((a, b) -> a + b)
+                            .orElse(null);
+                  }
+              }
+              """
+          )
+        );
+    }
+
+    @Test
+    void privateMethodsShouldNotBeAnnotated() {
+        rewriteRun(
+          spec -> spec.recipe(new AnnotateNullableMethods()),
+          //language=java
+          java(
+            """
+              public class Test {
+                  private String getString() {
+                      return null;
+                  }
+              }
+              """
+          )
+        );
+    }
+}