Ensure representative set for many-to-one mappings

Change-Id: I692245c6fb54f791c205782a9499b920e1df89b8
diff --git a/src/main/java/com/android/tools/r8/horizontalclassmerging/HorizontalClassMergerGraphLens.java b/src/main/java/com/android/tools/r8/horizontalclassmerging/HorizontalClassMergerGraphLens.java
index 12a4190..abadbc4 100644
--- a/src/main/java/com/android/tools/r8/horizontalclassmerging/HorizontalClassMergerGraphLens.java
+++ b/src/main/java/com/android/tools/r8/horizontalclassmerging/HorizontalClassMergerGraphLens.java
@@ -9,6 +9,7 @@
 import com.android.tools.r8.graph.DexMethod;
 import com.android.tools.r8.graph.DexType;
 import com.android.tools.r8.graph.NestedGraphLens;
+import com.android.tools.r8.graph.ProgramMethod;
 import com.android.tools.r8.ir.conversion.ExtraParameter;
 import com.android.tools.r8.utils.IterableUtils;
 import com.android.tools.r8.utils.collections.BidirectionalManyToOneHashMap;
@@ -103,6 +104,13 @@
         AppView<?> appView, HorizontallyMergedClasses mergedClasses) {
       assert pendingMethodMapUpdates.isEmpty();
       assert pendingNewMethodSignatureUpdates.isEmpty();
+      assert newMethodSignatures.values().stream()
+          .allMatch(
+              value -> {
+                assert newMethodSignatures.getKeys(value).size() == 1
+                    || newMethodSignatures.hasExplicitRepresentativeKey(value);
+                return true;
+              });
       return new HorizontalClassMergerGraphLens(
           appView,
           mergedClasses,
@@ -157,6 +165,17 @@
       recordNewMethodSignature(from, to, isRepresentative);
     }
 
+    void moveMethods(Iterable<ProgramMethod> methods, DexMethod to) {
+      moveMethods(methods, to, null);
+    }
+
+    void moveMethods(Iterable<ProgramMethod> methods, DexMethod to, ProgramMethod representative) {
+      for (ProgramMethod from : methods) {
+        boolean isRepresentative = representative != null && from == representative;
+        moveMethod(from.getReference(), to, isRepresentative);
+      }
+    }
+
     void recordNewMethodSignature(DexMethod oldMethodSignature, DexMethod newMethodSignature) {
       recordNewMethodSignature(oldMethodSignature, newMethodSignature, false);
     }
@@ -194,6 +213,10 @@
         for (DexMethod originalMethodSignature : oldMethodSignatures) {
           pendingNewMethodSignatureUpdates.put(originalMethodSignature, newMethodSignature);
         }
+        DexMethod representative = newMethodSignatures.getRepresentativeKey(oldMethodSignature);
+        if (representative != null) {
+          pendingNewMethodSignatureUpdates.setRepresentative(newMethodSignature, representative);
+        }
       }
     }
 
@@ -205,7 +228,13 @@
 
       // Commit pending original method signatures updates.
       newMethodSignatures.removeAll(pendingNewMethodSignatureUpdates.keySet());
-      pendingNewMethodSignatureUpdates.forEachManyToOneMapping(newMethodSignatures::put);
+      pendingNewMethodSignatureUpdates.forEachManyToOneMapping(
+          (keys, value, representative) -> {
+            newMethodSignatures.put(keys, value);
+            if (keys.size() > 1) {
+              newMethodSignatures.setRepresentative(value, representative);
+            }
+          });
       pendingNewMethodSignatureUpdates.clear();
     }
 
diff --git a/src/main/java/com/android/tools/r8/horizontalclassmerging/InstanceInitializerMerger.java b/src/main/java/com/android/tools/r8/horizontalclassmerging/InstanceInitializerMerger.java
index c2063f3..b50c71f 100644
--- a/src/main/java/com/android/tools/r8/horizontalclassmerging/InstanceInitializerMerger.java
+++ b/src/main/java/com/android/tools/r8/horizontalclassmerging/InstanceInitializerMerger.java
@@ -325,11 +325,7 @@
     ProgramMethod representative = ListUtils.first(instanceInitializers);
     DexMethod newMethodReference =
         representative.getReference().withHolder(group.getTarget(), dexItemFactory);
-
-    for (ProgramMethod constructor : instanceInitializers) {
-      boolean isRepresentative = constructor == representative;
-      lensBuilder.moveMethod(constructor.getReference(), newMethodReference, isRepresentative);
-    }
+    lensBuilder.moveMethods(instanceInitializers, newMethodReference, representative);
 
     DexEncodedMethod newMethod =
         representative.getHolder() == group.getTarget()
diff --git a/src/main/java/com/android/tools/r8/synthesis/SyntheticFinalization.java b/src/main/java/com/android/tools/r8/synthesis/SyntheticFinalization.java
index 89552fb..8c37bf1 100644
--- a/src/main/java/com/android/tools/r8/synthesis/SyntheticFinalization.java
+++ b/src/main/java/com/android/tools/r8/synthesis/SyntheticFinalization.java
@@ -125,6 +125,14 @@
       methodMap.put(from, to);
     }
 
