Add StreamFlatMapOptional error-prone check (#2946)

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.

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

@@ -0,0 +1,79 @@
+/*
+ * (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.fixes.SuggestedFixes;
+import com.google.errorprone.matchers.ChildMultiMatcher.MatchType;
+import com.google.errorprone.matchers.Description;
+import com.google.errorprone.matchers.Matcher;
+import com.google.errorprone.matchers.Matchers;
+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.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();
+
+    private static final Matcher<ExpressionTree> STREAM_FLATMAP_OPTIONAL_STREAM = Matchers.methodInvocation(
+            STREAM_FLAT_MAP,
+            // Any of the three MatchTypes are reasonable in this case, given a single arg
+            MatchType.LAST,
+            OPTIONAL_STREAM);
+
+    @Override
+    public Description matchMethodInvocation(MethodInvocationTree tree, VisitorState state) {
+        if (STREAM_FLATMAP_OPTIONAL_STREAM.matches(tree, state)) {
+            ExpressionTree receiver = ASTHelpers.getReceiver(tree.getMethodSelect());
+            if (receiver != null) {
+                SuggestedFix.Builder fix = SuggestedFix.builder();
+                String optionalType = SuggestedFixes.qualifyType(state, fix, Optional.class.getCanonicalName());
+                String replacement = ".filter(" + optionalType + "::isPresent).map(" + optionalType + "::get)";
+                return buildDescription(tree)
+                        .addFix(fix.replace(state.getEndPosition(receiver), state.getEndPosition(tree), replacement)
+                                .build())
+                        .build();
+            }
+        }
+        return Description.NO_MATCH;
+    }
+}

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

@@ -0,0 +1,72 @@
+/*
+ * (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.Optional;",
+                        "import java.util.stream.Stream;",
+                        "public class Test {",
+                        "  Stream<String> f1(Stream<Collection<Optional<String>>> in) {",
+                        "    return in.flatMap(Collection::stream).flatMap(Optional::stream);",
+                        "  }",
+                        "  Stream<String> f2(Stream<Collection<Optional<String>>> in) {",
+                        "    return in.flatMap(list -> list.stream().flatMap(Optional::stream));",
+                        "  }",
+                        "  Stream<String> f3(Stream<Collection<Optional<String>>> in) {",
+                        "    return in.flatMap(list -> list.stream()).flatMap(Optional::stream);",
+                        "  }",
+                        "  Stream<String> f4(Stream<Optional<Optional<String>>> in) {",
+                        "      return in.flatMap(Optional::stream).flatMap(Optional::stream);",
+                        "  }",
+                        "}")
+                .addOutputLines(
+                        "Test.java",
+                        "import java.util.Collection;",
+                        "import java.util.Optional;",
+                        "import java.util.stream.Stream;",
+                        "public class Test {",
+                        "  Stream<String> f1(Stream<Collection<Optional<String>>> in) {",
+                        "    return in.flatMap(Collection::stream)"
+                                + ".filter(Optional::isPresent).map(Optional::get);",
+                        "  }",
+                        "  Stream<String> f2(Stream<Collection<Optional<String>>> in) {",
+                        "    return in.flatMap(list -> list.stream()"
+                                + ".filter(Optional::isPresent).map(Optional::get));",
+                        "  }",
+                        "  Stream<String> f3(Stream<Collection<Optional<String>>> in) {",
+                        "    return in.flatMap(list -> list.stream())"
+                                + ".filter(Optional::isPresent).map(Optional::get);",
+                        "  }",
+                        "  Stream<String> f4(Stream<Optional<Optional<String>>> in) {",
+                        "      return in.filter(Optional::isPresent).map(Optional::get)"
+                                + ".filter(Optional::isPresent).map(Optional::get);",
+                        "  }",
+                        "}")
+                .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",