Make ImmutableMap Stream collect()ion nicer (#654)

This adds an entriesToImmutableMap() collector that can be used in place of
toImmutableMap(Map.Entry::getkey, Map.Entry::getValue()).

It also fixes up some existing calls that use toImmutableMap() when terser
alternatives exist.

buildSrc/src/main/java/google/registry/gradle/plugin/CoverPageGenerator.java

@@ -29,6 +29,7 @@ import google.registry.gradle.plugin.ProjectData.TaskData;
 import java.io.File;
 import java.nio.file.Path;
 import java.nio.file.Paths;
+import java.util.Map;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
 
@@ -162,7 +163,7 @@ final class CoverPageGenerator {
             task.reports().entrySet().stream()
                 .collect(
                     toImmutableMap(
-                        entry -> entry.getKey(),
+                        Map.Entry::getKey,
                         entry ->
                             entry.getValue().files().isEmpty()
                                 ? ""

buildSrc/src/main/java/google/registry/gradle/plugin/GcsPluginUtils.java

@@ -168,7 +168,7 @@ final class GcsPluginUtils {
             .collect(
                 toImmutableMap(
                     file -> rootDir.relativize(toNormalizedPath(file)),
-                    file -> toByteArraySupplier(file)));
+                    GcsPluginUtils::toByteArraySupplier));
 
     if (files.isEmpty()) {
       // The directory exists, but is empty. Return empty FilesWithEntryPoint

core/src/main/java/google/registry/flows/domain/token/AllocationTokenCustomLogic.java

@@ -16,12 +16,12 @@ package google.registry.flows.domain.token;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Maps;
 import com.google.common.net.InternetDomainName;
 import google.registry.flows.EppException;
 import google.registry.model.domain.DomainCommand;
 import google.registry.model.domain.token.AllocationToken;
 import google.registry.model.registry.Registry;
-import java.util.function.Function;
 import org.joda.time.DateTime;
 
 /**
@@ -50,7 +50,6 @@ public class AllocationTokenCustomLogic {
       String clientId,
       DateTime now) {
     // Do nothing.
-    return domainNames.stream()
+    return Maps.toMap(domainNames, k -> "");
-        .collect(ImmutableMap.toImmutableMap(Function.identity(), ignored -> ""));
   }
 }

core/src/main/java/google/registry/model/domain/rgp/GracePeriodStatus.java

@@ -86,10 +86,7 @@ public enum GracePeriodStatus implements EppEnum {
   /** Provide a quick lookup of GracePeriodStatus from XML name. */
   private static final ImmutableMap<String, GracePeriodStatus> XML_NAME_TO_GRACE_PERIOD_STATUS =
       Stream.of(GracePeriodStatus.values())
-          .collect(
+          .collect(toImmutableMap(GracePeriodStatus::getXmlName, value -> value));
-              toImmutableMap(
-                  (GracePeriodStatus gracePeriodStatus) -> gracePeriodStatus.xmlName,
-                  value -> value));
 
   @XmlAttribute(name = "s")
   private final String xmlName;

core/src/main/java/google/registry/model/registry/Registries.java

@@ -19,13 +19,13 @@ import static com.google.common.base.Predicates.equalTo;
 import static com.google.common.base.Predicates.in;
 import static com.google.common.base.Predicates.not;
 import static com.google.common.base.Strings.emptyToNull;
-import static com.google.common.collect.ImmutableMap.toImmutableMap;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static com.google.common.collect.Maps.filterValues;
 import static google.registry.model.CacheUtils.memoizeWithShortExpiration;
 import static google.registry.model.common.EntityGroupRoot.getCrossTldKey;
 import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.util.CollectionUtils.entriesToImmutableMap;
 import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
 
 import com.google.common.base.Joiner;
@@ -37,7 +37,6 @@ import com.google.common.collect.Streams;
 import com.google.common.net.InternetDomainName;
 import com.googlecode.objectify.Key;
 import google.registry.model.registry.Registry.TldType;
-import java.util.Map;
 import java.util.Optional;
 
 /** Utilities for finding and listing {@link Registry} entities. */
@@ -71,7 +70,7 @@ public final class Registries {
                               .collect(toImmutableSet());
                       return Registry.getAll(tlds).stream()
                           .map(e -> Maps.immutableEntry(e.getTldStr(), e.getTldType()))
-                          .collect(toImmutableMap(Map.Entry::getKey, Map.Entry::getValue));
+                          .collect(entriesToImmutableMap());
                     }));
   }
 

core/src/main/java/google/registry/model/registry/label/ReservedList.java

@@ -34,7 +34,6 @@ import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.MapDifference;
-import com.google.common.collect.MapDifference.ValueDifference;
 import com.google.common.collect.Maps;
 import com.google.common.flogger.FluentLogger;
 import com.google.common.util.concurrent.UncheckedExecutionException;
@@ -265,7 +264,7 @@ public final class ReservedList
           datastoreList.reservedListMap.entrySet().parallelStream()
               .collect(
                   toImmutableMap(
-                      entry -> entry.getKey(),
+                      Map.Entry::getKey,
                       entry ->
                           ReservedEntry.create(
                               entry.getValue().reservationType, entry.getValue().comment)));
@@ -283,17 +282,14 @@ public final class ReservedList
                   cloudSqlList.getRevisionId(), diff.entriesDiffering().size()));
         } else {
           StringBuilder diffMessage = new StringBuilder("Unequal reserved lists detected:\n");
-          diff.entriesDiffering().entrySet().stream()
+          diff.entriesDiffering()
               .forEach(
-                  entry -> {
+                  (label, valueDiff) ->
-                    String label = entry.getKey();
-                    ValueDifference<ReservedEntry> valueDiff = entry.getValue();
                       diffMessage.append(
                           String.format(
                               "Domain label %s has entry %s in Datastore and entry"
                                   + " %s in Cloud SQL.\n",
-                            label, valueDiff.leftValue(), valueDiff.rightValue()));
+                              label, valueDiff.leftValue(), valueDiff.rightValue())));
-                  });
           logger.atWarning().log(diffMessage.toString());
         }
       }

