Add StreamFlatMapOptional error-prone check

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

@@ -0,0 +1,76 @@
+/*
+ * (c) Copyright 2024 Palantir Technologies Inc. All rights reserved.
+ *
+ * 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
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * 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 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.ExpressionTree;
+import com.sun.source.tree.MethodInvocationTree;
+import java.util.List;
+import java.util.Optional;
+import java.util.function.Function;
+import java.util.stream.Stream;
+
+@AutoService(BugChecker.class)
+@BugPattern(
+        link = "https://github.com/palantir/gradle-baseline#baseline-error-prone-checks",
+        linkType = BugPattern.LinkType.CUSTOM,
+        severity = BugPattern.SeverityLevel.WARNING,
+        summary = "`Stream.filter(Optional::isPresent).map(Optional::get)` is more efficient than "
+                + "`Stream.flatMap(Optional::stream)`")
+public final class StreamFlatMapOptional extends BugChecker implements BugChecker.MethodInvocationTreeMatcher {
+    private static final long serialVersionUID = 1L;
+
+    private static final Matcher<ExpressionTree> STREAM_FLAT_MAP = MethodMatchers.instanceMethod()
+            .onDescendantOf(Stream.class.getName())
+            .namedAnyOf("flatMap")
+            .withParameters(Function.class.getName());
+
+    private static final Matcher<ExpressionTree> OPTIONAL_STREAM = MethodMatchers.instanceMethod()
+            .onDescendantOf(Optional.class.getName())
+            .named("stream")
+            .withNoParameters();
+
+    @Override
+    public Description matchMethodInvocation(MethodInvocationTree tree, VisitorState state) {
+        if (STREAM_FLAT_MAP.matches(tree, state)) {
+            ExpressionTree stream = ASTHelpers.getReceiver(tree);
+            if (stream != null) {
+                List<? extends ExpressionTree> arguments = tree.getArguments();
+                if (arguments != null && arguments.size() == 1) {
+                    if (OPTIONAL_STREAM.matches(arguments.get(0), state)) {
+                        String replacement = state.getSourceForNode(ASTHelpers.getReceiver(tree.getMethodSelect()))
+                                + ".filter(Optional::isPresent).map(Optional::get)";
+                        SuggestedFix fix = SuggestedFix.builder()
+                                .addImport("java.util.Optional")
+                                .replace(tree, replacement)
+                                .build();
+                        return buildDescription(tree).addFix(fix).build();
+                    }
+                }
+            }
+        }
+        return Description.NO_MATCH;
+    }
+}

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

@@ -0,0 +1,62 @@
+/*
+ * (c) Copyright 2024 Palantir Technologies Inc. All rights reserved.
+ *
+ * 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
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * 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 com.palantir.baseline.errorprone;
+
+import org.junit.jupiter.api.Test;
+
+class StreamFlatmapOptionalTest {
+    @Test
+    public void test() {
+        fix().addInputLines(
+                        "Test.java",
+                        "import java.util.Collection;",
+                        "import java.util.List;",
+                        "import java.util.Optional;",
+                        "import java.util.stream.Collectors;",
+                        "public class Test {",
+                        "  List<String> f1(List<List<Optional<String>>> in) {",
+                        "    return in.stream().flatMap(Collection::stream)"
+                                + ".flatMap(Optional::stream).collect(Collectors.toList());",
+                        "  }",
+                        "  List<String> f2(List<List<Optional<String>>> in) {",
+                        "    return in.stream().flatMap(list -> list.stream().flatMap(Optional::stream))"
+                                + ".collect(Collectors.toList());",
+                        "  }",
+                        "}")
+                .addOutputLines(
+                        "Test.java",
+                        "import java.util.Collection;",
+                        "import java.util.List;",
+                        "import java.util.Optional;",
+                        "import java.util.stream.Collectors;",
+                        "public class Test {",
+                        "  List<String> f1(List<List<Optional<String>>> in) {",
+                        "    return in.stream().flatMap(Collection::stream)"
+                                + ".filter(Optional::isPresent).map(Optional::get).collect(Collectors.toList());",
+                        "  }",
+                        "  List<String> f2(List<List<Optional<String>>> in) {",
+                        "    return in.stream().flatMap(list -> list.stream()"
+                                + ".filter(Optional::isPresent).map(Optional::get)).collect(Collectors.toList());",
+                        "  }",
+                        "}")
+                .doTest();
+    }
+
+    private RefactoringValidator fix() {
+        return RefactoringValidator.of(StreamFlatMapOptional.class, getClass());
+    }
+}

changelog/@unreleased/pr-2946.v2.yml

@@ -0,0 +1,8 @@
+type: improvement
+improvement:
+  description: |-
+    If one has a `Stream<Optional<T>> stream` of size `N` and does `stream.flatMap(Optional::stream)`, you’ll end up allocating `N` extra streams — one for each `Optional` input element. When `N` is large, those allocations can cause extra GC cycles and pauses if allocation rate is high enough leading to issues with latency, throughput, and allocation sensitive code paths.
+
+    `Stream.filter(Optional::isPresent).map(Optional::get)` is more efficient than `Stream.flatMap(Optional::stream)` as it does not allocate a new `Stream` for every element in the stream.
+  links:
+  - https://github.com/palantir/gradle-baseline/pull/2946

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

@@ -72,6 +72,7 @@ public class BaselineErrorProneExtension {
             "Slf4jLogsafeArgs",
             "Slf4jThrowable",
             "StreamOfEmpty",
+            "StreamFlatMapOptional",
             "StrictUnusedVariable",
             "StringBuilderConstantParameters",
             "ThrowError",