Cleanup computation of new argument index after argument removal

Change-Id: I36d17cc7ef3ccced1f48501287d9178147c5ff7f
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 7d65dc9..dffe2a3 100644
--- a/src/main/java/com/android/tools/r8/graph/RewrittenPrototypeDescription.java
+++ b/src/main/java/com/android/tools/r8/graph/RewrittenPrototypeDescription.java
@@ -14,6 +14,8 @@
 import com.android.tools.r8.ir.optimize.info.CallSiteOptimizationInfo;
 import com.android.tools.r8.ir.optimize.info.ConcreteCallSiteOptimizationInfo;
 import com.android.tools.r8.shaking.AppInfoWithLiveness;
+import com.android.tools.r8.utils.ConsumerUtils;
+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;
@@ -405,7 +407,26 @@
               oldIndex -> getArgumentInfo(oldIndex + firstArgumentIndex).isRemovedArgumentInfo());
         };
       }
-      return null;
+      return ConsumerUtils.emptyConsumer();
+    }
+
+    public int getNewArgumentIndex(int argumentIndex) {
+      int numberOfArgumentsRemovedBeforeArgument = 0;
+      Iterator<Entry<ArgumentInfo>> iterator = iterator();
+      while (iterator.hasNext()) {
+        Entry<ArgumentInfo> entry = iterator.next();
+        int argumentIndexForInfo = entry.getIntKey();
+        if (argumentIndexForInfo >= argumentIndex) {
+          break;
+        }
+        ArgumentInfo argumentInfo = entry.getValue();
+        if (argumentInfo.isRemovedArgumentInfo()) {
+          numberOfArgumentsRemovedBeforeArgument++;
+        }
+      }
+      assert IteratorUtils.allRemainingMatchDestructive(
+          iterator, entry -> entry.getIntKey() >= argumentIndex);
+      return argumentIndex - numberOfArgumentsRemovedBeforeArgument;
     }
   }
 
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/enums/classification/CheckNotNullEnumUnboxerMethodClassification.java b/src/main/java/com/android/tools/r8/ir/optimize/enums/classification/CheckNotNullEnumUnboxerMethodClassification.java
index 77d6178..1a466ed 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/enums/classification/CheckNotNullEnumUnboxerMethodClassification.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/enums/classification/CheckNotNullEnumUnboxerMethodClassification.java
@@ -4,13 +4,9 @@
 
 package com.android.tools.r8.ir.optimize.enums.classification;
 
-import com.android.tools.r8.graph.RewrittenPrototypeDescription.ArgumentInfo;
 import com.android.tools.r8.graph.RewrittenPrototypeDescription.ArgumentInfoCollection;
 import com.android.tools.r8.ir.code.InvokeStatic;
 import com.android.tools.r8.ir.code.Value;
-import com.android.tools.r8.utils.IteratorUtils;
-import it.unimi.dsi.fastutil.ints.Int2ReferenceMap.Entry;
-import java.util.Iterator;
 
 public final class CheckNotNullEnumUnboxerMethodClassification
     extends EnumUnboxerMethodClassification {
@@ -52,27 +48,9 @@
       return unknown();
     }
 
-    int numberOfArgumentsRemovedBeforeThis = 0;
-
-    Iterator<Entry<ArgumentInfo>> iterator = removedParameters.iterator();
-    while (iterator.hasNext()) {
-      Entry<ArgumentInfo> entry = iterator.next();
-      int argumentIndexForInfo = entry.getIntKey();
-      if (argumentIndexForInfo >= getArgumentIndex()) {
-        break;
-      }
-      ArgumentInfo argumentInfo = entry.getValue();
-      if (argumentInfo.isRemovedArgumentInfo()) {
-        numberOfArgumentsRemovedBeforeThis++;
-      }
-    }
-
-    assert IteratorUtils.allRemainingMatchDestructive(
-        iterator, entry -> entry.getIntKey() >= getArgumentIndex());
-
-    return numberOfArgumentsRemovedBeforeThis > 0
-        ? new CheckNotNullEnumUnboxerMethodClassification(
-            getArgumentIndex() - numberOfArgumentsRemovedBeforeThis)
+    int newArgumentIndex = removedParameters.getNewArgumentIndex(argumentIndex);
+    return newArgumentIndex != argumentIndex
+        ? new CheckNotNullEnumUnboxerMethodClassification(newArgumentIndex)
         : this;
   }