core/src/main/java/google/registry/persistence/converter/StringMapConverterBase.java

@@ -14,7 +14,7 @@
 
 package google.registry.persistence.converter;
 
-import static com.google.common.collect.ImmutableMap.toImmutableMap;
+import static google.registry.util.CollectionUtils.entriesToImmutableMap;
 
 import google.registry.persistence.converter.StringMapDescriptor.StringMap;
 import java.util.Map;
@@ -38,7 +38,7 @@ public abstract class StringMapConverterBase<K, V>
         : StringMap.create(
             attribute.entrySet().stream()
                 .map(this::convertToDatabaseMapEntry)
-                .collect(toImmutableMap(Map.Entry::getKey, Map.Entry::getValue)));
+                .collect(entriesToImmutableMap()));
   }
 
   @Override
@@ -47,6 +47,6 @@ public abstract class StringMapConverterBase<K, V>
         ? null
         : dbData.getMap().entrySet().stream()
             .map(this::convertToEntityMapEntry)
-            .collect(toImmutableMap(Map.Entry::getKey, Map.Entry::getValue));
+            .collect(entriesToImmutableMap());
   }
 }

core/src/main/java/google/registry/persistence/converter/TimedTransitionPropertyConverterBase.java

@@ -14,8 +14,9 @@
 
 package google.registry.persistence.converter;
 
-import static com.google.common.collect.ImmutableMap.toImmutableMap;
+import static google.registry.util.CollectionUtils.entriesToImmutableMap;
 
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSortedMap;
 import google.registry.model.common.TimedTransitionProperty;
 import google.registry.model.common.TimedTransitionProperty.TimedTransition;
@@ -45,7 +46,7 @@ public abstract class TimedTransitionPropertyConverterBase<K, V extends TimedTra
         : StringMap.create(
             attribute.entrySet().stream()
                 .map(this::convertToDatabaseMapEntry)
-                .collect(toImmutableMap(Map.Entry::getKey, Map.Entry::getValue)));
+                .collect(entriesToImmutableMap()));
   }
 
   @Override
@@ -53,10 +54,10 @@ public abstract class TimedTransitionPropertyConverterBase<K, V extends TimedTra
     if (dbData == null) {
       return null;
     }
-    Map<DateTime, K> map =
+    ImmutableMap<DateTime, K> map =
         dbData.getMap().entrySet().stream()
             .map(this::convertToEntityMapEntry)
-            .collect(toImmutableMap(Map.Entry::getKey, Map.Entry::getValue));
+            .collect(entriesToImmutableMap());
     return TimedTransitionProperty.fromValueMap(
         ImmutableSortedMap.copyOf(map), getTimedTransitionSubclass());
   }

core/src/main/java/google/registry/reporting/icann/IcannReportingUploadAction.java

@@ -15,7 +15,6 @@
 package google.registry.reporting.icann;
 
 import static com.google.common.base.Preconditions.checkArgument;
-import static com.google.common.collect.ImmutableMap.toImmutableMap;
 import static com.google.common.net.MediaType.PLAIN_TEXT_UTF_8;
 import static google.registry.model.common.Cursor.getCursorTimeOrStartOfTime;
 import static google.registry.model.ofy.ObjectifyService.ofy;
