Add equals() to RewrittenPrototypeDescription

Change-Id: I5431504aa99be87b504e7f090d00bad8fe9b06b3
diff --git a/src/main/java/com/android/tools/r8/graph/RewrittenPrototypeDescription.java b/src/main/java/com/android/tools/r8/graph/RewrittenPrototypeDescription.java
index 3c1f6e0..a4a886f 100644
--- a/src/main/java/com/android/tools/r8/graph/RewrittenPrototypeDescription.java
+++ b/src/main/java/com/android/tools/r8/graph/RewrittenPrototypeDescription.java
@@ -18,9 +18,9 @@
 import com.android.tools.r8.utils.IteratorUtils;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Ordering;
-import it.unimi.dsi.fastutil.ints.Int2ReferenceMap.Entry;
-import it.unimi.dsi.fastutil.ints.Int2ReferenceRBTreeMap;
-import it.unimi.dsi.fastutil.ints.Int2ReferenceSortedMap;
+import it.unimi.dsi.fastutil.ints.Int2ObjectMap.Entry;
+import it.unimi.dsi.fastutil.ints.Int2ObjectRBTreeMap;
+import it.unimi.dsi.fastutil.ints.Int2ObjectSortedMap;
 import it.unimi.dsi.fastutil.ints.IntArrayList;
 import it.unimi.dsi.fastutil.ints.IntBidirectionalIterator;
 import it.unimi.dsi.fastutil.ints.IntCollection;
