From 5d362c27a1f72e7bcff2d8e9818d003f77943a0a Mon Sep 17 00:00:00 2001
From: Rick Ossendrijver <rick.ossies@hotmail.com>
Date: Sat, 15 Jan 2022 11:43:06 +0100
Subject: [PATCH] Introduce template to rewrite `flatMapIterable` to
 `concatMapIterable`

---
 .../refastertemplates/ReactorTemplates.java   | 20 +++++++++++++++++++
 .../ReactorTemplatesTestInput.java            |  5 +++++
 .../ReactorTemplatesTestOutput.java           |  5 +++++
 3 files changed, 30 insertions(+)

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 27e7f6fa5b..c58b67f7a4 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
@@ -21,6 +21,7 @@
 import reactor.core.publisher.Mono;
 import reactor.test.StepVerifier;
 import reactor.test.publisher.PublisherProbe;
+import tech.picnic.errorprone.bugpatterns.FluxFlatMapUsageCheck;
 
 /** Refaster templates related to Reactor expressions and statements. */
 final class ReactorTemplates {
@@ -155,6 +156,25 @@ Flux<S> after(Flux<T> flux, Function<? super T, ? extends Publisher<? extends S>
     }
   }
 
+  /**
+   * Prefer {@link Flux#concatMapIterable(Function)} over {@link Flux#concatMapIterable(Function)}
+   * to be consistent with {@link FluxFlatMapUsageCheck}.
+   *
+   * <p>NB: Both implementations emit values in a deterministic order and there is no difference
+   * with eager or lazy inner subscriptions. This means that both implementations are *equivalent*.
+   */
+  static final class FluxConcatMapIterable<T, S> {
+    @BeforeTemplate
+    Flux<S> before(Flux<T> flux, Function<? super T, ? extends Iterable<? extends S>> function) {
+      return flux.flatMapIterable(function);
+    }
+
+    @AfterTemplate
+    Flux<S> after(Flux<T> flux, Function<? super T, ? extends Iterable<? extends S>> function) {
+      return flux.concatMapIterable(function);
+    }
+  }
+
   /**
    * Don't use {@link Mono#flatMapMany(Function)} to implicitly convert a {@link Mono} to a {@link
    * Flux}.
diff --git a/error-prone-contrib/src/test/resources/tech/picnic/errorprone/bugpatterns/ReactorTemplatesTestInput.java b/error-prone-contrib/src/test/resources/tech/picnic/errorprone/bugpatterns/ReactorTemplatesTestInput.java
index 26dd52a1be..add57fb20c 100644
--- a/error-prone-contrib/src/test/resources/tech/picnic/errorprone/bugpatterns/ReactorTemplatesTestInput.java
+++ b/error-prone-contrib/src/test/resources/tech/picnic/errorprone/bugpatterns/ReactorTemplatesTestInput.java
@@ -1,5 +1,6 @@
 package tech.picnic.errorprone.bugpatterns;
 
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import java.time.Duration;
 import java.util.Optional;
@@ -50,6 +51,10 @@ ImmutableSet<Flux<Integer>> testFluxConcatMap() {
         Flux.just(1).flatMap(Mono::just, 1), Flux.just(2).flatMapSequential(Mono::just, 1));
   }
 
+  Flux<Integer> testFluxConcatMapIterable() {
+    return Flux.just(1, 2).flatMapIterable(ImmutableList::of);
+  }
+
   Flux<String> testMonoFlatMapToFlux() {
     return Mono.just("foo").flatMapMany(s -> Mono.just(s + s));
   }
diff --git a/error-prone-contrib/src/test/resources/tech/picnic/errorprone/bugpatterns/ReactorTemplatesTestOutput.java b/error-prone-contrib/src/test/resources/tech/picnic/errorprone/bugpatterns/ReactorTemplatesTestOutput.java
index cfc97c2977..3dcff32c5f 100644
--- a/error-prone-contrib/src/test/resources/tech/picnic/errorprone/bugpatterns/ReactorTemplatesTestOutput.java
+++ b/error-prone-contrib/src/test/resources/tech/picnic/errorprone/bugpatterns/ReactorTemplatesTestOutput.java
@@ -2,6 +2,7 @@
 
 import static com.google.common.collect.MoreCollectors.toOptional;
 
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import java.time.Duration;
 import java.util.Optional;
@@ -50,6 +51,10 @@ ImmutableSet<Flux<Integer>> testFluxConcatMap() {
     return ImmutableSet.of(Flux.just(1).concatMap(Mono::just), Flux.just(2).concatMap(Mono::just));
   }
 
+  Flux<Integer> testFluxConcatMapIterable() {
+    return Flux.just(1, 2).concatMapIterable(ImmutableList::of);
+  }
+
   Flux<String> testMonoFlatMapToFlux() {
     return Mono.just("foo").flatMap(s -> Mono.just(s + s)).flux();
   }