From ea5a7060058252e16d8f2b963be0e55dc6f11c3c Mon Sep 17 00:00:00 2001
From: Dima Legeza <dmytro.legeza@gmail.com>
Date: Wed, 29 Nov 2023 17:40:16 +0100
Subject: [PATCH] * Implement the Mono type matcher preserving generic type

---
 .../bugpatterns/MonoZipOfMonoVoidUsage.java   | 40 ++++++++++++++-----
 .../MonoZipOfMonoVoidUsageTest.java           |  2 -
 2 files changed, 30 insertions(+), 12 deletions(-)

diff --git a/error-prone-contrib/src/main/java/tech/picnic/errorprone/bugpatterns/MonoZipOfMonoVoidUsage.java b/error-prone-contrib/src/main/java/tech/picnic/errorprone/bugpatterns/MonoZipOfMonoVoidUsage.java
index 0e1f1ac9151..498207ec027 100644
--- a/error-prone-contrib/src/main/java/tech/picnic/errorprone/bugpatterns/MonoZipOfMonoVoidUsage.java
+++ b/error-prone-contrib/src/main/java/tech/picnic/errorprone/bugpatterns/MonoZipOfMonoVoidUsage.java
@@ -14,6 +14,7 @@
 import static tech.picnic.errorprone.bugpatterns.util.MoreTypes.type;
 
 import com.google.auto.service.AutoService;
+import com.google.common.collect.Iterables;
 import com.google.errorprone.BugPattern;
 import com.google.errorprone.VisitorState;
 import com.google.errorprone.bugpatterns.BugChecker;
@@ -26,6 +27,7 @@
 import com.sun.source.tree.ExpressionTree;
 import com.sun.source.tree.MethodInvocationTree;
 import com.sun.tools.javac.code.Type;
+import com.sun.tools.javac.tree.JCTree;
 
 /**
  * A {@link BugChecker} that flags usages of Mono.zip(Mono, Mono) and Mono.zipWith(Mono) with
@@ -68,7 +70,7 @@ public final class MonoZipOfMonoVoidUsage extends BugChecker
               instanceMethod().onDescendantOf(MONO).named("zipWith"),
               toType(MethodInvocationTree.class, staticMethod().onClass(MONO).named("empty"))),
           allOf(
-              instanceMethod().onExactClass(MONO_VOID_TYPE).named("zipWith"),
+              onClassWithMethodName(MONO_VOID_TYPE, "zipWith"),
               toType(MethodInvocationTree.class, hasArgumentOfType(MONO))));
 
   // On Mono.zip, at least one element should match empty in order to proceed.
@@ -96,6 +98,28 @@ public Description matchMethodInvocation(MethodInvocationTree tree, VisitorState
         .build();
   }
 
+  private static Matcher<ExpressionTree> onClassWithMethodName(
+      Supplier<Type> genericDesiredType, String methodName) {
+    return (tree, state) -> {
+      JCTree.JCFieldAccess methodExecuted =
+          (JCTree.JCFieldAccess) ((JCTree.JCMethodInvocation) tree).getMethodSelect();
+      Type type = methodExecuted.selected.type;
+      String invokedMethodName = methodExecuted.getIdentifier().toString();
+      return invokedMethodName.equals(methodName)
+          && hasSameGenericType(type, genericDesiredType.get(state), state);
+    };
+  }
+
+  private static Matcher<MethodInvocationTree> hasArgumentOfType(
+      Supplier<Type> genericDesiredType) {
+    return (tree, state) ->
+        tree.getArguments().stream()
+            .anyMatch(
+                arg ->
+                    hasSameGenericType(
+                        ASTHelpers.getType(arg), genericDesiredType.get(state), state));
+  }
+
   /**
    * We need to extract real types from the generics because {@link ASTHelpers} cannot distinguish
    * Mono&lt;Integer&gt; and Mono&lt;Void&gt; and reports those being the same.
@@ -112,14 +136,10 @@ public Description matchMethodInvocation(MethodInvocationTree tree, VisitorState
    *
    * <p>In this case we will always have only one parameter.
    */
-  private static Matcher<MethodInvocationTree> hasArgumentOfType(Supplier<Type> type) {
-    return (tree, state) ->
-        tree.getArguments().stream()
-            .anyMatch(
-                arg -> {
-                  Type argumentType = ASTHelpers.getType(arg).allparams().get(0);
-                  Type requiredType = type.get(state).allparams().get(0);
-                  return isSameType(argumentType, requiredType, state);
-                });
+  private static boolean hasSameGenericType(
+      Type genericArgumentType, Type genericDesiredType, VisitorState state) {
+    Type argumentType = Iterables.getOnlyElement(genericArgumentType.allparams());
+    Type requiredType = Iterables.getOnlyElement(genericDesiredType.allparams());
+    return isSameType(argumentType, requiredType, state);
   }
 }
diff --git a/error-prone-contrib/src/test/java/tech/picnic/errorprone/bugpatterns/MonoZipOfMonoVoidUsageTest.java b/error-prone-contrib/src/test/java/tech/picnic/errorprone/bugpatterns/MonoZipOfMonoVoidUsageTest.java
index cdda2dbebc4..24daf8054b6 100644
--- a/error-prone-contrib/src/test/java/tech/picnic/errorprone/bugpatterns/MonoZipOfMonoVoidUsageTest.java
+++ b/error-prone-contrib/src/test/java/tech/picnic/errorprone/bugpatterns/MonoZipOfMonoVoidUsageTest.java
@@ -21,8 +21,6 @@ void identification() {
             "    // BUG: Diagnostic contains:",
             "    Mono.zip(a, a);",
             "    Mono.zip(e, e);",
-            // TODO: Should not be reported though
-            "    // BUG: Diagnostic contains:",
             "    e.zipWith(e);",
             "    // BUG: Diagnostic contains:",
             "    Mono.zip(d, c, b, a);",