From 78b0417422c4523bbe390396b1001ee639f3feb0 Mon Sep 17 00:00:00 2001
From: Stephan Schroevers <stephan.schroevers@teampicnic.com>
Date: Thu, 11 May 2023 09:14:27 +0200
Subject: [PATCH] Suggestions

---
 .../picnic/errorprone/refasterrules/StreamRules.java     | 4 ++--
 .../errorprone/refasterrules/StreamRulesTestInput.java   | 9 ++++++---
 .../errorprone/refasterrules/StreamRulesTestOutput.java  | 7 +++++--
 3 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/error-prone-contrib/src/main/java/tech/picnic/errorprone/refasterrules/StreamRules.java b/error-prone-contrib/src/main/java/tech/picnic/errorprone/refasterrules/StreamRules.java
index 1eaafc41e29..fdd5f2e77fb 100644
--- a/error-prone-contrib/src/main/java/tech/picnic/errorprone/refasterrules/StreamRules.java
+++ b/error-prone-contrib/src/main/java/tech/picnic/errorprone/refasterrules/StreamRules.java
@@ -3,6 +3,7 @@
 import static com.google.errorprone.refaster.ImportPolicy.STATIC_IMPORT_ALWAYS;
 import static java.util.Comparator.naturalOrder;
 import static java.util.Comparator.reverseOrder;
+import static java.util.function.Function.identity;
 import static java.util.function.Predicate.not;
 import static java.util.stream.Collectors.counting;
 import static java.util.stream.Collectors.filtering;
@@ -627,9 +628,8 @@ R after(
 
   static final class StreamsConcat<T> {
     @BeforeTemplate
-    @SuppressWarnings("StreamOfArray")
     Stream<T> before(@Repeated Stream<T> stream) {
-      return Stream.of(Refaster.asVarargs(stream)).flatMap(Function.identity());
+      return Stream.of(Refaster.asVarargs(stream)).flatMap(Refaster.anyOf(identity(), s -> s));
     }
 
     @AfterTemplate
diff --git a/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/StreamRulesTestInput.java b/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/StreamRulesTestInput.java
index b3c3bfee94a..a9ec0e4ca1d 100644
--- a/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/StreamRulesTestInput.java
+++ b/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/StreamRulesTestInput.java
@@ -3,6 +3,7 @@
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static java.util.Comparator.comparingInt;
 import static java.util.Comparator.reverseOrder;
+import static java.util.function.Function.identity;
 import static java.util.function.Predicate.not;
 import static java.util.stream.Collectors.counting;
 import static java.util.stream.Collectors.filtering;
@@ -40,6 +41,7 @@ public ImmutableSet<Object> elidedTypesAndStaticImports() {
         counting(),
         filtering(null, null),
         flatMapping(null, null),
+        identity(),
         mapping(null, null),
         maxBy(null),
         minBy(null),
@@ -251,8 +253,9 @@ ImmutableSet<Integer> testStreamFlatMapCollect() {
     return Stream.of(1).collect(flatMapping(n -> Stream.of(n, n), toImmutableSet()));
   }
 
-  Stream<Integer> testStreamsConcat() {
-    return Stream.of(Stream.of(1), Stream.of(2), Stream.of(3), Stream.of(4))
-        .flatMap(Function.identity());
+  ImmutableSet<Stream<Integer>> testStreamsConcat() {
+    return ImmutableSet.of(
+        Stream.of(Stream.of(1), Stream.of(2)).flatMap(identity()),
+        Stream.of(Stream.of(3), Stream.of(4)).flatMap(v -> v));
   }
 }
diff --git a/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/StreamRulesTestOutput.java b/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/StreamRulesTestOutput.java
index 14066851e8f..ab69c763bc1 100644
--- a/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/StreamRulesTestOutput.java
+++ b/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/StreamRulesTestOutput.java
@@ -4,6 +4,7 @@
 import static java.util.Comparator.comparingInt;
 import static java.util.Comparator.naturalOrder;
 import static java.util.Comparator.reverseOrder;
+import static java.util.function.Function.identity;
 import static java.util.function.Predicate.not;
 import static java.util.stream.Collectors.counting;
 import static java.util.stream.Collectors.filtering;
@@ -42,6 +43,7 @@ public ImmutableSet<Object> elidedTypesAndStaticImports() {
         counting(),
         filtering(null, null),
         flatMapping(null, null),
+        identity(),
         mapping(null, null),
         maxBy(null),
         minBy(null),
@@ -252,7 +254,8 @@ ImmutableSet<Integer> testStreamFlatMapCollect() {
     return Stream.of(1).flatMap(n -> Stream.of(n, n)).collect(toImmutableSet());
   }
 
-  Stream<Integer> testStreamsConcat() {
-    return Streams.concat(Stream.of(1), Stream.of(2), Stream.of(3), Stream.of(4));
+  ImmutableSet<Stream<Integer>> testStreamsConcat() {
+    return ImmutableSet.of(
+        Streams.concat(Stream.of(1), Stream.of(2)), Streams.concat(Stream.of(3), Stream.of(4)));
   }
 }