fix #1326: Implement new OptionalFlatmapOfNullable check

README.md

@@ -189,6 +189,7 @@ Safe Logging can be found at [github.com/palantir/safe-logging](https://github.c
 - `UnnecessaryLambdaArgumentParentheses`: Lambdas with a single parameter do not require argument parentheses.
 - `RawTypes`: Avoid raw types; add appropriate type parameters if possible.
 - `VisibleForTestingPackagePrivate`: `@VisibleForTesting` members should be package-private.
+- `OptionalFlatMapOfNullable`: Optional.map functions may return null to safely produce an empty result.
 
 ### Programmatic Application
 

baseline-error-prone/src/main/java/com/palantir/baseline/errorprone/OptionalFlatMapOfNullable.java

@@ -0,0 +1,94 @@
+package com.palantir.baseline.errorprone;
+
+import com.google.auto.service.AutoService;
+import com.google.errorprone.BugPattern;
+import com.google.errorprone.VisitorState;
+import com.google.errorprone.bugpatterns.BugChecker;
+import com.google.errorprone.fixes.SuggestedFix;
+import com.google.errorprone.matchers.Description;
+import com.google.errorprone.matchers.Matcher;
+import com.google.errorprone.matchers.method.MethodMatchers;
+import com.google.errorprone.util.ASTHelpers;
+import com.sun.source.tree.BlockTree;
+import com.sun.source.tree.ExpressionTree;
+import com.sun.source.tree.LambdaExpressionTree;
+import com.sun.source.tree.MethodInvocationTree;
+import com.sun.source.tree.ReturnTree;
+import com.sun.source.tree.StatementTree;
+import com.sun.source.tree.Tree;
+
+import java.util.List;
+import java.util.Optional;
+import java.util.function.Function;
+
+@AutoService(BugChecker.class)
+@BugPattern(
+        name = "OptionalFlatMapOfNullable",
+        link = "https://github.com/palantir/gradle-baseline#baseline-error-prone-checks",
+        linkType = BugPattern.LinkType.CUSTOM,
+        severity = BugPattern.SeverityLevel.WARNING,
+        summary = "Optional.map functions may return null to safely produce an empty result.")
+public final class OptionalFlatMapOfNullable extends BugChecker implements BugChecker.MethodInvocationTreeMatcher {
+
+    private static final Matcher<ExpressionTree> flatMap = MethodMatchers.instanceMethod()
+            .onDescendantOf(Optional.class.getName())
+            .named("flatMap")
+            .withParameters(Function.class.getName());
+
+    private static final Matcher<ExpressionTree> ofNullable = MethodMatchers.staticMethod()
+            .onClass(Optional.class.getName())
+            .named("ofNullable")
+            .withParameters(Object.class.getName());
+
+    @Override
+    public Description matchMethodInvocation(MethodInvocationTree tree, VisitorState state) {
+        if (!flatMap.matches(tree, state)) {
+            return Description.NO_MATCH;
+        }
+        ExpressionTree functionParameter =
+                ASTHelpers.stripParentheses(tree.getArguments().get(0));
+        if (functionParameter.getKind() != Tree.Kind.LAMBDA_EXPRESSION) {
+            return Description.NO_MATCH;
+        }
+        LambdaExpressionTree lambdaExpressionTree = (LambdaExpressionTree) functionParameter;
+        Optional<ExpressionTree> maybeExpression = finalExpression(lambdaExpressionTree);
+        if (!maybeExpression.isPresent()) {
+            return Description.NO_MATCH;
+        }
+        ExpressionTree expression = maybeExpression.get();
+        if (!ofNullable.matches(expression, state)) {
+            return Description.NO_MATCH;
+        }
+        MethodInvocationTree ofNullableInvocation = (MethodInvocationTree) expression;
+        ExpressionTree ofNullableArg = ofNullableInvocation.getArguments().get(0);
+        return buildDescription(tree)
+                .addFix(SuggestedFix.builder()
+                        .merge(MoreSuggestedFixes.renameMethodInvocation(tree, "map", state))
+                        .replace(ofNullableInvocation, state.getSourceForNode(ofNullableArg))
+                        .build())
+                .build();
+    }
+
+    private Optional<ExpressionTree> finalExpression(LambdaExpressionTree lambdaExpressionTree) {
+        Tree body = lambdaExpressionTree.getBody();
+        switch (lambdaExpressionTree.getBodyKind()) {
+            case EXPRESSION:
+                return Optional.of((ExpressionTree) body);
+            case STATEMENT:
+                if (body instanceof BlockTree) {
+                    BlockTree block = (BlockTree) body;
+                    List<? extends StatementTree> statements = block.getStatements();
+                    if (!statements.isEmpty()) {
+                        StatementTree finalStatement = statements.get(statements.size() - 1);
+                        if (finalStatement instanceof ReturnTree) {
+                            ReturnTree returnTree = (ReturnTree) finalStatement;
+                            return Optional.ofNullable(returnTree.getExpression());
+                        }
+                    }
+                }
+                break;
+        }
+        // Don't break compilation when new language features are introduced.
+        return Optional.empty();
+    }
+}

baseline-error-prone/src/test/java/com/palantir/baseline/errorprone/OptionalFlatMapOfNullableTest.java

@@ -0,0 +1,79 @@
+package com.palantir.baseline.errorprone;
+
+import org.junit.jupiter.api.Test;
+
+class OptionalFlatMapOfNullableTest {
+
+    @Test
+    void testFix_expression() {
+        fix().addInputLines(
+                        "Test.java",
+                        "import java.util.Optional;",
+                        "class Test {",
+                        "   Optional<?> f(Optional<String> in) {",
+                        "       return in.flatMap(x -> Optional.ofNullable(x));",
+                        "   }",
+                        "}")
+                .addOutputLines(
+                        "Test.java",
+                        "import java.util.Optional;",
+                        "class Test {",
+                        "   Optional<?> f(Optional<String> in) {",
+                        "       return in.map(x -> x);",
+                        "   }",
+                        "}")
+                .doTest();
+    }
+
+    @Test
+    void testFix_statement() {
+        fix().addInputLines(
+                        "Test.java",
+                        "import java.util.Optional;",
+                        "class Test {",
+                        "   Optional<?> f(Optional<String> in) {",
+                        "       return in.flatMap(x -> { return Optional.ofNullable(x); });",
+                        "   }",
+                        "}")
+                .addOutputLines(
+                        "Test.java",
+                        "import java.util.Optional;",
+                        "class Test {",
+                        "   Optional<?> f(Optional<String> in) {",
+                        "       return in.map(x -> { return x; });",
+                        "   }",
+                        "}")
+                .doTest();
+    }
+
+    @Test
+    void testFix_statement_additionalStatements() {
+        fix().addInputLines(
+                        "Test.java",
+                        "import java.util.Optional;",
+                        "class Test {",
+                        "   Optional<?> f(Optional<String> in) {",
+                        "       return in.flatMap(x -> {",
+                        "         String y = x + x;",
+                        "         return Optional.ofNullable(y);",
+                        "       });",
+                        "   }",
+                        "}")
+                .addOutputLines(
+                        "Test.java",
+                        "import java.util.Optional;",
+                        "class Test {",
+                        "   Optional<?> f(Optional<String> in) {",
+                        "       return in.map(x -> {",
+                        "         String y = x + x;",
+                        "         return y;",
+                        "       });",
+                        "   }",
+                        "}")
+                .doTest();
+    }
+
+    private RefactoringValidator fix() {
+        return RefactoringValidator.of(new OptionalFlatMapOfNullable(), getClass());
+    }
+}

gradle-baseline-java/src/main/groovy/com/palantir/baseline/extensions/BaselineErrorProneExtension.java

@@ -32,6 +32,7 @@ public class BaselineErrorProneExtension {
             "FinalClass",
             "LambdaMethodReference",
             "LoggerEnclosingClass",
+            "OptionalFlatMapOfNullable",
             "OptionalOrElseMethodInvocation",
             "PreferBuiltInConcurrentKeySet",
             "PreferCollectionTransform",