Use a signature bimap in ConcurrentMethodFixup

Change-Id: I50b57777cfaa5a7299c9cf2c0fc3be770fd9dbd8

src/main/java/com/android/tools/r8/graph/fixup/ConcurrentMethodFixup.java

diff --git a/src/main/java/com/android/tools/r8/graph/fixup/ConcurrentMethodFixup.java b/src/main/java/com/android/tools/r8/graph/fixup/ConcurrentMethodFixup.java
index 81d780f..ce17027 100644
--- a/src/main/java/com/android/tools/r8/graph/fixup/ConcurrentMethodFixup.java
+++ b/src/main/java/com/android/tools/r8/graph/fixup/ConcurrentMethodFixup.java
@@ -22,6 +22,7 @@
 import com.android.tools.r8.shaking.KeepMethodInfo;
 import com.android.tools.r8.utils.ThreadUtils;
 import com.android.tools.r8.utils.Timing;
+import com.android.tools.r8.utils.collections.DexMethodSignatureBiMap;
 import com.google.common.collect.BiMap;
 import com.google.common.collect.HashBiMap;
 import com.google.common.collect.Sets;
@@ -76,7 +77,8 @@
   private void processConnectedProgramComponents(Set<DexProgramClass> classes) {
     List<DexProgramClass> sorted = new ArrayList<>(classes);
     sorted.sort(Comparator.comparing(DexClass::getType));
-    BiMap<DexMethodSignature, DexMethodSignature> componentSignatures = HashBiMap.create();
+    DexMethodSignatureBiMap<DexMethodSignature> componentSignatures =
+        new DexMethodSignatureBiMap<>();
 
     // 1) Reserve all library overrides and pinned virtual methods.
     reserveComponentPinnedAndInterfaceMethodSignatures(sorted, componentSignatures);
@@ -104,7 +106,7 @@
   private void processClass(
       DexProgramClass clazz,
       Set<DexProgramClass> processedClasses,
-      BiMap<DexMethodSignature, DexMethodSignature> componentSignatures) {
+      DexMethodSignatureBiMap<DexMethodSignature> componentSignatures) {
     assert !clazz.isInterface();
     if (!processedClasses.add(clazz)) {
       return;
@@ -121,7 +123,7 @@
   private void processInterface(
       DexProgramClass clazz,
       Set<DexProgramClass> processedInterfaces,
-      BiMap<DexMethodSignature, DexMethodSignature> componentSignatures) {
+      DexMethodSignatureBiMap<DexMethodSignature> componentSignatures) {
     assert clazz.isInterface();
     if (!processedInterfaces.add(clazz)) {
       return;
@@ -147,7 +149,7 @@
   }
 
   private MethodNamingUtility createMethodNamingUtility(
-      BiMap<DexMethodSignature, DexMethodSignature> inheritedSignatures, DexProgramClass clazz) {
+      DexMethodSignatureBiMap<DexMethodSignature> componentSignatures, DexProgramClass clazz) {
     BiMap<DexMethod, DexMethod> localSignatures = HashBiMap.create();
     clazz.forEachProgramInstanceInitializer(
         method -> {
@@ -155,12 +157,12 @@
             localSignatures.put(method.getReference(), method.getReference());
           }
         });
-    return new MethodNamingUtility(appView.dexItemFactory(), inheritedSignatures, localSignatures);
+    return new MethodNamingUtility(appView.dexItemFactory(), componentSignatures, localSignatures);
   }
 
   private void reserveComponentPinnedAndInterfaceMethodSignatures(
       List<DexProgramClass> stronglyConnectedProgramClasses,
-      BiMap<DexMethodSignature, DexMethodSignature> componentSignatures) {
+      DexMethodSignatureBiMap<DexMethodSignature> componentSignatures) {
     Set<ClasspathOrLibraryClass> seenNonProgramClasses = Sets.newIdentityHashSet();
     for (DexProgramClass clazz : stronglyConnectedProgramClasses) {
       // If a private or static method is pinned, we need to reserve the mapping to avoid creating
@@ -178,13 +180,10 @@
               superclass != null
                   && !superclass.isProgramClass()
                   && seenNonProgramClasses.add(superclass.asClasspathOrLibraryClass()),
-          (supertype, superclass) -> {
-            for (DexMethodSignature vMethod :
-                nonProgramVirtualMethods.getOrComputeNonProgramMethods(
-                    superclass.asClasspathOrLibraryClass())) {
-              componentSignatures.put(vMethod, vMethod);
-            }
-          });
+          (supertype, superclass) ->
+              componentSignatures.putAllToIdentity(
+                  nonProgramVirtualMethods.getOrComputeNonProgramMethods(
+                      superclass.asClasspathOrLibraryClass())));
     }
   }
 }

