Don't erase dynamic argument types not needed for reprocessing

If the dynamic type of a given parameter in the program was deemed to not have any effect on reprocessing, we removed the dynamic type information.

This CL changes the argument propagator to not prune such dynamic type information (as it is needed for strengthening parameter types), meanwhile making sure methods are not reprocessed due to the newly retained dynamic types.

Change-Id: Iacdb66d0cb5b56958982c1969047a364b9783067
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/info/ConcreteCallSiteOptimizationInfo.java b/src/main/java/com/android/tools/r8/ir/optimize/info/ConcreteCallSiteOptimizationInfo.java
index 2ee5a6e..e855aa1 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/info/ConcreteCallSiteOptimizationInfo.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/info/ConcreteCallSiteOptimizationInfo.java
@@ -119,6 +119,10 @@
     return constants.getOrDefault(argIndex, UnknownValue.getInstance());
   }
 
+  public Nullability getNullability(int argIndex) {
+    return getDynamicType(argIndex).getNullability();
+  }
+
   public static CallSiteOptimizationInfo fromMethodState(
       AppView<AppInfoWithLiveness> appView,
       ProgramMethod method,
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagator.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagator.java
index 4764635..5f43dd3 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagator.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagator.java
@@ -107,7 +107,7 @@
       reprocessingCriteriaCollection.analyzeArgumentUses(method, code);
     } else {
       assert !methodProcessor.isPrimaryMethodProcessor();
-      assert reprocessingCriteriaCollection == null;
+      assert !methodProcessor.isPostMethodProcessor() || reprocessingCriteriaCollection == null;
     }
   }
 
@@ -170,9 +170,10 @@
             .run(stronglyConnectedProgramComponents, affectedClasses::add, executorService, timing);
 
     // Find all the code objects that need reprocessing.
-    new ArgumentPropagatorMethodReprocessingEnqueuer(appView)
+    new ArgumentPropagatorMethodReprocessingEnqueuer(appView, reprocessingCriteriaCollection)
         .enqueueMethodForReprocessing(
             graphLens, postMethodProcessorBuilder, executorService, timing);
+    reprocessingCriteriaCollection = null;
 
     // Finally, apply the graph lens to the program (i.e., remove constant parameters from method
     // definitions).
@@ -209,7 +210,6 @@
             stronglyConnectedProgramComponents,
             interfaceDispatchOutsideProgram)
         .populateOptimizationInfo(converter, executorService, timing);
