From 37f7f17f0233b3602abca31e91acd6a9cfc2bffd Mon Sep 17 00:00:00 2001
From: Phil Werli <phil.werli@teampicnic.com>
Date: Wed, 4 Oct 2023 18:44:52 +0200
Subject: [PATCH] Introduce `{Mono,Flux}OfType` Refaster rules

---
 .../refasterrules/ReactorRules.java           | 34 +++++++++++++++++++
 .../refasterrules/ReactorRulesTestInput.java  | 12 +++++++
 .../refasterrules/ReactorRulesTestOutput.java |  8 +++++
 3 files changed, 54 insertions(+)

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 c92668319f3..06ed6643d3d 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
@@ -794,6 +794,40 @@ Flux<S> after(Flux<T> flux) {
     }
   }
 
+  /** Prefer {@link Mono#ofType(Class)} over more contrived alternatives. */
+  static final class MonoOfType<T, S> {
+    @BeforeTemplate
+    Mono<S> before(Mono<T> mono) {
+      return mono.filter(
+              Refaster.anyOf(
+                  Refaster.<S>clazz()::isInstance,
+                  x -> Refaster.<S>clazz().isAssignableFrom(x.getClass())))
+          .cast(Refaster.<S>clazz());
+    }
+
+    @AfterTemplate
+    Mono<S> after(Mono<T> mono) {
+      return mono.ofType(Refaster.<S>clazz());
+    }
+  }
+
+  /** Prefer {@link Flux#ofType(Class)} over more contrived alternatives. */
+  static final class FluxOfType<T, S> {
+    @BeforeTemplate
+    Flux<S> before(Flux<T> flux) {
+      return flux.filter(
+              Refaster.anyOf(
+                  Refaster.<S>clazz()::isInstance,
+                  x -> Refaster.<S>clazz().isAssignableFrom(x.getClass())))
+          .cast(Refaster.<S>clazz());
+    }
+
+    @AfterTemplate
+    Flux<S> after(Flux<T> flux) {
+      return flux.ofType(Refaster.<S>clazz());
+    }
+  }
+
   /** Prefer {@link Mono#flatMap(Function)} over more contrived alternatives. */
   static final class MonoFlatMap<S, T> {
     @BeforeTemplate
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 7f8021f70ca..441e04be6a8 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
@@ -269,6 +269,18 @@ Flux<Number> testFluxCast() {
     return Flux.just(1).map(Number.class::cast);
   }
 
+  ImmutableSet<Mono<Number>> testMonoOfType() {
+    return ImmutableSet.of(
+        Mono.just(1).filter(Number.class::isInstance).cast(Number.class),
+        Mono.just(2).filter(n -> Number.class.isAssignableFrom(n.getClass())).cast(Number.class));
+  }
+
+  ImmutableSet<Flux<Number>> testFluxOfType() {
+    return ImmutableSet.of(
+        Flux.just(1).filter(Number.class::isInstance).cast(Number.class),
+        Flux.just(2).filter(n -> Number.class.isAssignableFrom(n.getClass())).cast(Number.class));
+  }
+
   Mono<String> testMonoFlatMap() {
     return Mono.just("foo").map(Mono::just).flatMap(identity());
   }
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 b12af508c48..2a30cdf99f7 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
@@ -266,6 +266,14 @@ Flux<Number> testFluxCast() {
     return Flux.just(1).cast(Number.class);
   }
 
+  ImmutableSet<Mono<Number>> testMonoOfType() {
+    return ImmutableSet.of(Mono.just(1).ofType(Number.class), Mono.just(2).ofType(Number.class));
+  }
+
+  ImmutableSet<Flux<Number>> testFluxOfType() {
+    return ImmutableSet.of(Flux.just(1).ofType(Number.class), Flux.just(2).ofType(Number.class));
+  }
+
   Mono<String> testMonoFlatMap() {
     return Mono.just("foo").flatMap(Mono::just);
   }