@@ -27,6 +26,7 @@ import static javax.servlet.http.HttpServletResponse.SC_OK;
 import com.google.appengine.tools.cloudstorage.GcsFilename;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Maps;
 import com.google.common.flogger.FluentLogger;
 import com.google.common.io.ByteStreams;
 import com.googlecode.objectify.Key;
@@ -209,9 +209,9 @@ public final class IcannReportingUploadAction implements Runnable {
 
     ImmutableSet<Registry> registries = Registries.getTldEntitiesOfType(TldType.REAL);
 
-    Map<Key<Cursor>, Registry> activityKeyMap =
+    ImmutableMap<Key<Cursor>, Registry> activityKeyMap =
         loadKeyMap(registries, CursorType.ICANN_UPLOAD_ACTIVITY);
-    Map<Key<Cursor>, Registry> transactionKeyMap =
+    ImmutableMap<Key<Cursor>, Registry> transactionKeyMap =
         loadKeyMap(registries, CursorType.ICANN_UPLOAD_TX);
 
     ImmutableSet.Builder<Key<Cursor>> keys = new ImmutableSet.Builder<>();
@@ -229,9 +229,9 @@ public final class IcannReportingUploadAction implements Runnable {
     return cursors.build();
   }
 
-  private Map<Key<Cursor>, Registry> loadKeyMap(
+  private ImmutableMap<Key<Cursor>, Registry> loadKeyMap(
       ImmutableSet<Registry> registries, CursorType type) {
-    return registries.stream().collect(toImmutableMap(r -> Cursor.createKey(type, r), r -> r));
+    return Maps.uniqueIndex(registries, r -> Cursor.createKey(type, r));
   }
 
   /**

core/src/main/java/google/registry/schema/cursor/CursorDao.java

@@ -15,7 +15,6 @@
 package google.registry.schema.cursor;
 
 import static com.google.appengine.api.search.checkers.Preconditions.checkNotNull;
-import static com.google.common.collect.ImmutableMap.toImmutableMap;
 import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
@@ -23,6 +22,7 @@ import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
 
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Maps;
 import com.google.common.flogger.FluentLogger;
 import google.registry.model.common.Cursor.CursorType;
 import google.registry.schema.cursor.Cursor.CursorId;
@@ -171,11 +171,11 @@ public class CursorDao {
       // Load all the cursors of that type from Cloud SQL
       List<Cursor> cloudSqlCursors = loadByType(type);
 
-      // Create a map of each tld to its cursor if one exists
+      // Create a map of each TLD to its cursor if one exists.
       ImmutableMap<String, Cursor> cloudSqlCursorMap =
-          cloudSqlCursors.stream().collect(toImmutableMap(c -> c.getScope(), c -> c));
+          Maps.uniqueIndex(cloudSqlCursors, Cursor::getScope);
 
-      // Compare each Datastore cursor with its corresponding Cloud SQL cursor
+      // Compare each Datastore cursor with its corresponding Cloud SQL cursor.
       for (google.registry.model.common.Cursor cursor : cursors.keySet()) {
         Cursor cloudSqlCursor = cloudSqlCursorMap.get(cursors.get(cursor));
         compare(cursor, cloudSqlCursor, cursors.get(cursor));

core/src/main/java/google/registry/tools/SetNumInstancesCommand.java

@@ -15,7 +15,6 @@
 package google.registry.tools;
 
 import static com.google.common.base.Preconditions.checkArgument;
-import static com.google.common.collect.ImmutableMap.toImmutableMap;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static com.google.common.collect.ImmutableSetMultimap.flatteningToImmutableSetMultimap;
 import static google.registry.util.CollectionUtils.nullToEmpty;
@@ -29,6 +28,7 @@ import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSetMultimap;
+import com.google.common.collect.Maps;
 import com.google.common.flogger.FluentLogger;
 import google.registry.config.RegistryConfig.Config;
 import google.registry.request.Action.Service;
@@ -53,8 +53,7 @@ final class SetNumInstancesCommand implements CommandWithRemoteApi {
       ImmutableSet.copyOf(Service.values());
 
   private static final ImmutableMap<String, Service> SERVICE_ID_TO_SERVICE =
-      ALL_DEPLOYED_SERVICES.stream()
+      Maps.uniqueIndex(ALL_DEPLOYED_SERVICES, Service::getServiceId);
-          .collect(toImmutableMap(service -> service.getServiceId(), service -> service));
 
   // TODO(b/119629679): Use List<Service> after upgrading jcommander to latest version.
   @Parameter(

core/src/test/java/google/registry/flows/domain/token/AllocationTokenFlowUtilsTest.java

@@ -34,6 +34,7 @@ import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSortedMap;
+import com.google.common.collect.Maps;
 import com.google.common.net.InternetDomainName;
 import com.googlecode.objectify.Key;
 import google.registry.flows.EppException;
@@ -48,7 +49,6 @@ import google.registry.model.registry.Registry;
 import google.registry.model.reporting.HistoryEntry;
 import google.registry.testing.AppEngineRule;
 import google.registry.testing.ShardableTestCase;
-import java.util.function.Function;
 import org.joda.time.DateTime;
 import org.junit.Before;
 import org.junit.Rule;
@@ -324,11 +324,7 @@ public class AllocationTokenFlowUtilsTest extends ShardableTestCase {
         AllocationToken tokenEntity,
         String clientId,
         DateTime now) {
-      return domainNames.stream()
+      return Maps.toMap(domainNames, domain -> domain.toString().contains("bunny") ? "fufu" : "");
-          .collect(
-              ImmutableMap.toImmutableMap(
-                  Function.identity(),
-                  domainName -> domainName.toString().contains("bunny") ? "fufu" : ""));
     }
   }
 }

core/src/test/java/google/registry/testing/AppEngineRuleTest.java

@@ -15,11 +15,11 @@
 package google.registry.testing;
 
 import static com.google.common.truth.Truth.assertWithMessage;
+import static google.registry.util.CollectionUtils.entriesToImmutableMap;
 import static java.nio.charset.StandardCharsets.UTF_8;
 import static org.junit.Assert.assertThrows;
 
 import com.google.common.base.Joiner;
-import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Multimap;
 import com.google.common.collect.MultimapBuilder;
 import com.google.common.collect.Multimaps;
@@ -136,7 +136,7 @@ public class AppEngineRuleTest {
       Map<String, Collection<Class<?>>> conflictingKinds =
           kindToEntityMultiMap.asMap().entrySet().stream()
               .filter(e -> e.getValue().size() > 1)
-              .collect(ImmutableMap.toImmutableMap(Map.Entry::getKey, Map.Entry::getValue));
+              .collect(entriesToImmutableMap());
       assertWithMessage(
               "Conflicting Ofy kinds found. Tests will break if they are registered with "
                   + " AppEngineRule in the same test executor.")

core/src/test/java/google/registry/tools/GenerateAllocationTokensCommandTest.java

@@ -34,6 +34,7 @@ import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSortedMap;
 import com.google.common.collect.Iterables;
+import com.google.common.collect.Maps;
 import com.google.common.io.Files;
 import com.googlecode.objectify.Key;
 import google.registry.model.domain.token.AllocationToken;
@@ -47,7 +48,6 @@ import google.registry.util.Retrier;
 import google.registry.util.StringGenerator.Alphabets;
 import java.io.File;
 import java.util.Collection;
-import java.util.function.Function;
 import javax.annotation.Nullable;
 import org.joda.time.DateTime;
 import org.junit.Before;
@@ -311,8 +311,7 @@ public class GenerateAllocationTokensCommandTest
     // Using ImmutableObject comparison here is tricky because the creation/updated timestamps are
     // neither easy nor valuable to test here.
     ImmutableMap<String, AllocationToken> actualTokens =
-        ofy().load().type(AllocationToken.class).list().stream()
+        Maps.uniqueIndex(ofy().load().type(AllocationToken.class), AllocationToken::getToken);
-            .collect(ImmutableMap.toImmutableMap(AllocationToken::getToken, Function.identity()));
     assertThat(actualTokens).hasSize(expectedTokens.length);
     for (AllocationToken expectedToken : expectedTokens) {
       AllocationToken match = actualTokens.get(expectedToken.getToken());

util/src/main/java/google/registry/util/CollectionUtils.java

@@ -33,6 +33,7 @@ import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.SortedMap;
+import java.util.stream.Collector;
 import javax.annotation.Nullable;
 
 /** Utility methods related to collections. */
@@ -148,4 +149,11 @@ public class CollectionUtils {
     }
     return shards.build();
   }
+
+  /**
+   * Returns a {@link Collector} that accumulates {@link Map.Entry}s into an {@code ImmutableMap}.
+   */
+  public static <K, V> Collector<Map.Entry<K, V>, ?, ImmutableMap<K, V>> entriesToImmutableMap() {
+    return ImmutableMap.toImmutableMap(Map.Entry::getKey, Map.Entry::getValue);
+  }
 }