From 62481aa342b0c84b58da8beb08d992ee3e25678f Mon Sep 17 00:00:00 2001
From: mohamedsamehsalah <mohamed.sameh@teampicnic.com>
Date: Wed, 10 Jul 2024 20:53:52 +0200
Subject: [PATCH] Missed a spot

---
 .../picnic/errorprone/refasterrules/EqualityRules.java    | 8 +++-----
 .../errorprone/refasterrules/EqualityRulesTestInput.java  | 7 +++++--
 .../errorprone/refasterrules/EqualityRulesTestOutput.java | 7 +++++--
 3 files changed, 13 insertions(+), 9 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 80f5bcc2f0..ce1ae28f95 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
@@ -44,13 +44,11 @@ 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>> {
+  /** Prefer reference-based equality for enums over more contrived equality checks. */
+  static final class EnumsReferenceEqualityLambda<T extends Enum<T>> {
     @BeforeTemplate
     Predicate<T> before(T a) {
-      return isEqual(a);
+      return Refaster.anyOf(isEqual(a), a::equals);
     }
 
     @AfterTemplate
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 af0ffcf305..fe50095996 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
@@ -34,8 +34,11 @@ boolean testEqualsPredicate() {
     return Stream.of("foo").anyMatch(s -> "bar".equals(s));
   }
 
-  boolean testPredicateIsEqualEnums() {
-    return Stream.of(RoundingMode.UP).anyMatch(isEqual(RoundingMode.DOWN));
+  ImmutableSet<Boolean> testEnumsReferenceEqualityLambda() {
+
+    return ImmutableSet.of(
+        Stream.of(RoundingMode.UP).anyMatch(isEqual(RoundingMode.DOWN)),
+        Stream.of(RoundingMode.UP).anyMatch(RoundingMode.DOWN::equals));
   }
 
   boolean testDoubleNegation() {
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 c56e21a4f0..9884f8efb7 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
@@ -34,8 +34,11 @@ boolean testEqualsPredicate() {
     return Stream.of("foo").anyMatch("bar"::equals);
   }
 
-  boolean testPredicateIsEqualEnums() {
-    return Stream.of(RoundingMode.UP).anyMatch(v -> v == RoundingMode.DOWN);
+  ImmutableSet<Boolean> testEnumsReferenceEqualityLambda() {
+
+    return ImmutableSet.of(
+        Stream.of(RoundingMode.UP).anyMatch(v -> v == RoundingMode.DOWN),
+        Stream.of(RoundingMode.UP).anyMatch(v -> v == RoundingMode.DOWN));
   }
 
   boolean testDoubleNegation() {