src/main/java/com/android/tools/r8/graph/fixup/MethodNamingUtility.java

diff --git a/src/main/java/com/android/tools/r8/graph/fixup/MethodNamingUtility.java b/src/main/java/com/android/tools/r8/graph/fixup/MethodNamingUtility.java
index 977e01c..688853d 100644
--- a/src/main/java/com/android/tools/r8/graph/fixup/MethodNamingUtility.java
+++ b/src/main/java/com/android/tools/r8/graph/fixup/MethodNamingUtility.java
@@ -10,18 +10,19 @@
 import com.android.tools.r8.graph.DexMethodSignature;
 import com.android.tools.r8.graph.DexProto;
 import com.android.tools.r8.graph.DexType;
+import com.android.tools.r8.utils.collections.DexMethodSignatureBiMap;
 import com.google.common.collect.BiMap;
 import java.util.function.BiConsumer;
 
 public class MethodNamingUtility {
 
   private final DexItemFactory factory;
-  private final BiMap<DexMethodSignature, DexMethodSignature> inheritedSignatures;
+  private final DexMethodSignatureBiMap<DexMethodSignature> inheritedSignatures;
   private final BiMap<DexMethod, DexMethod> localSignatures;
 
   public MethodNamingUtility(
       DexItemFactory factory,
-      BiMap<DexMethodSignature, DexMethodSignature> inheritedSignatures,
+      DexMethodSignatureBiMap<DexMethodSignature> inheritedSignatures,
       BiMap<DexMethod, DexMethod> localSignatures) {
     this.factory = factory;
     this.inheritedSignatures = inheritedSignatures;

src/main/java/com/android/tools/r8/utils/collections/DexMethodSignatureBiMap.java

diff --git a/src/main/java/com/android/tools/r8/utils/collections/DexMethodSignatureBiMap.java b/src/main/java/com/android/tools/r8/utils/collections/DexMethodSignatureBiMap.java
new file mode 100644
index 0000000..64ce1ec
--- /dev/null
+++ b/src/main/java/com/android/tools/r8/utils/collections/DexMethodSignatureBiMap.java
@@ -0,0 +1,14 @@
+// Copyright (c) 2023, the R8 project authors. Please see the AUTHORS file
+// for details. All rights reserved. Use of this source code is governed by a
+// BSD-style license that can be found in the LICENSE file.
+
+package com.android.tools.r8.utils.collections;
+
+import com.google.common.collect.HashBiMap;
+
+public class DexMethodSignatureBiMap<T> extends DexMethodSignatureMap<T> {
+
+  public DexMethodSignatureBiMap() {
+    super(HashBiMap.create());
+  }
+}

src/main/java/com/android/tools/r8/utils/collections/DexMethodSignatureMap.java

diff --git a/src/main/java/com/android/tools/r8/utils/collections/DexMethodSignatureMap.java b/src/main/java/com/android/tools/r8/utils/collections/DexMethodSignatureMap.java
index 84d1cae..5289f92 100644
--- a/src/main/java/com/android/tools/r8/utils/collections/DexMethodSignatureMap.java
+++ b/src/main/java/com/android/tools/r8/utils/collections/DexMethodSignatureMap.java
@@ -20,7 +20,7 @@
 
   private final Map<DexMethodSignature, T> backing;
 
-  private DexMethodSignatureMap(Map<DexMethodSignature, T> backing) {
+  DexMethodSignatureMap(Map<DexMethodSignature, T> backing) {
     this.backing = backing;
   }
 
@@ -189,7 +189,16 @@
   }
 
   @Override
-  public void putAll(Map<? extends DexMethodSignature, ? extends T> m) {}
+  public void putAll(Map<? extends DexMethodSignature, ? extends T> map) {
+    map.forEach(this::put);
+  }
+
+  @SuppressWarnings("unchecked")
+  public void putAllToIdentity(Collection<? extends DexMethodSignature> signatures) {
+    for (DexMethodSignature signature : signatures) {
+      put(signature, (T) signature);
+    }
+  }
 
   public T remove(DexMethodSignature signature) {
     return backing.remove(signature);