@@ -31,14 +31,35 @@
 import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Objects;
 import java.util.function.Consumer;
 
 public class RewrittenPrototypeDescription {
 
-  public interface ArgumentInfo {
+  public abstract static class ArgumentInfo {
+
+    static final ArgumentInfo NO_INFO =
+        new ArgumentInfo() {
+
+          @Override
+          public ArgumentInfo combine(ArgumentInfo info) {
+            assert false : "ArgumentInfo NO_INFO should not be combined";
+            return info;
+          }
+
+          @Override
+          public boolean equals(Object obj) {
+            return obj == this;
+          }
+
+          @Override
+          public int hashCode() {
+            return System.identityHashCode(this);
+          }
+        };
 
     @SuppressWarnings("ConstantConditions")
-    static ArgumentInfo combine(ArgumentInfo arg1, ArgumentInfo arg2) {
+    public static ArgumentInfo combine(ArgumentInfo arg1, ArgumentInfo arg2) {
       if (arg1 == null) {
         assert arg2 != null;
         return arg2;
@@ -50,33 +71,33 @@
       return arg1.combine(arg2);
     }
 
-    ArgumentInfo NO_INFO =
-        info -> {
-          assert false : "ArgumentInfo NO_INFO should not be combined";
-          return info;
-        };
-
-    default boolean isRemovedArgumentInfo() {
+    public boolean isRemovedArgumentInfo() {
       return false;
     }
 
-    default RemovedArgumentInfo asRemovedArgumentInfo() {
+    public RemovedArgumentInfo asRemovedArgumentInfo() {
       return null;
     }
 
-    default boolean isRewrittenTypeInfo() {
+    public boolean isRewrittenTypeInfo() {
       return false;
     }
 
-    default RewrittenTypeInfo asRewrittenTypeInfo() {
+    public RewrittenTypeInfo asRewrittenTypeInfo() {
       return null;
     }
 
     // ArgumentInfo are combined with `this` first, and the `info` argument second.
-    ArgumentInfo combine(ArgumentInfo info);
+    public abstract ArgumentInfo combine(ArgumentInfo info);
+
+    @Override
+    public abstract boolean equals(Object obj);
+
+    @Override
+    public abstract int hashCode();
   }
 
-  public static class RemovedArgumentInfo implements ArgumentInfo {
+  public static class RemovedArgumentInfo extends ArgumentInfo {
 
     public static class Builder {
 
@@ -142,9 +163,23 @@
       assert false : "Once the argument is removed one cannot modify it any further.";
       return this;
     }
+
+    @Override
+    public boolean equals(Object obj) {
+      if (obj == null || getClass() != obj.getClass()) {
+        return false;
+      }
+      RemovedArgumentInfo other = (RemovedArgumentInfo) obj;
+      return type == other.type && Objects.equals(singleValue, other.singleValue);
+    }
+
+    @Override
+    public int hashCode() {
+      return Objects.hash(singleValue, type);
+    }
   }
 
-  public static class RewrittenTypeInfo implements ArgumentInfo {
+  public static class RewrittenTypeInfo extends ArgumentInfo {
 
     private final DexType oldType;
     private final DexType newType;
@@ -201,6 +236,20 @@
       return new RewrittenTypeInfo(oldType, rewrittenTypeInfo.newType);
     }
 
+    @Override
+    public boolean equals(Object obj) {
+      if (obj == null || getClass() != obj.getClass()) {
+        return false;
+      }
+      RewrittenTypeInfo other = (RewrittenTypeInfo) obj;
+      return oldType == other.oldType && newType == other.newType;
+    }
+
+    @Override
+    public int hashCode() {
+      return Objects.hash(oldType, newType);
+    }
+
     public boolean verifyIsDueToUnboxing(DexItemFactory dexItemFactory) {
       assert oldType.toBaseType(dexItemFactory).isClassType();
       assert newType.toBaseType(dexItemFactory).isIntType();
@@ -212,14 +261,14 @@
 
     private static final ArgumentInfoCollection EMPTY = new ArgumentInfoCollection();
 
-    private final Int2ReferenceSortedMap<ArgumentInfo> argumentInfos;
+    private final Int2ObjectSortedMap<ArgumentInfo> argumentInfos;
 
     // Specific constructor for empty.
     private ArgumentInfoCollection() {
-      this.argumentInfos = new Int2ReferenceRBTreeMap<>();
+      this.argumentInfos = new Int2ObjectRBTreeMap<>();
     }
 
-    private ArgumentInfoCollection(Int2ReferenceSortedMap<ArgumentInfo> argumentInfos) {
+    private ArgumentInfoCollection(Int2ObjectSortedMap<ArgumentInfo> argumentInfos) {
       assert argumentInfos != null : "should use empty.";
       assert !argumentInfos.isEmpty() : "should use empty.";
       this.argumentInfos = argumentInfos;
@@ -230,7 +279,7 @@
     }
 
     public void forEach(IntObjConsumer<ArgumentInfo> consumer) {
-      for (Entry<ArgumentInfo> entry : argumentInfos.int2ReferenceEntrySet()) {
+      for (Entry<ArgumentInfo> entry : argumentInfos.int2ObjectEntrySet()) {
         consumer.accept(entry.getIntKey(), entry.getValue());
       }
     }
@@ -267,7 +316,7 @@
     }
 
     public Iterator<Entry<ArgumentInfo>> iterator() {
-      return argumentInfos.int2ReferenceEntrySet().iterator();
+      return argumentInfos.int2ObjectEntrySet().iterator();
     }
 
     public boolean hasRemovedArguments() {
@@ -301,17 +350,31 @@
       return argumentInfos.size();
     }
 
+    @Override
+    public boolean equals(Object obj) {
+      if (obj == null || getClass() != obj.getClass()) {
+        return false;
+      }
+      ArgumentInfoCollection other = (ArgumentInfoCollection) obj;
+      return argumentInfos.equals(other.argumentInfos);
+    }
+
+    @Override
+    public int hashCode() {
+      return argumentInfos.hashCode();
+    }
+
     public static Builder builder() {
       return new Builder();
     }
 
     public static class Builder {
 
-      private Int2ReferenceSortedMap<ArgumentInfo> argumentInfos;
+      private Int2ObjectSortedMap<ArgumentInfo> argumentInfos;
 
       public Builder addArgumentInfo(int argIndex, ArgumentInfo argInfo) {
         if (argumentInfos == null) {
-          argumentInfos = new Int2ReferenceRBTreeMap<>();
+          argumentInfos = new Int2ObjectRBTreeMap<>();
         }
         assert !argumentInfos.containsKey(argIndex);
         argumentInfos.put(argIndex, argInfo);
@@ -377,7 +440,7 @@
         }
       }
 
-      Int2ReferenceSortedMap<ArgumentInfo> newArgInfos = new Int2ReferenceRBTreeMap<>();
+      Int2ObjectSortedMap<ArgumentInfo> newArgInfos = new Int2ObjectRBTreeMap<>();
       newArgInfos.putAll(argumentInfos);
       IntBidirectionalIterator iterator = argumentInfos.keySet().iterator();
       int offset = 0;
@@ -666,4 +729,20 @@
     return new RewrittenPrototypeDescription(
         newExtraParameters, rewrittenReturnInfo, argumentInfoCollection);
   }
+
+  @Override
+  public boolean equals(Object obj) {
+    if (obj == null || getClass() != obj.getClass()) {
+      return false;
+    }
+    RewrittenPrototypeDescription other = (RewrittenPrototypeDescription) obj;
+    return extraParameters.equals(other.extraParameters)
+        && Objects.equals(rewrittenReturnInfo, other.rewrittenReturnInfo)
+        && argumentInfoCollection.equals(other.argumentInfoCollection);
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hash(extraParameters, rewrittenReturnInfo, argumentInfoCollection);
+  }
 }
diff --git a/src/main/java/com/android/tools/r8/ir/conversion/ExtraConstantIntParameter.java b/src/main/java/com/android/tools/r8/ir/conversion/ExtraConstantIntParameter.java
index d698e7a..b0544d9 100644
--- a/src/main/java/com/android/tools/r8/ir/conversion/ExtraConstantIntParameter.java
+++ b/src/main/java/com/android/tools/r8/ir/conversion/ExtraConstantIntParameter.java
@@ -10,7 +10,8 @@
 import com.android.tools.r8.ir.analysis.value.SingleNumberValue;
 
 public class ExtraConstantIntParameter extends ExtraParameter {
-  long value;
+
+  private final long value;
 
   public ExtraConstantIntParameter(long value) {
     this.value = value;
@@ -26,4 +27,18 @@
   public SingleNumberValue getValue(AppView<?> appView) {
     return appView.abstractValueFactory().createSingleNumberValue(value);
   }
+
+  @Override
+  public boolean equals(Object obj) {
+    if (obj == null || getClass() != obj.getClass()) {
+      return false;
+    }
+    ExtraConstantIntParameter other = (ExtraConstantIntParameter) obj;
+    return value == other.value;
+  }
+
+  @Override
+  public int hashCode() {
+    return Long.hashCode(value);
+  }
 }
diff --git a/src/main/java/com/android/tools/r8/ir/conversion/ExtraParameter.java b/src/main/java/com/android/tools/r8/ir/conversion/ExtraParameter.java
index 93f52df..6211fa3 100644
--- a/src/main/java/com/android/tools/r8/ir/conversion/ExtraParameter.java
+++ b/src/main/java/com/android/tools/r8/ir/conversion/ExtraParameter.java
@@ -10,7 +10,14 @@
 import com.android.tools.r8.ir.analysis.value.SingleNumberValue;
 
 public abstract class ExtraParameter {
+
   public abstract TypeElement getTypeElement(AppView<?> appView, DexType argType);
 
   public abstract SingleNumberValue getValue(AppView<?> appView);
+
+  @Override
+  public abstract boolean equals(Object obj);
+
+  @Override
+  public abstract int hashCode();
 }
diff --git a/src/main/java/com/android/tools/r8/ir/conversion/ExtraUnusedNullParameter.java b/src/main/java/com/android/tools/r8/ir/conversion/ExtraUnusedNullParameter.java
index f478661..a543a0c 100644
--- a/src/main/java/com/android/tools/r8/ir/conversion/ExtraUnusedNullParameter.java
+++ b/src/main/java/com/android/tools/r8/ir/conversion/ExtraUnusedNullParameter.java
@@ -21,4 +21,14 @@
   public SingleNumberValue getValue(AppView<?> appView) {
     return appView.abstractValueFactory().createNullValue();
   }
+
+  @Override
+  public boolean equals(Object obj) {
+    return obj != null && getClass() == obj.getClass();
+  }
+
+  @Override
+  public int hashCode() {
+    return 0;
+  }
 }
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagatorProgramOptimizer.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagatorProgramOptimizer.java
index 483f6c4..055013f 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagatorProgramOptimizer.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagatorProgramOptimizer.java
@@ -331,13 +331,13 @@
 
     private DexMethod getNewMethodSignature(
         ProgramMethod method, RewrittenPrototypeDescription prototypeChanges) {
-      DexMethodSignature methodSignatureWithoutParametersRemoved = method.getMethodSignature();
+      DexMethodSignature methodSignatureWithoutPrototypeChanges = method.getMethodSignature();
       IntSet removableParameterIndices = prototypeChanges.getArgumentInfoCollection().getKeys();
 
       // Check if there is a reserved signature for this already.
       DexMethodSignature reservedSignature =
           newMethodSignatures
-              .getOrDefault(methodSignatureWithoutParametersRemoved, Collections.emptyMap())
+              .getOrDefault(methodSignatureWithoutPrototypeChanges, Collections.emptyMap())
               .get(removableParameterIndices);
       if (reservedSignature != null) {
         return reservedSignature.withHolder(method.getHolderType(), dexItemFactory);
@@ -353,7 +353,7 @@
         if (!method.getDefinition().isInstanceInitializer()) {
           reserveMethodSignature(
               methodSignatureWithParametersRemoved,
-              methodSignatureWithoutParametersRemoved,
+              methodSignatureWithoutPrototypeChanges,
               removableParameterIndices);
         }
         return methodReferenceWithParametersRemoved;
@@ -363,7 +363,7 @@
           occupiedMethodSignatures.get(methodSignatureWithParametersRemoved);
       // In this case we should have found a reserved method signature above.
       assert !(occupant.getFirst().equals(removableParameterIndices)
-          && occupant.getSecond().equals(methodSignatureWithoutParametersRemoved));
+          && occupant.getSecond().equals(methodSignatureWithoutPrototypeChanges));
 
       // We need to find a new name for this method, since the signature is already occupied.
       // TODO(b/190154391): Instead of generating a new name, we could also try permuting the order
@@ -380,16 +380,14 @@
                   return true;
                 }
                 return candidateOccupant.getFirst().equals(removableParameterIndices)
-                    && candidateOccupant
-                        .getSecond()
-                        .equals(methodSignatureWithoutParametersRemoved);
+                    && candidateOccupant.getSecond().equals(methodSignatureWithoutPrototypeChanges);
               });
 
       // Reserve the newly generated method signature.
       if (!method.getDefinition().isInstanceInitializer()) {
         reserveMethodSignature(
             newMethod.getSignature(),
-            methodSignatureWithoutParametersRemoved,
+            methodSignatureWithoutPrototypeChanges,
             removableParameterIndices);
       }