From 193979fb34b0ea3b82f98d58b2784e4e259acfc4 Mon Sep 17 00:00:00 2001
From: Stephan Schroevers <stephan.schroevers@teampicnic.com>
Date: Sun, 27 Nov 2022 15:40:15 +0100
Subject: [PATCH] Suggestions

---
 .../refasterrules/ReactorRules.java           | 47 +++++--------------
 .../refasterrules/ReactorRulesTestInput.java  | 28 ++++-------
 .../refasterrules/ReactorRulesTestOutput.java | 28 ++++-------
 3 files changed, 32 insertions(+), 71 deletions(-)

diff --git a/error-prone-contrib/src/main/java/tech/picnic/errorprone/refasterrules/ReactorRules.java b/error-prone-contrib/src/main/java/tech/picnic/errorprone/refasterrules/ReactorRules.java
index 34fcc11e03..d3695bec7f 100644
--- a/error-prone-contrib/src/main/java/tech/picnic/errorprone/refasterrules/ReactorRules.java
+++ b/error-prone-contrib/src/main/java/tech/picnic/errorprone/refasterrules/ReactorRules.java
@@ -320,7 +320,10 @@ Flux<T> after(Flux<T> flux) {
   static final class FluxConcatMap<T, S> {
     @BeforeTemplate
     Flux<S> before(Flux<T> flux, Function<? super T, ? extends Publisher<? extends S>> function) {
-      return Refaster.anyOf(flux.flatMap(function, 1), flux.flatMapSequential(function, 1));
+      return Refaster.anyOf(
+          flux.flatMap(function, 1),
+          flux.flatMapSequential(function, 1),
+          flux.map(function).concatMap(identity()));
     }
 
     @AfterTemplate
@@ -337,7 +340,9 @@ Flux<S> before(
         Function<? super T, ? extends Publisher<? extends S>> function,
         int prefetch) {
       return Refaster.anyOf(
-          flux.flatMap(function, 1, prefetch), flux.flatMapSequential(function, 1, prefetch));
+          flux.flatMap(function, 1, prefetch),
+          flux.flatMapSequential(function, 1, prefetch),
+          flux.map(function).concatMap(identity(), prefetch));
     }
 
     @AfterTemplate
@@ -640,57 +645,31 @@ Flux<S> after(Flux<T> flux) {
   }
 
   /** Prefer {@link Mono#flatMap(Function)} over more contrived alternatives. */
-  static final class MonoFlatMapIdentity<S, T> {
+  static final class MonoFlatMap<S, T> {
     @BeforeTemplate
-    Mono<T> before(Mono<S> mono, Function<S, Mono<T>> function) {
+    Mono<T> before(Mono<S> mono, Function<? super S, ? extends Mono<? extends T>> function) {
       return mono.map(function).flatMap(identity());
     }
 
     @AfterTemplate
-    Mono<T> after(Mono<S> mono, Function<S, Mono<T>> function) {
+    Mono<T> after(Mono<S> mono, Function<? super S, ? extends Mono<? extends T>> function) {
       return mono.flatMap(function);
     }
   }
 
   /** Prefer {@link Mono#flatMapMany(Function)} over more contrived alternatives. */
-  static final class FlatMapManyIdentity<S, T> {
+  static final class MonoFlatMapMany<S, T> {
     @BeforeTemplate
-    Flux<T> before(Mono<S> mono, Function<S, ? extends Publisher<T>> function) {
+    Flux<T> before(Mono<S> mono, Function<? super S, ? extends Publisher<? extends T>> function) {
       return mono.map(function).flatMapMany(identity());
     }
 
     @AfterTemplate
-    Flux<T> after(Mono<S> mono, Function<S, ? extends Publisher<T>> function) {
+    Flux<T> after(Mono<S> mono, Function<? super S, ? extends Publisher<? extends T>> function) {
       return mono.flatMapMany(function);
     }
   }
 
-  /** Prefer {@link Flux#concatMap(Function)} over more contrived alternatives. */
-  static final class ConcatMapIdentity<S, T> {
-    @BeforeTemplate
-    Flux<T> before(Flux<S> flux, Function<S, ? extends Publisher<T>> function) {
-      return flux.map(function).concatMap(identity());
-    }
-
-    @AfterTemplate
-    Flux<T> after(Flux<S> flux, Function<S, ? extends Publisher<T>> function) {
-      return flux.concatMap(function);
-    }
-  }
-
-  /** Prefer {@link Flux#concatMap(Function, int)} over more contrived alternatives. */
-  static final class ConcatMapIdentityWithPrefetch<S, T> {
-    @BeforeTemplate
-    Flux<T> before(Flux<S> flux, Function<S, ? extends Publisher<T>> function, int prefetch) {
-      return flux.map(function).concatMap(identity(), prefetch);
-    }
-
-    @AfterTemplate
-    Flux<T> after(Flux<S> flux, Function<S, ? extends Publisher<T>> function, int prefetch) {
-      return flux.concatMap(function, prefetch);
-    }
-  }
-
   /**
    * Prefer {@link Flux#concatMapIterable(Function)} over alternatives that require an additional
    * subscription.
diff --git a/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/ReactorRulesTestInput.java b/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/ReactorRulesTestInput.java
index f7c5543000..daf40d369f 100644
--- a/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/ReactorRulesTestInput.java
+++ b/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/ReactorRulesTestInput.java
@@ -105,12 +105,16 @@ ImmutableSet<Flux<Integer>> testFluxSwitchIfEmptyOfEmptyPublisher() {
 
   ImmutableSet<Flux<Integer>> testFluxConcatMap() {
     return ImmutableSet.of(
-        Flux.just(1).flatMap(Mono::just, 1), Flux.just(2).flatMapSequential(Mono::just, 1));
+        Flux.just(1).flatMap(Mono::just, 1),
+        Flux.just(2).flatMapSequential(Mono::just, 1),
+        Flux.just(3).map(Mono::just).concatMap(identity()));
   }
 
   ImmutableSet<Flux<Integer>> testFluxConcatMapWithPrefetch() {
     return ImmutableSet.of(
-        Flux.just(1).flatMap(Mono::just, 1, 3), Flux.just(2).flatMapSequential(Mono::just, 1, 4));
+        Flux.just(1).flatMap(Mono::just, 1, 3),
+        Flux.just(2).flatMapSequential(Mono::just, 1, 4),
+        Flux.just(3).map(Mono::just).concatMap(identity(), 5));
   }
 
   Flux<Integer> testFluxConcatMapIterable() {
@@ -208,26 +212,12 @@ Flux<Number> testFluxCast() {
     return Flux.just(1).map(Number.class::cast);
   }
 
-  Mono<String> testMonoFlatMapIdentity() {
+  Mono<String> testMonoFlatMap() {
     return Mono.just("foo").map(Mono::just).flatMap(identity());
   }
 
-  ImmutableSet<Flux<String>> testFlatMapManyIdentity() {
-    return ImmutableSet.of(
-        Mono.just("foo").map(Mono::just).flatMapMany(identity()),
-        Mono.just("foo").map(Flux::just).flatMapMany(identity()));
-  }
-
-  ImmutableSet<Flux<String>> testConcatMapIdentity() {
-    return ImmutableSet.of(
-        Flux.just("foo", "bar").map(Mono::just).concatMap(identity()),
-        Flux.just("foo", "bar").map(Flux::just).concatMap(identity()));
-  }
-
-  ImmutableSet<Flux<String>> testConcatMapIdentityWithPrefetch() {
-    return ImmutableSet.of(
-        Flux.just("foo", "bar").map(Mono::just).concatMap(identity(), 1),
-        Flux.just("foo", "bar").map(Flux::just).concatMap(identity(), 1));
+  Flux<String> testMonoFlatMapMany() {
+    return Mono.just("foo").map(Mono::just).flatMapMany(identity());
   }
 
   ImmutableSet<Flux<String>> testConcatMapIterableIdentity() {
diff --git a/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/ReactorRulesTestOutput.java b/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/ReactorRulesTestOutput.java
index ddf1bc5fb6..df6fb70af4 100644
--- a/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/ReactorRulesTestOutput.java
+++ b/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/ReactorRulesTestOutput.java
@@ -108,12 +108,17 @@ ImmutableSet<Flux<Integer>> testFluxSwitchIfEmptyOfEmptyPublisher() {
   }
 
   ImmutableSet<Flux<Integer>> testFluxConcatMap() {
-    return ImmutableSet.of(Flux.just(1).concatMap(Mono::just), Flux.just(2).concatMap(Mono::just));
+    return ImmutableSet.of(
+        Flux.just(1).concatMap(Mono::just),
+        Flux.just(2).concatMap(Mono::just),
+        Flux.just(3).concatMap(Mono::just));
   }
 
   ImmutableSet<Flux<Integer>> testFluxConcatMapWithPrefetch() {
     return ImmutableSet.of(
-        Flux.just(1).concatMap(Mono::just, 3), Flux.just(2).concatMap(Mono::just, 4));
+        Flux.just(1).concatMap(Mono::just, 3),
+        Flux.just(2).concatMap(Mono::just, 4),
+        Flux.just(3).concatMap(Mono::just, 5));
   }
 
   Flux<Integer> testFluxConcatMapIterable() {
@@ -206,25 +211,12 @@ Flux<Number> testFluxCast() {
     return Flux.just(1).cast(Number.class);
   }
 
-  Mono<String> testMonoFlatMapIdentity() {
+  Mono<String> testMonoFlatMap() {
     return Mono.just("foo").flatMap(Mono::just);
   }
 
-  ImmutableSet<Flux<String>> testFlatMapManyIdentity() {
-    return ImmutableSet.of(
-        Mono.just("foo").flatMapMany(Mono::just), Mono.just("foo").flatMapMany(Flux::just));
-  }
-
-  ImmutableSet<Flux<String>> testConcatMapIdentity() {
-    return ImmutableSet.of(
-        Flux.just("foo", "bar").concatMap(Mono::just),
-        Flux.just("foo", "bar").concatMap(Flux::just));
-  }
-
-  ImmutableSet<Flux<String>> testConcatMapIdentityWithPrefetch() {
-    return ImmutableSet.of(
-        Flux.just("foo", "bar").concatMap(Mono::just, 1),
-        Flux.just("foo", "bar").concatMap(Flux::just, 1));
+  Flux<String> testMonoFlatMapMany() {
+    return Mono.just("foo").flatMapMany(Mono::just);
   }
 
   ImmutableSet<Flux<String>> testConcatMapIterableIdentity() {