Introduce IsIdentity Matcher and use it in ImmutableMapTemplates

The Matcher is used for `Maps#toMap` and `Maps#UniqueIndex`.

error-prone-contrib/src/main/java/tech/picnic/errorprone/refastertemplates/ImmutableMapTemplates.java

@@ -10,7 +10,9 @@ import com.google.errorprone.refaster.ImportPolicy;
 import com.google.errorprone.refaster.Refaster;
 import com.google.errorprone.refaster.annotation.AfterTemplate;
 import com.google.errorprone.refaster.annotation.BeforeTemplate;
+import com.google.errorprone.refaster.annotation.Matches;
 import com.google.errorprone.refaster.annotation.MayOptionallyUse;
+import com.google.errorprone.refaster.annotation.NotMatches;
 import com.google.errorprone.refaster.annotation.Placeholder;
 import com.google.errorprone.refaster.annotation.UseImportPolicy;
 import java.util.Collection;
@@ -20,6 +22,7 @@ import java.util.Map;
 import java.util.Set;
 import java.util.function.Function;
 import java.util.stream.Stream;
+import tech.picnic.errorprone.refaster.util.IsIdentity;
 
 /** Refaster templates related to expressions dealing with {@link ImmutableMap}s. */
 final class ImmutableMapTemplates {
@@ -94,32 +97,35 @@ final class ImmutableMapTemplates {
    * Prefer {@link Maps#toMap(Iterable, com.google.common.base.Function)} over more contrived
    * alternatives.
    */
-  // XXX: Drop the `Refaster.anyOf` if/when we decide to rewrite one to the other.
   static final class IterableToImmutableMap<K, V> {
     @BeforeTemplate
     ImmutableMap<K, V> before(
-        Iterator<K> iterable, Function<? super K, ? extends V> valueFunction) {
+        Iterator<K> iterable,
+        @NotMatches(IsIdentity.class) Function<? super K, ? extends V> valueFunction) {
       return Streams.stream(iterable)
           .collect(toImmutableMap(Refaster.anyOf(identity(), k -> k), valueFunction));
     }
 
     @BeforeTemplate
     ImmutableMap<K, V> before(
-        Iterable<K> iterable, Function<? super K, ? extends V> valueFunction) {
+        Iterable<K> iterable,
+        @NotMatches(IsIdentity.class) Function<? super K, ? extends V> valueFunction) {
       return Streams.stream(iterable)
           .collect(toImmutableMap(Refaster.anyOf(identity(), k -> k), valueFunction));
     }
 
     @BeforeTemplate
     ImmutableMap<K, V> before(
-        Collection<K> iterable, Function<? super K, ? extends V> valueFunction) {
+        Collection<K> iterable,
+        @NotMatches(IsIdentity.class) Function<? super K, ? extends V> valueFunction) {
       return iterable.stream()
           .collect(toImmutableMap(Refaster.anyOf(identity(), k -> k), valueFunction));
     }
 
     @BeforeTemplate
     ImmutableMap<K, V> before(
-        Set<K> iterable, com.google.common.base.Function<? super K, V> valueFunction) {
+        Set<K> iterable,
+        @NotMatches(IsIdentity.class) com.google.common.base.Function<? super K, V> valueFunction) {
       return ImmutableMap.copyOf(Maps.asMap(iterable, valueFunction));
     }
 
@@ -130,6 +136,49 @@ final class ImmutableMapTemplates {
     }
   }
 
+  /**
+   * Prefer {@link Maps#toMap(Iterable, com.google.common.base.Function)} over more contrived
+   * alternatives and rewrite {@link Function#identity()} to {@code k -> k}.
+   */
+  static final class IterableToImmutableMapIdentity<K, V> {
+    @BeforeTemplate
+    ImmutableMap<K, V> before(
+        Iterator<K> iterable,
+        @Matches(IsIdentity.class) Function<? super K, ? extends V> valueFunction) {
+      return Streams.stream(iterable)
+          .collect(toImmutableMap(Refaster.anyOf(identity(), k -> k), valueFunction));
+    }
+
+    @BeforeTemplate
+    ImmutableMap<K, V> before(
+        Iterable<K> iterable,
+        @Matches(IsIdentity.class) Function<? super K, ? extends V> valueFunction) {
+      return Streams.stream(iterable)
+          .collect(toImmutableMap(Refaster.anyOf(identity(), k -> k), valueFunction));
+    }
+
+    @BeforeTemplate
+    ImmutableMap<K, V> before(
+        Collection<K> iterable,
+        @Matches(IsIdentity.class) Function<? super K, ? extends V> valueFunction) {
+      return iterable.stream()
+          .collect(toImmutableMap(Refaster.anyOf(identity(), k -> k), valueFunction));
+    }
+
+    @BeforeTemplate
+    ImmutableMap<K, V> before(
+        Set<K> iterable,
+        @Matches(IsIdentity.class) com.google.common.base.Function<? super K, V> valueFunction) {
+      return ImmutableMap.copyOf(Maps.asMap(iterable, valueFunction));
+    }
+
+    @AfterTemplate
+    ImmutableMap<K, K> after(
+        Iterable<K> iterable, com.google.common.base.Function<? super K, V> valueFunction) {
+      return Maps.toMap(iterable, k -> k);
+    }
+  }
+
   /** Prefer {@link ImmutableMap#copyOf(Iterable)} over more contrived alternatives. */
   static final class EntryIterableToImmutableMap<K, V> {
     @BeforeTemplate
@@ -188,23 +237,27 @@ final class ImmutableMapTemplates {
    * Prefer {@link Maps#uniqueIndex(Iterable, com.google.common.base.Function)} over the
    * stream-based alternative.
    */
-  // XXX: Drop the `Refaster.anyOf` if/when we decide to rewrite one to the other.
   static final class IndexIterableToImmutableMap<K, V> {
     @BeforeTemplate
-    ImmutableMap<K, V> before(Iterator<V> iterable, Function<? super V, ? extends K> keyFunction) {
-      return Streams.stream(iterable)
-          .collect(toImmutableMap(keyFunction, Refaster.anyOf(identity(), v -> v)));
-    }
-
-    @BeforeTemplate
-    ImmutableMap<K, V> before(Iterable<V> iterable, Function<? super V, ? extends K> keyFunction) {
+    ImmutableMap<K, V> before(
+        Iterator<V> iterable,
+        @NotMatches(IsIdentity.class) Function<? super V, ? extends K> keyFunction) {
       return Streams.stream(iterable)
           .collect(toImmutableMap(keyFunction, Refaster.anyOf(identity(), v -> v)));
     }
 
     @BeforeTemplate
     ImmutableMap<K, V> before(
-        Collection<V> iterable, Function<? super V, ? extends K> keyFunction) {
+        Iterable<V> iterable,
+        @NotMatches(IsIdentity.class) Function<? super V, ? extends K> keyFunction) {
+      return Streams.stream(iterable)
+          .collect(toImmutableMap(keyFunction, Refaster.anyOf(identity(), v -> v)));
+    }
+
+    @BeforeTemplate
+    ImmutableMap<K, V> before(
+        Collection<V> iterable,
+        @NotMatches(IsIdentity.class) Function<? super V, ? extends K> keyFunction) {
       return iterable.stream()
           .collect(toImmutableMap(keyFunction, Refaster.anyOf(identity(), v -> v)));
     }
@@ -216,6 +269,42 @@ final class ImmutableMapTemplates {
     }
   }
 
+  /**
+   * Prefer {@link Maps#uniqueIndex(Iterable, com.google.common.base.Function)} over the
+   * stream-based alternatives and rewrite {@link Function#identity()} to {@code k -> k}.
+   */
+  static final class IndexIterableToImmutableMapIdentity<K, V> {
+    @BeforeTemplate
+    ImmutableMap<K, V> before(
+        Iterator<V> iterable,
+        @Matches(IsIdentity.class) Function<? super V, ? extends K> keyFunction) {
+      return Streams.stream(iterable)
+          .collect(toImmutableMap(keyFunction, Refaster.anyOf(identity(), v -> v)));
+    }
+
+    @BeforeTemplate
+    ImmutableMap<K, V> before(
+        Iterable<V> iterable,
+        @Matches(IsIdentity.class) Function<? super V, ? extends K> keyFunction) {
+      return Streams.stream(iterable)
+          .collect(toImmutableMap(keyFunction, Refaster.anyOf(identity(), v -> v)));
+    }
+
+    @BeforeTemplate
+    ImmutableMap<K, V> before(
+        Collection<V> iterable,
+        @Matches(IsIdentity.class) Function<? super V, ? extends K> keyFunction) {
+      return iterable.stream()
+          .collect(toImmutableMap(keyFunction, Refaster.anyOf(identity(), v -> v)));
+    }
+
+    @AfterTemplate
+    ImmutableMap<V, V> after(
+        Iterable<V> iterable, com.google.common.base.Function<? super V, K> keyFunction) {
+      return Maps.uniqueIndex(iterable, v -> v);
+    }
+  }
+
   /**
    * Prefer creating an immutable copy of the result of {@link Maps#transformValues(Map,
    * com.google.common.base.Function)} over more contrived alternatives.

error-prone-contrib/src/test/resources/tech/picnic/errorprone/bugpatterns/ImmutableMapTemplatesTestInput.java

@@ -40,6 +40,7 @@ final class ImmutableMapTemplatesTest implements RefasterTemplateTestCase {
             .collect(toImmutableMap(Map.Entry::getKey, Map.Entry::getValue)));
   }
 
+  @SuppressWarnings({"IterableToImmutableMapIdentity", "IndexIterableToImmutableMapIdentity"})
   ImmutableSet<ImmutableMap<Integer, Integer>> testIterableToImmutableMap() {
     return ImmutableSet.of(
         ImmutableList.of(1).stream().collect(toImmutableMap(identity(), n -> n * 2)),
@@ -47,7 +48,14 @@ final class ImmutableMapTemplatesTest implements RefasterTemplateTestCase {
             .collect(toImmutableMap(n -> n, Integer::valueOf)),
         Streams.stream(ImmutableList.of(3).iterator())
             .collect(toImmutableMap(identity(), n -> n.intValue())),
-        ImmutableMap.copyOf(Maps.asMap(ImmutableSet.of(4), Integer::valueOf)));
+        ImmutableMap.copyOf(Maps.asMap(ImmutableSet.of(4), Integer::valueOf)),
+        ImmutableList.of(1).stream().collect(toImmutableMap(identity(), identity())));
+  }
+
+  ImmutableSet<ImmutableMap<Integer, Integer>> testIterableToImmutableMapIdentity() {
+    return ImmutableSet.of(
+        ImmutableList.of(1).stream().collect(toImmutableMap(identity(), identity())),
+        Streams.stream(ImmutableList.of(2)::iterator).collect(toImmutableMap(n -> n, identity())));
   }
 
   ImmutableSet<ImmutableMap<String, Integer>> testEntryIterableToImmutableMap() {
@@ -78,6 +86,16 @@ final class ImmutableMapTemplatesTest implements RefasterTemplateTestCase {
             .collect(toImmutableMap(n -> n.intValue(), identity())));
   }
 
+  @SuppressWarnings("IterableToImmutableMapIdentity")
+  ImmutableSet<ImmutableMap<Integer, Integer>> testIndexIterableToImmutableMapIdentity() {
+    return ImmutableSet.of(
+        ImmutableList.of(1).stream().collect(toImmutableMap(identity(), identity())),
+        Streams.stream(ImmutableList.of(2)::iterator)
+            .collect(toImmutableMap(identity(), identity())),
+        Streams.stream(ImmutableList.of(3).iterator())
+            .collect(toImmutableMap(identity(), identity())));
+  }
+
   ImmutableSet<ImmutableMap<String, Integer>> testTransformMapValuesToImmutableMap() {
     return ImmutableSet.of(
         ImmutableMap.of("foo", 1L).entrySet().stream()

error-prone-contrib/src/test/resources/tech/picnic/errorprone/bugpatterns/ImmutableMapTemplatesTestOutput.java

@@ -37,12 +37,19 @@ final class ImmutableMapTemplatesTest implements RefasterTemplateTestCase {
         ImmutableMap.of(Map.entry("foo", 1).getKey(), Map.entry("foo", 1).getValue()));
   }
 
+  @SuppressWarnings({"IterableToImmutableMapIdentity", "IndexIterableToImmutableMapIdentity"})
   ImmutableSet<ImmutableMap<Integer, Integer>> testIterableToImmutableMap() {
     return ImmutableSet.of(
         Maps.toMap(ImmutableList.of(1), n -> n * 2),
         Maps.toMap(ImmutableList.of(2)::iterator, Integer::valueOf),
         Maps.toMap(ImmutableList.of(3).iterator(), n -> n.intValue()),
-        Maps.toMap(ImmutableSet.of(4), Integer::valueOf));
+        Maps.toMap(ImmutableSet.of(4), Integer::valueOf),
+        ImmutableList.of(1).stream().collect(toImmutableMap(identity(), identity())));
+  }
+
+  ImmutableSet<ImmutableMap<Integer, Integer>> testIterableToImmutableMapIdentity() {
+    return ImmutableSet.of(
+        Maps.toMap(ImmutableList.of(1), k -> k), Maps.toMap(ImmutableList.of(2)::iterator, n -> n));
   }
 
   ImmutableSet<ImmutableMap<String, Integer>> testEntryIterableToImmutableMap() {
@@ -65,6 +72,14 @@ final class ImmutableMapTemplatesTest implements RefasterTemplateTestCase {
         Maps.uniqueIndex(ImmutableList.of(3).iterator(), n -> n.intValue()));
   }
 
+  @SuppressWarnings("IterableToImmutableMapIdentity")
+  ImmutableSet<ImmutableMap<Integer, Integer>> testIndexIterableToImmutableMapIdentity() {
+    return ImmutableSet.of(
+        Maps.uniqueIndex(ImmutableList.of(1), v -> v),
+        Maps.uniqueIndex(ImmutableList.of(2)::iterator, v -> v),
+        Maps.uniqueIndex(ImmutableList.of(3).iterator(), v -> v));
+  }
+
   ImmutableSet<ImmutableMap<String, Integer>> testTransformMapValuesToImmutableMap() {
     return ImmutableSet.of(
         ImmutableMap.copyOf(

refaster-support/src/main/java/tech/picnic/errorprone/refaster/util/IsIdentity.java

@@ -0,0 +1,23 @@
+package tech.picnic.errorprone.refaster.util;
+
+import static com.google.errorprone.matchers.Matchers.staticMethod;
+
+import com.google.errorprone.VisitorState;
+import com.google.errorprone.matchers.Matcher;
+import com.sun.source.tree.ExpressionTree;
+
+/**
+ * A matcher for `identity` operations, for use with Refaster's {@code {Not,}Matches} annotation.
+ */
+public final class IsIdentity implements Matcher<ExpressionTree> {
+  private static final long serialVersionUID = 1L;
+  private static final Matcher<ExpressionTree> IDENTITY =
+      staticMethod()
+          .onClass((type, state) -> type.toString().contains("Function"))
+          .named("identity");
+
+  @Override
+  public boolean matches(ExpressionTree tree, VisitorState state) {
+    return IDENTITY.matches(tree, state);
+  }
+}

refaster-support/src/test/java/tech/picnic/errorprone/refaster/util/IsIdentityTest.java

@@ -0,0 +1,55 @@
+package tech.picnic.errorprone.refaster.util;
+
+import static com.google.errorprone.BugPattern.SeverityLevel.ERROR;
+
+import com.google.errorprone.BugPattern;
+import com.google.errorprone.CompilationTestHelper;
+import com.google.errorprone.VisitorState;
+import com.google.errorprone.bugpatterns.BugChecker;
+import com.google.errorprone.bugpatterns.BugChecker.MethodInvocationTreeMatcher;
+import com.google.errorprone.matchers.Description;
+import com.sun.source.tree.MethodInvocationTree;
+import org.junit.jupiter.api.Test;
+
+final class IsIdentityTest {
+  @Test
+  void matches() {
+    CompilationTestHelper.newInstance(TestChecker.class, getClass())
+        .addSourceLines(
+            "A.java",
+            "import static com.google.common.collect.ImmutableMap.toImmutableMap;",
+            "import static java.util.function.Function.identity;",
+            "import com.google.common.collect.ImmutableList;",
+            "import com.google.common.collect.ImmutableMap;",
+            "import com.google.common.collect.ImmutableSet;",
+            "import com.google.common.collect.Streams;",
+            "",
+            "class A {",
+            "  void positive() {",
+            "    // BUG: Diagnostic contains:",
+            "    java.util.function.Function.identity();",
+            "    // BUG: Diagnostic contains:",
+            "    com.google.common.base.Functions.identity();",
+            "  }",
+            "",
+            "  public static void identity() { }",
+            "",
+            "  static void negative() {",
+            "    identity();",
+            "    A.identity();",
+            "  }",
+            "}")
+        .doTest();
+  }
+
+  /** A {@link BugChecker} which simply delegates to {@link IsIdentity}. */
+  @BugPattern(name = "TestChecker", summary = "Flags identity method invocations", severity = ERROR)
+  public static final class TestChecker extends BugChecker implements MethodInvocationTreeMatcher {
+    private static final long serialVersionUID = 1L;
+
+    @Override
+    public Description matchMethodInvocation(MethodInvocationTree tree, VisitorState state) {
+      return new IsIdentity().matches(tree, state) ? describeMatch(tree) : Description.NO_MATCH;
+    }
+  }
+}