From 2c79bde919a8b8712478f4d458201aac34223296 Mon Sep 17 00:00:00 2001
From: mohamedsamehsalah <mohamed.sameh@teampicnic.com>
Date: Mon, 8 Jul 2024 16:30:29 +0200
Subject: [PATCH] Introduce `PredicateIsEqualEnums`

---
 .../errorprone/refasterrules/EqualityRules.java  | 16 ++++++++++++++++
 .../refasterrules/EqualityRulesTestInput.java    |  7 ++++++-
 .../refasterrules/EqualityRulesTestOutput.java   |  7 ++++++-
 3 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/error-prone-contrib/src/main/java/tech/picnic/errorprone/refasterrules/EqualityRules.java b/error-prone-contrib/src/main/java/tech/picnic/errorprone/refasterrules/EqualityRules.java
index 7d5f57c2346..80f5bcc2f0d 100644
--- a/error-prone-contrib/src/main/java/tech/picnic/errorprone/refasterrules/EqualityRules.java
+++ b/error-prone-contrib/src/main/java/tech/picnic/errorprone/refasterrules/EqualityRules.java
@@ -1,5 +1,6 @@
 package tech.picnic.errorprone.refasterrules;
 
+import static java.util.function.Predicate.isEqual;
 import static java.util.function.Predicate.not;
 
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
@@ -43,6 +44,21 @@ boolean after(T a, T b) {
     }
   }
 
+  /**
+   * Prefer reference-based equality for enums over {@link Predicate#isEqual(Object)} comparison.
+   */
+  static final class PredicateIsEqualEnums<T extends Enum<T>> {
+    @BeforeTemplate
+    Predicate<T> before(T a) {
+      return isEqual(a);
+    }
+
+    @AfterTemplate
+    Predicate<T> after(T a) {
+      return v -> v == a;
+    }
+  }
+
   /** Prefer {@link Object#equals(Object)} over the equivalent lambda function. */
   // XXX: As it stands, this rule is a special case of what `MethodReferenceUsage` tries to achieve.
   // If/when `MethodReferenceUsage` becomes production ready, we should simply drop this check.
diff --git a/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/EqualityRulesTestInput.java b/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/EqualityRulesTestInput.java
index 7cd4e3bfb4c..af0ffcf3058 100644
--- a/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/EqualityRulesTestInput.java
+++ b/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/EqualityRulesTestInput.java
@@ -1,5 +1,6 @@
 package tech.picnic.errorprone.refasterrules;
 
+import static java.util.function.Predicate.isEqual;
 import static java.util.function.Predicate.not;
 
 import com.google.common.collect.BoundType;
@@ -14,7 +15,7 @@
 final class EqualityRulesTest implements RefasterRuleCollectionTestCase {
   @Override
   public ImmutableSet<Object> elidedTypesAndStaticImports() {
-    return ImmutableSet.of(Objects.class, Optional.class, not(null));
+    return ImmutableSet.of(Objects.class, Optional.class, isEqual(null), not(null));
   }
 
   ImmutableSet<Boolean> testPrimitiveOrReferenceEquality() {
@@ -33,6 +34,10 @@ boolean testEqualsPredicate() {
     return Stream.of("foo").anyMatch(s -> "bar".equals(s));
   }
 
+  boolean testPredicateIsEqualEnums() {
+    return Stream.of(RoundingMode.UP).anyMatch(isEqual(RoundingMode.DOWN));
+  }
+
   boolean testDoubleNegation() {
     return !!Boolean.TRUE;
   }
diff --git a/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/EqualityRulesTestOutput.java b/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/EqualityRulesTestOutput.java
index 0b8af555a71..c56e21a4f0a 100644
--- a/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/EqualityRulesTestOutput.java
+++ b/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/EqualityRulesTestOutput.java
@@ -1,5 +1,6 @@
 package tech.picnic.errorprone.refasterrules;
 
+import static java.util.function.Predicate.isEqual;
 import static java.util.function.Predicate.not;
 
 import com.google.common.collect.BoundType;
@@ -14,7 +15,7 @@
 final class EqualityRulesTest implements RefasterRuleCollectionTestCase {
   @Override
   public ImmutableSet<Object> elidedTypesAndStaticImports() {
-    return ImmutableSet.of(Objects.class, Optional.class, not(null));
+    return ImmutableSet.of(Objects.class, Optional.class, isEqual(null), not(null));
   }
 
   ImmutableSet<Boolean> testPrimitiveOrReferenceEquality() {
@@ -33,6 +34,10 @@ boolean testEqualsPredicate() {
     return Stream.of("foo").anyMatch("bar"::equals);
   }
 
+  boolean testPredicateIsEqualEnums() {
+    return Stream.of(RoundingMode.UP).anyMatch(v -> v == RoundingMode.DOWN);
+  }
+
   boolean testDoubleNegation() {
     return Boolean.TRUE;
   }