From 990b11c46a9eb1911874aa4beda5818774f6d91d Mon Sep 17 00:00:00 2001
From: Bastien Diederichs <bastien.diederichs@teampicnic.com>
Date: Wed, 5 Oct 2022 17:59:17 +0200
Subject: [PATCH] ANS-1289 Extend to concatMap

---
 .../refastertemplates/ReactorTemplates.java   | 28 ++++++++++++++++---
 .../ReactorTemplatesTestInput.java            | 12 ++++++--
 .../ReactorTemplatesTestOutput.java           | 12 ++++++--
 3 files changed, 44 insertions(+), 8 deletions(-)

diff --git a/error-prone-contrib/src/main/java/tech/picnic/errorprone/refastertemplates/ReactorTemplates.java b/error-prone-contrib/src/main/java/tech/picnic/errorprone/refastertemplates/ReactorTemplates.java
index ac66bdbf0f2..981537fbb69 100644
--- a/error-prone-contrib/src/main/java/tech/picnic/errorprone/refastertemplates/ReactorTemplates.java
+++ b/error-prone-contrib/src/main/java/tech/picnic/errorprone/refastertemplates/ReactorTemplates.java
@@ -273,19 +273,39 @@ Flux<S> after(Flux<T> flux) {
   }
 
   /**
-   * Prefer {@link Flux#flatMapIterable(Function)} (Class)} over {@link Flux#flatMap(Function)} with
+   * Prefer {@link Flux#concatMapIterable(Function)} over {@link Flux#concatMap(Function)} with
    * {@link Flux#fromIterable(Iterable)}.
    */
-  static final class FluxFlatMapIterable<S> {
+  static final class FluxConcatMapFromIterable<S> {
     @BeforeTemplate
     Flux<S> before(Flux<? extends Iterable<S>> flux) {
-      return flux.flatMap(list -> Flux.fromIterable(list));
+      return Refaster.anyOf(
+          flux.concatMap(list -> Flux.fromIterable(list)), flux.concatMap(Flux::fromIterable));
     }
 
     @AfterTemplate
     @UseImportPolicy(STATIC_IMPORT_ALWAYS)
     Flux<S> after(Flux<? extends Iterable<S>> flux) {
-      return flux.flatMapIterable(identity());
+      return flux.concatMapIterable(identity());
+    }
+  }
+
+  /**
+   * Prefer {@link Flux#flatMapIterable(Function, int)} over {@link Flux#flatMap(Function, int)}
+   * with {@link Flux#fromIterable(Iterable)}.
+   */
+  static final class FluxFlatMapFromIterable<S> {
+    @BeforeTemplate
+    Flux<S> before(Flux<? extends Iterable<S>> flux, int concurrency) {
+      return Refaster.anyOf(
+          flux.flatMap(list -> Flux.fromIterable(list), concurrency),
+          flux.flatMap(Flux::fromIterable, concurrency));
+    }
+
+    @AfterTemplate
+    @UseImportPolicy(STATIC_IMPORT_ALWAYS)
+    Flux<S> after(Flux<? extends Iterable<S>> flux, int concurrency) {
+      return flux.flatMapIterable(identity(), concurrency);
     }
   }
 
diff --git a/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refastertemplates/ReactorTemplatesTestInput.java b/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refastertemplates/ReactorTemplatesTestInput.java
index 3aa63d8a5c8..570e667981e 100644
--- a/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refastertemplates/ReactorTemplatesTestInput.java
+++ b/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refastertemplates/ReactorTemplatesTestInput.java
@@ -95,8 +95,16 @@ Flux<Number> testFluxCast() {
     return Flux.just(1).map(Number.class::cast);
   }
 
-  Flux<String> testFluxFlatMapIterable() {
-    return Flux.just(ImmutableList.of("1")).flatMap(list -> Flux.fromIterable(list));
+  ImmutableSet<Flux<String>> testFluxConcatMapFromIterable() {
+    return ImmutableSet.of(
+        Flux.just(ImmutableList.of("1")).concatMap(list -> Flux.fromIterable(list)),
+        Flux.just(ImmutableList.of("1")).concatMap(Flux::fromIterable));
+  }
+
+  ImmutableSet<Flux<String>> testFluxFlatMapFromIterable() {
+    return ImmutableSet.of(
+        Flux.just(ImmutableList.of("1")).flatMap(list -> Flux.fromIterable(list), 1),
+        Flux.just(ImmutableList.of("1")).flatMap(Flux::fromIterable, 2));
   }
 
   ImmutableSet<PublisherProbe<Void>> testPublisherProbeEmpty() {
diff --git a/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refastertemplates/ReactorTemplatesTestOutput.java b/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refastertemplates/ReactorTemplatesTestOutput.java
index 34b39e12c81..7bda8462528 100644
--- a/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refastertemplates/ReactorTemplatesTestOutput.java
+++ b/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refastertemplates/ReactorTemplatesTestOutput.java
@@ -95,8 +95,16 @@ Flux<Number> testFluxCast() {
     return Flux.just(1).cast(Number.class);
   }
 
-  Flux<String> testFluxFlatMapIterable() {
-    return Flux.just(ImmutableList.of("1")).flatMapIterable(identity());
+  ImmutableSet<Flux<String>> testFluxConcatMapFromIterable() {
+    return ImmutableSet.of(
+        Flux.just(ImmutableList.of("1")).concatMapIterable(identity()),
+        Flux.just(ImmutableList.of("1")).concatMapIterable(identity()));
+  }
+
+  ImmutableSet<Flux<String>> testFluxFlatMapFromIterable() {
+    return ImmutableSet.of(
+        Flux.just(ImmutableList.of("1")).flatMapIterable(identity(), 1),
+        Flux.just(ImmutableList.of("1")).flatMapIterable(identity(), 2));
   }
 
   ImmutableSet<PublisherProbe<Void>> testPublisherProbeEmpty() {