-    reprocessingCriteriaCollection = null;
     timing.end();
   }
 
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagatorCodeScanner.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagatorCodeScanner.java
index eef092d..2839aeb 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagatorCodeScanner.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagatorCodeScanner.java
@@ -502,12 +502,6 @@
     // then use UnknownParameterState.
     if (parameterTypeElement.isClassType()) {
       DynamicType dynamicType = argument.getDynamicType(appView);
-      if (!parameterReprocessingCriteria.shouldReprocessDueToDynamicType()) {
-        dynamicType =
-            parameterReprocessingCriteria.widenDynamicClassType(
-                appView, dynamicType, parameterTypeElement.asClassType());
-      }
-
       DynamicType widenedDynamicType =
           WideningUtils.widenDynamicNonReceiverType(appView, dynamicType, parameterType);
       return abstractValue.isUnknown() && widenedDynamicType.isUnknown()
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagatorMethodReprocessingEnqueuer.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagatorMethodReprocessingEnqueuer.java
index f3e9000..d79d23e 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagatorMethodReprocessingEnqueuer.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagatorMethodReprocessingEnqueuer.java
@@ -19,6 +19,7 @@
 import com.android.tools.r8.ir.conversion.IRConverter;
 import com.android.tools.r8.ir.conversion.PostMethodProcessor;
 import com.android.tools.r8.ir.optimize.info.CallSiteOptimizationInfo;
+import com.android.tools.r8.optimize.argumentpropagation.reprocessingcriteria.ArgumentPropagatorReprocessingCriteriaCollection;
 import com.android.tools.r8.shaking.AppInfoWithLiveness;
 import com.android.tools.r8.utils.ThreadUtils;
 import com.android.tools.r8.utils.Timing;
@@ -34,9 +35,13 @@
 public class ArgumentPropagatorMethodReprocessingEnqueuer {
 
   private final AppView<AppInfoWithLiveness> appView;
+  private final ArgumentPropagatorReprocessingCriteriaCollection reprocessingCriteriaCollection;
 
-  public ArgumentPropagatorMethodReprocessingEnqueuer(AppView<AppInfoWithLiveness> appView) {
+  public ArgumentPropagatorMethodReprocessingEnqueuer(
+      AppView<AppInfoWithLiveness> appView,
+      ArgumentPropagatorReprocessingCriteriaCollection reprocessingCriteriaCollection) {
     this.appView = appView;
+    this.reprocessingCriteriaCollection = reprocessingCriteriaCollection;
   }
 
   /**
@@ -61,7 +66,7 @@
     timing.end();
 
     timing.begin("Enqueue methods with non-trivial info");
-    enqueueMethodsWithNonTrivialOptimizationInfo(methodsToReprocessBuilder);
+    enqueueAffectedCallees(graphLens, methodsToReprocessBuilder);
     timing.end();
 
     timing.begin("Enqueue affected methods");
@@ -73,16 +78,34 @@
     timing.end();
   }
 
-  private void enqueueMethodsWithNonTrivialOptimizationInfo(
+  private void enqueueAffectedCallees(
+      ArgumentPropagatorGraphLens graphLens,
       LongLivedProgramMethodSetBuilder<ProgramMethodSet> methodsToReprocessBuilder) {
     GraphLens currentGraphLens = appView.graphLens();
     for (DexProgramClass clazz : appView.appInfo().classes()) {
       clazz.forEachProgramMethodMatching(
           DexEncodedMethod::hasCode,
           method -> {
+            if (method.getDefinition().getCode().isSharedCodeObject()) {
+              return;
+            }
+
+            if (graphLens != null) {
+              DexMethod rewrittenMethodSignature =
+                  graphLens.internalGetNextMethodSignature(method.getReference());
+              if (graphLens.hasPrototypeChanges(rewrittenMethodSignature)) {
+                assert !appView.appInfo().isNeverReprocessMethod(method);
+                methodsToReprocessBuilder.add(method, currentGraphLens);
+                appView.testing().callSiteOptimizationInfoInspector.accept(method);
+                return;
+              }
+            }
+
             CallSiteOptimizationInfo callSiteOptimizationInfo =
                 method.getOptimizationInfo().getArgumentInfos();
-            if (callSiteOptimizationInfo.isConcreteCallSiteOptimizationInfo()
+            if (reprocessingCriteriaCollection
+                    .getReprocessingCriteria(method)
+                    .shouldReprocess(appView, method, callSiteOptimizationInfo)
                 && !appView.appInfo().isNeverReprocessMethod(method)) {
               methodsToReprocessBuilder.add(method, currentGraphLens);
               appView.testing().callSiteOptimizationInfoInspector.accept(method);
@@ -99,6 +122,7 @@
       LongLivedProgramMethodSetBuilder<ProgramMethodSet> methodsToReprocessBuilder,
       ExecutorService executorService)
       throws ExecutionException {
+    GraphLens currentGraphLens = appView.graphLens();
     Collection<List<ProgramMethod>> methodsToReprocess =
         ThreadUtils.processItemsWithResults(
             appView.appInfo().classes(),
@@ -107,12 +131,7 @@
               clazz.forEachProgramMethodMatching(
                   DexEncodedMethod::hasCode,
                   method -> {
-                    if (graphLens.internalGetNextMethodSignature(method.getReference())
-                        != method.getReference()) {
-                      if (!method.getOptimizationInfo().hasBeenInlinedIntoSingleCallSite()) {
-                        methodsToReprocessInClass.add(method);
-                      }
-                    } else {
+                    if (!methodsToReprocessBuilder.contains(method, currentGraphLens)) {
                       AffectedMethodUseRegistry registry =
                           new AffectedMethodUseRegistry(appView, method, graphLens);
                       if (method.registerCodeReferencesWithResult(registry)) {
@@ -124,10 +143,9 @@
               return methodsToReprocessInClass;
             },
             executorService);
-    GraphLens currentGraphLens = appView.graphLens();
     methodsToReprocess.forEach(
-        methodsToReprocessForClass ->
-            methodsToReprocessBuilder.addAll(methodsToReprocessForClass, currentGraphLens));
+        methodsToReprocessInClass ->
+            methodsToReprocessBuilder.addAll(methodsToReprocessInClass, currentGraphLens));
   }
 
   static class AffectedMethodUseRegistry extends UseRegistryWithResult<Boolean, ProgramMethod> {
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagatorOptimizationInfoPopulator.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagatorOptimizationInfoPopulator.java
index 94da691..bd70946 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagatorOptimizationInfoPopulator.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagatorOptimizationInfoPopulator.java
@@ -22,7 +22,6 @@
 import com.android.tools.r8.optimize.argumentpropagation.codescanner.ConcreteClassTypeParameterState;
 import com.android.tools.r8.optimize.argumentpropagation.codescanner.ConcreteMethodState;
 import com.android.tools.r8.optimize.argumentpropagation.codescanner.ConcreteMonomorphicMethodState;
-import com.android.tools.r8.optimize.argumentpropagation.codescanner.ConcreteMonomorphicMethodStateOrUnknown;
 import com.android.tools.r8.optimize.argumentpropagation.codescanner.ConcreteParameterState;
 import com.android.tools.r8.optimize.argumentpropagation.codescanner.ConcretePrimitiveTypeParameterState;
 import com.android.tools.r8.optimize.argumentpropagation.codescanner.MethodState;
@@ -32,7 +31,6 @@
 import com.android.tools.r8.optimize.argumentpropagation.propagation.InterfaceMethodArgumentPropagator;
 import com.android.tools.r8.optimize.argumentpropagation.propagation.VirtualDispatchMethodArgumentPropagator;
 import com.android.tools.r8.optimize.argumentpropagation.reprocessingcriteria.ArgumentPropagatorReprocessingCriteriaCollection;
-import com.android.tools.r8.optimize.argumentpropagation.reprocessingcriteria.MethodReprocessingCriteria;
 import com.android.tools.r8.optimize.argumentpropagation.utils.WideningUtils;
 import com.android.tools.r8.shaking.AppInfoWithLiveness;
 import com.android.tools.r8.utils.ListUtils;
@@ -218,28 +216,18 @@
               return true;
             });
 
-    // If we have any reprocessing criteria for the given method, check that they are satisfied
-    // before reenqueing.
-    MethodReprocessingCriteria reprocessingCriteria =
-        reprocessingCriteriaCollection.getReprocessingCriteria(method);
-    ConcreteMonomorphicMethodStateOrUnknown widenedMethodState =
-        reprocessingCriteria.widenMethodState(appView, method, monomorphicMethodState);
-    if (widenedMethodState.isUnknown()) {
-      return;
-    }
-
-    ConcreteMonomorphicMethodState finalMethodState = widenedMethodState.asMonomorphic();
     getSimpleFeedback()
         .setArgumentInfos(
             method,
-            ConcreteCallSiteOptimizationInfo.fromMethodState(appView, method, finalMethodState));
+            ConcreteCallSiteOptimizationInfo.fromMethodState(
+                appView, method, monomorphicMethodState));
 
     // Strengthen the return value of the method if the method is known to return one of the
     // arguments.
     MethodOptimizationInfo optimizationInfo = method.getOptimizationInfo();
     if (optimizationInfo.returnsArgument()) {
       ParameterState returnedArgumentState =
-          finalMethodState.getParameterState(optimizationInfo.getReturnedArgument());
+          monomorphicMethodState.getParameterState(optimizationInfo.getReturnedArgument());
       OptimizationFeedback.getSimple()
           .methodReturnsAbstractValue(
               method.getDefinition(), appView, returnedArgumentState.getAbstractValue(appView));
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/reprocessingcriteria/AlwaysFalseParameterReprocessingCriteria.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/reprocessingcriteria/AlwaysFalseParameterReprocessingCriteria.java
index 90d7e4a..cfec454 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/reprocessingcriteria/AlwaysFalseParameterReprocessingCriteria.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/reprocessingcriteria/AlwaysFalseParameterReprocessingCriteria.java
@@ -6,7 +6,8 @@
 
 import com.android.tools.r8.graph.AppView;
 import com.android.tools.r8.graph.DexType;
-import com.android.tools.r8.optimize.argumentpropagation.codescanner.ConcreteParameterState;
+import com.android.tools.r8.graph.ProgramMethod;
+import com.android.tools.r8.ir.optimize.info.ConcreteCallSiteOptimizationInfo;
 import com.android.tools.r8.shaking.AppInfoWithLiveness;
 
 /**
@@ -34,7 +35,9 @@
   @Override
   public boolean shouldReprocess(
       AppView<AppInfoWithLiveness> appView,
-      ConcreteParameterState parameterState,
+      ProgramMethod method,
+      ConcreteCallSiteOptimizationInfo methodState,
+      int parameterIndex,
       DexType parameterType) {
     return false;
   }
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/reprocessingcriteria/AlwaysTrueParameterReprocessingCriteria.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/reprocessingcriteria/AlwaysTrueParameterReprocessingCriteria.java
index 72d382f..ff85df7 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/reprocessingcriteria/AlwaysTrueParameterReprocessingCriteria.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/reprocessingcriteria/AlwaysTrueParameterReprocessingCriteria.java
@@ -6,7 +6,8 @@
 
 import com.android.tools.r8.graph.AppView;
 import com.android.tools.r8.graph.DexType;
-import com.android.tools.r8.optimize.argumentpropagation.codescanner.ConcreteParameterState;
+import com.android.tools.r8.graph.ProgramMethod;
+import com.android.tools.r8.ir.optimize.info.ConcreteCallSiteOptimizationInfo;
 import com.android.tools.r8.shaking.AppInfoWithLiveness;
 
 /**
@@ -32,7 +33,9 @@
   @Override
   public boolean shouldReprocess(
       AppView<AppInfoWithLiveness> appView,
-      ConcreteParameterState parameterState,
+      ProgramMethod method,
+      ConcreteCallSiteOptimizationInfo methodState,
+      int parameterIndex,
       DexType parameterType) {
     return true;
   }
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/reprocessingcriteria/MethodReprocessingCriteria.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/reprocessingcriteria/MethodReprocessingCriteria.java
index 658383f..69e0354 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/reprocessingcriteria/MethodReprocessingCriteria.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/reprocessingcriteria/MethodReprocessingCriteria.java
@@ -7,13 +7,9 @@
 import com.android.tools.r8.graph.AppView;
 import com.android.tools.r8.graph.DexType;
 import com.android.tools.r8.graph.ProgramMethod;
-import com.android.tools.r8.optimize.argumentpropagation.codescanner.ConcreteMonomorphicMethodState;
-import com.android.tools.r8.optimize.argumentpropagation.codescanner.ConcreteMonomorphicMethodStateOrUnknown;
-import com.android.tools.r8.optimize.argumentpropagation.codescanner.MethodState;
-import com.android.tools.r8.optimize.argumentpropagation.codescanner.ParameterState;
-import com.android.tools.r8.optimize.argumentpropagation.codescanner.UnknownParameterState;
+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.google.common.collect.Iterables;
 import it.unimi.dsi.fastutil.ints.Int2ReferenceMap;
 import it.unimi.dsi.fastutil.ints.Int2ReferenceOpenHashMap;
 
@@ -43,35 +39,29 @@
         parameterIndex, ParameterReprocessingCriteria.alwaysReprocess());
   }
 
-  public ConcreteMonomorphicMethodStateOrUnknown widenMethodState(
+  public boolean shouldReprocess(
       AppView<AppInfoWithLiveness> appView,
       ProgramMethod method,
-      ConcreteMonomorphicMethodState methodState) {
-    for (int parameterIndex = 0; parameterIndex < methodState.size(); parameterIndex++) {
-      ParameterState parameterState = methodState.getParameterState(parameterIndex);
-      assert !parameterState.isBottom();
-      if (parameterState.isUnknown()) {
-        continue;
+      CallSiteOptimizationInfo methodState) {
+    if (!methodState.isConcreteCallSiteOptimizationInfo()) {
+      return false;
+    }
+    ConcreteCallSiteOptimizationInfo concreteMethodState =
+        methodState.asConcreteCallSiteOptimizationInfo();
+    for (int parameterIndex = 0;
+        parameterIndex < method.getDefinition().getNumberOfArguments();
+        parameterIndex++) {
+      if (methodState.getAbstractArgumentValue(parameterIndex).isSingleValue()) {
+        return true;
       }
-
-      if (parameterState.getAbstractValue(appView).isSingleValue()) {
-        // Don't widen when we have information that can be used for parameter removal.
-        continue;
-      }
-
       ParameterReprocessingCriteria parameterReprocessingCriteria =
           getParameterReprocessingCriteria(parameterIndex);
       DexType parameterType = method.getArgumentType(parameterIndex);
       if (parameterReprocessingCriteria.shouldReprocess(
-          appView, parameterState.asConcrete(), parameterType)) {
-        continue;
+          appView, method, concreteMethodState, parameterIndex, parameterType)) {
+        return true;
       }
-
-      methodState.setParameterState(parameterIndex, UnknownParameterState.get());
     }
-
-    return Iterables.all(methodState.getParameterStates(), ParameterState::isUnknown)
-        ? MethodState.unknown()
-        : methodState;
+    return false;
   }
 }
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/reprocessingcriteria/NonTrivialParameterReprocessingCriteria.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/reprocessingcriteria/NonTrivialParameterReprocessingCriteria.java
index b55bf13..47f01c1 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/reprocessingcriteria/NonTrivialParameterReprocessingCriteria.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/reprocessingcriteria/NonTrivialParameterReprocessingCriteria.java
@@ -6,11 +6,10 @@
 
 import com.android.tools.r8.graph.AppView;
 import com.android.tools.r8.graph.DexType;
+import com.android.tools.r8.graph.ProgramMethod;
 import com.android.tools.r8.ir.analysis.type.DynamicType;
 import com.android.tools.r8.ir.analysis.type.Nullability;
-import com.android.tools.r8.optimize.argumentpropagation.codescanner.ConcreteParameterState;
-import com.android.tools.r8.optimize.argumentpropagation.codescanner.ConcretePrimitiveTypeParameterState;
-import com.android.tools.r8.optimize.argumentpropagation.codescanner.ConcreteReferenceTypeParameterState;
+import com.android.tools.r8.ir.optimize.info.ConcreteCallSiteOptimizationInfo;
 import com.android.tools.r8.optimize.argumentpropagation.utils.WideningUtils;
 import com.android.tools.r8.shaking.AppInfoWithLiveness;
 
@@ -31,13 +30,16 @@
   @Override
   public boolean shouldReprocess(
       AppView<AppInfoWithLiveness> appView,
-      ConcreteParameterState parameterState,
+      ProgramMethod method,
+      ConcreteCallSiteOptimizationInfo methodState,
+      int parameterIndex,
       DexType parameterType) {
-    if (parameterState.isReferenceParameter()) {
-      return shouldReprocess(appView, parameterState.asReferenceParameter(), parameterType);
+    if (parameterType.isReferenceType()) {
+      return shouldReprocessReferenceParameter(
+          appView, method, methodState, parameterIndex, parameterType);
     } else {
-      assert parameterState.isPrimitiveParameter();
-      return shouldReprocess(appView, parameterState.asPrimitiveParameter(), parameterType);
+      assert parameterType.isPrimitiveType();
+      return shouldReprocessPrimitiveParameter(methodState, parameterIndex);
     }
   }
 
@@ -56,34 +58,35 @@
     return true;
   }
 
-  private boolean shouldReprocess(
-      AppView<AppInfoWithLiveness> appView,
-      ConcretePrimitiveTypeParameterState parameterState,
-      DexType parameterType) {
-    return true;
+  private boolean shouldReprocessPrimitiveParameter(
+      ConcreteCallSiteOptimizationInfo methodState, int parameterIndex) {
+    return methodState.getAbstractArgumentValue(parameterIndex).isNonTrivial();
   }
 
-  private boolean shouldReprocess(
+  private boolean shouldReprocessReferenceParameter(
       AppView<AppInfoWithLiveness> appView,
-      ConcreteReferenceTypeParameterState parameterState,
+      ProgramMethod method,
+      ConcreteCallSiteOptimizationInfo methodState,
+      int parameterIndex,
       DexType parameterType) {
     if (shouldReprocessDueToAbstractValue()
-        && !parameterState.getAbstractValue(appView).isUnknown()) {
+        && !methodState.getAbstractArgumentValue(parameterIndex).isUnknown()) {
       return true;
     }
     if (shouldReprocessDueToDynamicType()) {
       DynamicType widenedDynamicType =
           WideningUtils.widenDynamicNonReceiverType(
               appView,
-              parameterState.getDynamicType().withNullability(Nullability.maybeNull()),
+              methodState.getDynamicType(parameterIndex).withNullability(Nullability.maybeNull()),
               parameterType);
       if (!widenedDynamicType.isUnknown()) {
         return true;
       }
     }
+    boolean isReceiverParameter = parameterIndex == 0 && method.getDefinition().isInstance();
     if (shouldReprocessDueToNullability()
-        && !parameterState.isReceiverParameter()
-        && !parameterState.getNullability().isUnknown()) {
+        && !isReceiverParameter
+        && !methodState.getNullability(parameterIndex).isUnknown()) {
       return true;
     }
     return false;
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/reprocessingcriteria/ParameterReprocessingCriteria.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/reprocessingcriteria/ParameterReprocessingCriteria.java
index d7c8890..9fb773e 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/reprocessingcriteria/ParameterReprocessingCriteria.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/reprocessingcriteria/ParameterReprocessingCriteria.java
@@ -6,9 +6,10 @@
 
 import com.android.tools.r8.graph.AppView;
 import com.android.tools.r8.graph.DexType;
+import com.android.tools.r8.graph.ProgramMethod;
 import com.android.tools.r8.ir.analysis.type.ClassTypeElement;
 import com.android.tools.r8.ir.analysis.type.DynamicType;
-import com.android.tools.r8.optimize.argumentpropagation.codescanner.ConcreteParameterState;
+import com.android.tools.r8.ir.optimize.info.ConcreteCallSiteOptimizationInfo;
 import com.android.tools.r8.shaking.AppInfoWithLiveness;
 
 public abstract class ParameterReprocessingCriteria {
@@ -35,7 +36,9 @@
 
   public abstract boolean shouldReprocess(
       AppView<AppInfoWithLiveness> appView,
-      ConcreteParameterState parameterState,
+      ProgramMethod method,
+      ConcreteCallSiteOptimizationInfo methodState,
+      int parameterIndex,
       DexType parameterType);
 
   public abstract boolean shouldReprocessDueToAbstractValue();
diff --git a/src/test/java/com/android/tools/r8/ir/optimize/callsites/nullability/InvokeDirectPositiveTest.java b/src/test/java/com/android/tools/r8/ir/optimize/callsites/nullability/InvokeDirectPositiveTest.java
index fd605ac..11c44ad 100644
--- a/src/test/java/com/android/tools/r8/ir/optimize/callsites/nullability/InvokeDirectPositiveTest.java
+++ b/src/test/java/com/android/tools/r8/ir/optimize/callsites/nullability/InvokeDirectPositiveTest.java
@@ -57,7 +57,7 @@
         : "Unexpected revisit: " + method.toSourceString();
     CallSiteOptimizationInfo callSiteOptimizationInfo =
         method.getOptimizationInfo().getArgumentInfos();
-    assert callSiteOptimizationInfo.getDynamicType(1).isNotNullType();
+    assertTrue(callSiteOptimizationInfo.getDynamicType(1).getNullability().isDefinitelyNotNull());
   }
 
   private void inspect(CodeInspector inspector) {
diff --git a/src/test/java/com/android/tools/r8/ir/optimize/callsites/nullability/InvokeStaticPositiveTest.java b/src/test/java/com/android/tools/r8/ir/optimize/callsites/nullability/InvokeStaticPositiveTest.java
index cf5cbe1..d181d3b 100644
--- a/src/test/java/com/android/tools/r8/ir/optimize/callsites/nullability/InvokeStaticPositiveTest.java
+++ b/src/test/java/com/android/tools/r8/ir/optimize/callsites/nullability/InvokeStaticPositiveTest.java
@@ -55,7 +55,7 @@
         : "Unexpected revisit: " + method.toSourceString();
     CallSiteOptimizationInfo callSiteOptimizationInfo =
         method.getOptimizationInfo().getArgumentInfos();
-    assert callSiteOptimizationInfo.getDynamicType(0).isNotNullType();
+    assertTrue(callSiteOptimizationInfo.getDynamicType(0).getNullability().isDefinitelyNotNull());
   }
 
   private void inspect(CodeInspector inspector) {