From beb74e6bbee6cd626bb1304f5c83564f1f1f5af1 Mon Sep 17 00:00:00 2001
From: mohamedsamehsalah <mohamed.sameh@teampicnic.com>
Date: Mon, 30 Sep 2024 13:47:24 +0200
Subject: [PATCH] Extend `ImmutableCollectionIterator` refaster rule

By re-writing `Stream#iterator` statements.
---
 .../picnic/errorprone/refasterrules/CollectionRules.java   | 7 ++-----
 .../errorprone/refasterrules/CollectionRulesTestInput.java | 5 +++--
 .../refasterrules/CollectionRulesTestOutput.java           | 4 ++--
 3 files changed, 7 insertions(+), 9 deletions(-)

diff --git a/error-prone-contrib/src/main/java/tech/picnic/errorprone/refasterrules/CollectionRules.java b/error-prone-contrib/src/main/java/tech/picnic/errorprone/refasterrules/CollectionRules.java
index 4ee37189e1d..376a2c43ed9 100644
--- a/error-prone-contrib/src/main/java/tech/picnic/errorprone/refasterrules/CollectionRules.java
+++ b/error-prone-contrib/src/main/java/tech/picnic/errorprone/refasterrules/CollectionRules.java
@@ -365,14 +365,11 @@ S[] after(ImmutableCollection<T> collection, IntFunction<S[]> generator) {
     }
   }
 
-  /**
-   * Don't call {@link ImmutableCollection#asList()} if {@link ImmutableCollection#iterator()} is
-   * called on the result; call it directly.
-   */
+  /** Prefer {@link ImmutableCollection#iterator()} over more contrived alternatives. */
   static final class ImmutableCollectionIterator<T> {
     @BeforeTemplate
     Iterator<T> before(ImmutableCollection<T> collection) {
-      return collection.asList().iterator();
+      return Refaster.anyOf(collection.stream().iterator(), collection.asList().iterator());
     }
 
     @AfterTemplate
diff --git a/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/CollectionRulesTestInput.java b/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/CollectionRulesTestInput.java
index c667a80895f..75d9472d579 100644
--- a/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/CollectionRulesTestInput.java
+++ b/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/CollectionRulesTestInput.java
@@ -119,8 +119,9 @@ Integer[] testImmutableCollectionToArrayWithGenerator() {
     return ImmutableSet.of(1).asList().toArray(Integer[]::new);
   }
 
-  Iterator<Integer> testImmutableCollectionIterator() {
-    return ImmutableSet.of(1).asList().iterator();
+  ImmutableSet<Iterator<Integer>> testImmutableCollectionIterator() {
+    return ImmutableSet.of(
+        ImmutableSet.of(1).stream().iterator(), ImmutableSet.of(1).asList().iterator());
   }
 
   ImmutableSet<Optional<Integer>> testOptionalFirstCollectionElement() {
diff --git a/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/CollectionRulesTestOutput.java b/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/CollectionRulesTestOutput.java
index 756135c8890..359feab5083 100644
--- a/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/CollectionRulesTestOutput.java
+++ b/error-prone-contrib/src/test/resources/tech/picnic/errorprone/refasterrules/CollectionRulesTestOutput.java
@@ -109,8 +109,8 @@ Integer[] testImmutableCollectionToArrayWithGenerator() {
     return ImmutableSet.of(1).toArray(Integer[]::new);
   }
 
-  Iterator<Integer> testImmutableCollectionIterator() {
-    return ImmutableSet.of(1).iterator();
+  ImmutableSet<Iterator<Integer>> testImmutableCollectionIterator() {
+    return ImmutableSet.of(ImmutableSet.of(1).iterator(), ImmutableSet.of(1).iterator());
   }
 
   ImmutableSet<Optional<Integer>> testOptionalFirstCollectionElement() {