+    void setRepresentative(DexField field, DexField representative) {
+      fieldMap.setRepresentative(field, representative);
+    }
+
+    void setRepresentative(DexMethod method, DexMethod representative) {
+      methodMap.setRepresentative(method, representative);
+    }
+
     SyntheticFinalizationGraphLens build(AppView<?> appView) {
       if (typeMap.isEmpty() && fieldMap.isEmpty() && methodMap.isEmpty()) {
         return null;
@@ -477,6 +485,26 @@
                   syntheticMethodDefinition.getReference()));
         });
 
+    Iterables.<EquivalenceGroup<? extends SyntheticDefinition<?, ?, DexProgramClass>>>concat(
+            syntheticClassGroups.values(), syntheticMethodGroups.values())
+        .forEach(
+            syntheticGroup ->
+                syntheticGroup
+                    .getRepresentative()
+                    .getHolder()
+                    .forEachProgramMember(
+                        member -> {
+                          if (member.isProgramField()) {
+                            DexField field = member.asProgramField().getReference();
+                            DexField rewrittenField = treeFixer.fixupFieldReference(field);
+                            lensBuilder.setRepresentative(rewrittenField, field);
+                          } else {
+                            DexMethod method = member.asProgramMethod().getReference();
+                            DexMethod rewrittenMethod = treeFixer.fixupMethodReference(method);
+                            lensBuilder.setRepresentative(rewrittenMethod, method);
+                          }
+                        }));
+
     for (DexType key : syntheticMethodGroups.keySet()) {
       assert application.definitionFor(key) != null;
     }
diff --git a/src/main/java/com/android/tools/r8/utils/collections/BidirectionalManyToManyRepresentativeMap.java b/src/main/java/com/android/tools/r8/utils/collections/BidirectionalManyToManyRepresentativeMap.java
index 773c066..37b9d40 100644
--- a/src/main/java/com/android/tools/r8/utils/collections/BidirectionalManyToManyRepresentativeMap.java
+++ b/src/main/java/com/android/tools/r8/utils/collections/BidirectionalManyToManyRepresentativeMap.java
@@ -14,6 +14,8 @@
 public interface BidirectionalManyToManyRepresentativeMap<K, V>
     extends BidirectionalManyToManyMap<K, V> {
 
+  boolean hasExplicitRepresentativeKey(V value);
+
   K getRepresentativeKey(V value);
 
   default K getRepresentativeKeyOrDefault(V value, K defaultValue) {
diff --git a/src/main/java/com/android/tools/r8/utils/collections/BidirectionalManyToOneRepresentativeHashMap.java b/src/main/java/com/android/tools/r8/utils/collections/BidirectionalManyToOneRepresentativeHashMap.java
index ffd653e..bdd70c8 100644
--- a/src/main/java/com/android/tools/r8/utils/collections/BidirectionalManyToOneRepresentativeHashMap.java
+++ b/src/main/java/com/android/tools/r8/utils/collections/BidirectionalManyToOneRepresentativeHashMap.java
@@ -4,6 +4,7 @@
 
 package com.android.tools.r8.utils.collections;
 
+import com.android.tools.r8.utils.TriConsumer;
 import java.util.Collections;
 import java.util.IdentityHashMap;
 import java.util.Map;
@@ -33,6 +34,12 @@
   }
 
   @Override
+  public void forEachManyToOneMapping(TriConsumer<? super Set<K>, V, K> consumer) {
+    forEachManyToOneMapping(
+        (keys, value) -> consumer.accept(keys, value, getRepresentativeKey(value)));
+  }
+
+  @Override
   public K removeRepresentativeFor(V value) {
     return representatives.remove(value);
   }
@@ -43,10 +50,19 @@
   }
 
   @Override
+  public boolean hasExplicitRepresentativeKey(V value) {
+    return representatives.containsKey(value);
+  }
+
+  @Override
   public K getRepresentativeKey(V value) {
     Set<K> keys = getKeys(value);
     if (!keys.isEmpty()) {
-      return keys.size() == 1 ? keys.iterator().next() : representatives.get(value);
+      if (keys.size() == 1) {
+        return keys.iterator().next();
+      }
+      assert hasExplicitRepresentativeKey(value);
+      return representatives.get(value);
     }
     return null;
   }
@@ -67,8 +83,10 @@
   @Override
   public V remove(K key) {
     V value = super.remove(key);
-    if (getKeys(value).size() <= 1 || getRepresentativeKey(value) == key) {
-      removeRepresentativeFor(value);
+    if (hasExplicitRepresentativeKey(value)) {
+      if (getKeys(value).size() <= 1 || getRepresentativeKey(value) == key) {
+        removeRepresentativeFor(value);
+      }
     }
     return value;
   }
diff --git a/src/main/java/com/android/tools/r8/utils/collections/BidirectionalManyToOneRepresentativeMap.java b/src/main/java/com/android/tools/r8/utils/collections/BidirectionalManyToOneRepresentativeMap.java
index 3e7bf7b..9202cc1 100644
--- a/src/main/java/com/android/tools/r8/utils/collections/BidirectionalManyToOneRepresentativeMap.java
+++ b/src/main/java/com/android/tools/r8/utils/collections/BidirectionalManyToOneRepresentativeMap.java
@@ -4,6 +4,9 @@
 
 package com.android.tools.r8.utils.collections;
 
+import com.android.tools.r8.utils.TriConsumer;
+import java.util.Set;
+
 /**
  * Interface that accommodates many-to-one mappings.
  *
@@ -11,4 +14,7 @@
  * from {@link BidirectionalManyToManyRepresentativeMap}.
  */
 public interface BidirectionalManyToOneRepresentativeMap<K, V>
-    extends BidirectionalManyToOneMap<K, V>, BidirectionalManyToManyRepresentativeMap<K, V> {}
+    extends BidirectionalManyToOneMap<K, V>, BidirectionalManyToManyRepresentativeMap<K, V> {
+
+  void forEachManyToOneMapping(TriConsumer<? super Set<K>, V, K> consumer);
+}
diff --git a/src/main/java/com/android/tools/r8/utils/collections/BidirectionalOneToManyRepresentativeHashMap.java b/src/main/java/com/android/tools/r8/utils/collections/BidirectionalOneToManyRepresentativeHashMap.java
index 93562e7..a4974d8 100644
--- a/src/main/java/com/android/tools/r8/utils/collections/BidirectionalOneToManyRepresentativeHashMap.java
+++ b/src/main/java/com/android/tools/r8/utils/collections/BidirectionalOneToManyRepresentativeHashMap.java
@@ -22,6 +22,12 @@
   }
 
   @Override
+  public boolean hasExplicitRepresentativeKey(V value) {
+    assert containsValue(value);
+    return true;
+  }
+
+  @Override
   public K getRepresentativeKey(V value) {
     return getKey(value);
   }
diff --git a/src/main/java/com/android/tools/r8/utils/collections/BidirectionalOneToOneHashMap.java b/src/main/java/com/android/tools/r8/utils/collections/BidirectionalOneToOneHashMap.java
index e23a565..2961111 100644
--- a/src/main/java/com/android/tools/r8/utils/collections/BidirectionalOneToOneHashMap.java
+++ b/src/main/java/com/android/tools/r8/utils/collections/BidirectionalOneToOneHashMap.java
@@ -4,6 +4,7 @@
 
 package com.android.tools.r8.utils.collections;
 
+import com.android.tools.r8.utils.TriConsumer;
 import com.google.common.collect.BiMap;
 import com.google.common.collect.HashBiMap;
 import java.util.Collections;
@@ -61,6 +62,11 @@
   }
 
   @Override
+  public void forEachManyToOneMapping(TriConsumer<? super Set<K>, V, K> consumer) {
+    backing.forEach((key, value) -> consumer.accept(Collections.singleton(key), value, key));
+  }
+
+  @Override
   public void forEachValue(Consumer<? super V> consumer) {
     backing.values().forEach(consumer);
   }
@@ -92,6 +98,12 @@
   }
 
   @Override
+  public boolean hasExplicitRepresentativeKey(V value) {
+    assert containsValue(value);
+    return true;
+  }
+
+  @Override
   public K getRepresentativeKey(V value) {
     return getKey(value);
   }
diff --git a/src/main/java/com/android/tools/r8/utils/collections/EmptyBidirectionalOneToOneMap.java b/src/main/java/com/android/tools/r8/utils/collections/EmptyBidirectionalOneToOneMap.java
index 62fba9a..5f81584 100644
--- a/src/main/java/com/android/tools/r8/utils/collections/EmptyBidirectionalOneToOneMap.java
+++ b/src/main/java/com/android/tools/r8/utils/collections/EmptyBidirectionalOneToOneMap.java
@@ -4,6 +4,7 @@
 
 package com.android.tools.r8.utils.collections;
 
+import com.android.tools.r8.utils.TriConsumer;
 import com.google.common.collect.BiMap;
 import com.google.common.collect.HashBiMap;
 import java.util.Collections;
@@ -42,6 +43,11 @@
   }
 
   @Override
+  public void forEachManyToOneMapping(TriConsumer<? super Set<K>, V, K> consumer) {
+    // Intentionally empty.
+  }
+
+  @Override
   public void forEachValue(Consumer<? super V> consumer) {
     // Intentionally empty.
   }
@@ -67,6 +73,11 @@
   }
 
   @Override
+  public boolean hasExplicitRepresentativeKey(V value) {
+    return false;
+  }
+
+  @Override
   public K getRepresentativeKey(V value) {
     return null;
   }