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 4f374e6..8b19ac3 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
@@ -211,644 +211,610 @@
       IRCode code,
       AbstractValueSupplier abstractValueSupplier,
       Timing timing) {
-    timing.begin("Argument propagation scanner");
-    for (Instruction instruction : code.instructions()) {
-      if (instruction.isFieldPut()) {
-        scan(instruction.asFieldPut(), abstractValueSupplier, method, timing);
-      } else if (instruction.isInvokeMethod()) {
-        scan(instruction.asInvokeMethod(), abstractValueSupplier, method, timing);
-      } else if (instruction.isInvokeCustom()) {
-        scan(instruction.asInvokeCustom());
-      }
-    }
-    timing.end();
+    new CodeScanner(abstractValueSupplier, code, method).scan(timing);
   }
 
-  private void scan(
-      FieldPut fieldPut,
-      AbstractValueSupplier abstractValueSupplier,
-      ProgramMethod context,
-      Timing timing) {
-    ProgramField field = fieldPut.resolveField(appView, context).getProgramField();
-    if (field == null) {
-      // Nothing to propagate.
-      return;
-    }
-    addTemporaryFieldState(fieldPut, field, abstractValueSupplier, context, timing);
-  }
+  protected class CodeScanner {
 
-  private void addTemporaryFieldState(
-      FieldPut fieldPut,
-      ProgramField field,
-      AbstractValueSupplier abstractValueSupplier,
-      ProgramMethod context,
-      Timing timing) {
-    timing.begin("Add field state");
-    fieldStates.addTemporaryFieldState(
-        field,
-        () -> computeFieldState(fieldPut, field, abstractValueSupplier, context, timing),
-        timing,
-        (existingFieldState, fieldStateToAdd) -> {
-          DexType inStaticType = null;
-          NonEmptyValueState newFieldState =
-              existingFieldState.mutableJoin(
-                  appView,
-                  fieldStateToAdd,
-                  inStaticType,
-                  field.getType(),
-                  StateCloner.getCloner(),
-                  Action.empty());
-          return narrowFieldState(field, newFieldState);
-        });
-    timing.end();
-  }
+    protected final AbstractValueSupplier abstractValueSupplier;
+    protected final IRCode code;
+    protected final ProgramMethod context;
 
-  private NonEmptyValueState computeFieldState(
-      FieldPut fieldPut,
-      ProgramField resolvedField,
-      AbstractValueSupplier abstractValueSupplier,
-      ProgramMethod context,
-      Timing timing) {
-    timing.begin("Compute field state for field-put");
-    NonEmptyValueState result =
-        computeFieldState(fieldPut, resolvedField, abstractValueSupplier, context);
-    timing.end();
-    return result;
-  }
-
-  private NonEmptyValueState computeFieldState(
-      FieldPut fieldPut,
-      ProgramField field,
-      AbstractValueSupplier abstractValueSupplier,
-      ProgramMethod context) {
-    TypeElement fieldType = field.getType().toTypeElement(appView);
-    if (!fieldPut.value().getType().lessThanOrEqual(fieldType, appView)) {
-      return ValueState.unknown();
+    protected CodeScanner(
+        AbstractValueSupplier abstractValueSupplier, IRCode code, ProgramMethod method) {
+      this.abstractValueSupplier = abstractValueSupplier;
+      this.code = code;
+      this.context = method;
     }
 
-    NonEmptyValueState inFlowState = computeInFlowState(field.getType(), fieldPut.value(), context);
-    if (inFlowState != null) {
-      return inFlowState;
-    }
-
-    if (field.getType().isArrayType()) {
-      Nullability nullability = fieldPut.value().getType().nullability();
-      return ConcreteArrayTypeValueState.create(nullability);
-    }
-
-    AbstractValue abstractValue = abstractValueSupplier.getAbstractValue(fieldPut.value());
-    if (abstractValue.isUnknown()) {
-      abstractValue =
-          getFallbackAbstractValueForField(
-              field,
-              () -> ObjectStateAnalysis.computeObjectState(fieldPut.value(), appView, context));
-    }
-    if (field.getType().isClassType()) {
-      DynamicType dynamicType =
-          WideningUtils.widenDynamicNonReceiverType(
-              appView, fieldPut.value().getDynamicType(appView), field.getType());
-      return ConcreteClassTypeValueState.create(abstractValue, dynamicType);
-    } else {
-      assert field.getType().isPrimitiveType();
-      return ConcretePrimitiveTypeValueState.create(abstractValue);
-    }
-  }
-
-  // If the value is an argument of the enclosing method or defined by a field-get, then clearly we
-  // have no information about its abstract value (yet). Instead of treating this as having an
-  // unknown runtime value, we instead record a flow constraint.
-  private InFlow computeInFlow(DexType staticType, Value value, ProgramMethod context) {
-    Value valueRoot = value.getAliasedValue(aliasedValueConfiguration);
-    if (valueRoot.isArgument()) {
-      MethodParameter inParameter =
-          methodParameterFactory.create(context, valueRoot.getDefinition().asArgument().getIndex());
-      return castBaseInFlow(widenBaseInFlow(staticType, inParameter, context), value);
-    } else if (valueRoot.isDefinedByInstructionSatisfying(Instruction::isFieldGet)) {
-      FieldGet fieldGet = valueRoot.getDefinition().asFieldGet();
-      ProgramField field = fieldGet.resolveField(appView, context).getProgramField();
-      if (field == null) {
-        return null;
-      }
-      if (fieldGet.isInstanceGet()) {
-        Value receiverValue = fieldGet.asInstanceGet().object();
-        BaseInFlow receiverInFlow =
-            asBaseInFlowOrNull(computeInFlow(staticType, receiverValue, context));
-        if (receiverInFlow != null
-            && receiverInFlow.equals(widenBaseInFlow(staticType, receiverInFlow, context))) {
-          return new InstanceFieldReadAbstractFunction(receiverInFlow, field.getReference());
+    public void scan(Timing timing) {
+      timing.begin("Argument propagation scanner");
+      for (Instruction instruction : code.instructions()) {
+        if (instruction.isFieldPut()) {
+          scanFieldPut(instruction.asFieldPut(), timing);
+        } else if (instruction.isInvokeMethod()) {
+          scanInvoke(instruction.asInvokeMethod(), timing);
+        } else if (instruction.isInvokeCustom()) {
+          scanInvokeCustom(instruction.asInvokeCustom());
         }
       }
-      return castBaseInFlow(
-          widenBaseInFlow(staticType, fieldValueFactory.create(field), context), value);
+      timing.end();
     }
-    return null;
-  }
 
-  private InFlow castBaseInFlow(InFlow inFlow, Value value) {
-    if (inFlow.isUnknownAbstractFunction()) {
-      return inFlow;
-    }
-    assert inFlow.isBaseInFlow();
-    Value valueRoot = value.getAliasedValue();
-    if (!valueRoot.isDefinedByInstructionSatisfying(Instruction::isCheckCast)) {
-      return inFlow;
-    }
-    CheckCast checkCast = valueRoot.getDefinition().asCheckCast();
-    return new CastAbstractFunction(inFlow.asBaseInFlow(), checkCast.getType());
-  }
-
-  private InFlow widenBaseInFlow(DexType staticType, BaseInFlow inFlow, ProgramMethod context) {
-    if (inFlow.isFieldValue()) {
-      if (isFieldValueAlreadyUnknown(staticType, inFlow.asFieldValue().getField())) {
-        return AbstractFunction.unknown();
+    private void scanFieldPut(FieldPut fieldPut, Timing timing) {
+      ProgramField field = fieldPut.resolveField(appView, context).getProgramField();
+      if (field == null) {
+        // Nothing to propagate.
+        return;
       }
-    } else {
-      assert inFlow.isMethodParameter();
-      if (isMethodParameterAlreadyUnknown(staticType, inFlow.asMethodParameter(), context)) {
-        return AbstractFunction.unknown();
-      }
+      addTemporaryFieldState(fieldPut, field, timing);
     }
-    return inFlow;
-  }
 
-  private NonEmptyValueState computeInFlowState(
-      DexType staticType, Value value, ProgramMethod context) {
-    InFlow inFlow = computeInFlow(staticType, value, context);
-    if (inFlow == null) {
-      return null;
+    private void addTemporaryFieldState(FieldPut fieldPut, ProgramField field, Timing timing) {
+      timing.begin("Add field state");
+      fieldStates.addTemporaryFieldState(
+          field,
+          () -> computeFieldState(fieldPut, field, timing),
+          timing,
+          (existingFieldState, fieldStateToAdd) -> {
+            DexType inStaticType = null;
+            NonEmptyValueState newFieldState =
+                existingFieldState.mutableJoin(
+                    appView,
+                    fieldStateToAdd,
+                    inStaticType,
+                    field.getType(),
+                    StateCloner.getCloner(),
+                    Action.empty());
+            return narrowFieldState(field, newFieldState);
+          });
+      timing.end();
     }
-    if (inFlow.isUnknownAbstractFunction()) {
-      return ValueState.unknown();
-    }
-    assert inFlow.isBaseInFlow()
-        || inFlow.isCastAbstractFunction()
-        || inFlow.isInstanceFieldReadAbstractFunction();
-    return ConcreteValueState.create(staticType, inFlow);
-  }
 
-  // Strengthens the abstract value of static final fields to a (self-)SingleFieldValue when the
-  // abstract value is unknown. The soundness of this is based on the fact that static final fields
-  // will never have their value changed after the <clinit> finishes, so value in a static final
-  // field can always be rematerialized by reading the field.
-  private NonEmptyValueState narrowFieldState(ProgramField field, NonEmptyValueState fieldState) {
-    AbstractValue fallbackAbstractValue =
-        getFallbackAbstractValueForField(field, ObjectState::empty);
-    if (!fallbackAbstractValue.isUnknown()) {
-      AbstractValue abstractValue = fieldState.getAbstractValue(appView);
-      if (!abstractValue.isUnknown()) {
-        return fieldState;
+    private NonEmptyValueState computeFieldState(
+        FieldPut fieldPut, ProgramField resolvedField, Timing timing) {
+      timing.begin("Compute field state for field-put");
+      NonEmptyValueState result = computeFieldState(fieldPut, resolvedField);
+      timing.end();
+      return result;
+    }
+
+    private NonEmptyValueState computeFieldState(FieldPut fieldPut, ProgramField field) {
+      TypeElement fieldType = field.getType().toTypeElement(appView);
+      if (!fieldPut.value().getType().lessThanOrEqual(fieldType, appView)) {
+        return ValueState.unknown();
       }
+
+      NonEmptyValueState inFlowState =
+          computeInFlowState(field.getType(), fieldPut.value(), context);
+      if (inFlowState != null) {
+        return inFlowState;
+      }
+
       if (field.getType().isArrayType()) {
-        // We do not track an abstract value for array types.
-        return fieldState;
+        Nullability nullability = fieldPut.value().getType().nullability();
+        return ConcreteArrayTypeValueState.create(nullability);
+      }
+
+      AbstractValue abstractValue = abstractValueSupplier.getAbstractValue(fieldPut.value());
+      if (abstractValue.isUnknown()) {
+        abstractValue =
+            getFallbackAbstractValueForField(
+                field,
+                () -> ObjectStateAnalysis.computeObjectState(fieldPut.value(), appView, context));
       }
       if (field.getType().isClassType()) {
         DynamicType dynamicType =
-            fieldState.isReferenceState()
-                ? fieldState.asReferenceState().getDynamicType()
-                : DynamicType.unknown();
-        return new ConcreteClassTypeValueState(fallbackAbstractValue, dynamicType);
+            WideningUtils.widenDynamicNonReceiverType(
+                appView, fieldPut.value().getDynamicType(appView), field.getType());
+        return ConcreteClassTypeValueState.create(abstractValue, dynamicType);
       } else {
         assert field.getType().isPrimitiveType();
-        return new ConcretePrimitiveTypeValueState(fallbackAbstractValue);
+        return ConcretePrimitiveTypeValueState.create(abstractValue);
       }
     }
-    return fieldState;
-  }
 
-  private AbstractValue getFallbackAbstractValueForField(
-      ProgramField field, Supplier<ObjectState> objectStateSupplier) {
-    if (field.isFinalOrEffectivelyFinal(appView) && field.getAccessFlags().isStatic()) {
-      return appView
-          .abstractValueFactory()
-          .createSingleFieldValue(field.getReference(), objectStateSupplier.get());
-    }
-    return AbstractValue.unknown();
-  }
-
-  private void scan(
-      InvokeMethod invoke,
-      AbstractValueSupplier abstractValueSupplier,
-      ProgramMethod context,
-      Timing timing) {
-    DexMethod invokedMethod = invoke.getInvokedMethod();
-    if (invokedMethod.getHolderType().isArrayType()) {
-      // Nothing to propagate; the targeted method is not a program method.
-      return;
+    // If the value is an argument of the enclosing method or defined by a field-get, then clearly
+    // we have no information about its abstract value (yet). Instead of treating this as having an
+    // unknown runtime value, we instead record a flow constraint.
+    private InFlow computeInFlow(DexType staticType, Value value, ProgramMethod context) {
+      Value valueRoot = value.getAliasedValue(aliasedValueConfiguration);
+      if (valueRoot.isArgument()) {
+        MethodParameter inParameter =
+            methodParameterFactory.create(
+                context, valueRoot.getDefinition().asArgument().getIndex());
+        return castBaseInFlow(widenBaseInFlow(staticType, inParameter, context), value);
+      } else if (valueRoot.isDefinedByInstructionSatisfying(Instruction::isFieldGet)) {
+        FieldGet fieldGet = valueRoot.getDefinition().asFieldGet();
+        ProgramField field = fieldGet.resolveField(appView, context).getProgramField();
+        if (field == null) {
+          return null;
+        }
+        if (fieldGet.isInstanceGet()) {
+          Value receiverValue = fieldGet.asInstanceGet().object();
+          BaseInFlow receiverInFlow =
+              asBaseInFlowOrNull(computeInFlow(staticType, receiverValue, context));
+          if (receiverInFlow != null
+              && receiverInFlow.equals(widenBaseInFlow(staticType, receiverInFlow, context))) {
+            return new InstanceFieldReadAbstractFunction(receiverInFlow, field.getReference());
+          }
+        }
+        return castBaseInFlow(
+            widenBaseInFlow(staticType, fieldValueFactory.create(field), context), value);
+      }
+      return null;
     }
 
-    if (appView.options().testing.checkReceiverAlwaysNullInCallSiteOptimization
-        && invoke.isInvokeMethodWithReceiver()
-        && invoke.asInvokeMethodWithReceiver().getReceiver().isAlwaysNull(appView)) {
-      // Nothing to propagate; the invoke instruction always fails.
-      return;
+    private InFlow castBaseInFlow(InFlow inFlow, Value value) {
+      if (inFlow.isUnknownAbstractFunction()) {
+        return inFlow;
+      }
+      assert inFlow.isBaseInFlow();
+      Value valueRoot = value.getAliasedValue();
+      if (!valueRoot.isDefinedByInstructionSatisfying(Instruction::isCheckCast)) {
+        return inFlow;
+      }
+      CheckCast checkCast = valueRoot.getDefinition().asCheckCast();
+      return new CastAbstractFunction(inFlow.asBaseInFlow(), checkCast.getType());
     }
 
-    SingleResolutionResult<?> resolutionResult =
-        invoke.resolveMethod(appView, context).asSingleResolution();
-    if (resolutionResult == null) {
-      // Nothing to propagate; the invoke instruction fails.
-      return;
+    private InFlow widenBaseInFlow(DexType staticType, BaseInFlow inFlow, ProgramMethod context) {
+      if (inFlow.isFieldValue()) {
+        if (isFieldValueAlreadyUnknown(staticType, inFlow.asFieldValue().getField())) {
+          return AbstractFunction.unknown();
+        }
+      } else {
+        assert inFlow.isMethodParameter();
+        if (isMethodParameterAlreadyUnknown(staticType, inFlow.asMethodParameter(), context)) {
+          return AbstractFunction.unknown();
+        }
+      }
+      return inFlow;
     }
 
-    if (!resolutionResult.getResolvedHolder().isProgramClass()) {
-      // Nothing to propagate; this could dispatch to a program method, but we cannot optimize
-      // methods that override non-program methods.
-      return;
+    private NonEmptyValueState computeInFlowState(
+        DexType staticType, Value value, ProgramMethod context) {
+      InFlow inFlow = computeInFlow(staticType, value, context);
+      if (inFlow == null) {
+        return null;
+      }
+      if (inFlow.isUnknownAbstractFunction()) {
+        return ValueState.unknown();
+      }
+      assert inFlow.isBaseInFlow()
+          || inFlow.isCastAbstractFunction()
+          || inFlow.isInstanceFieldReadAbstractFunction();
+      return ConcreteValueState.create(staticType, inFlow);
     }
 
-    ProgramMethod resolvedMethod = resolutionResult.getResolvedProgramMethod();
-    if (resolvedMethod.getDefinition().isLibraryMethodOverride().isPossiblyTrue()) {
-      assert resolvedMethod.getDefinition().isLibraryMethodOverride().isTrue();
-      // Nothing to propagate; we don't know anything about methods that can be called from outside
-      // the program.
-      return;
+    // Strengthens the abstract value of static final fields to a (self-)SingleFieldValue when the
+    // abstract value is unknown. The soundness of this is based on the fact that static final
+    // fields will never have their value changed after the <clinit> finishes, so value in a static
+    // final field can always be rematerialized by reading the field.
+    private NonEmptyValueState narrowFieldState(ProgramField field, NonEmptyValueState fieldState) {
+      AbstractValue fallbackAbstractValue =
+          getFallbackAbstractValueForField(field, ObjectState::empty);
+      if (!fallbackAbstractValue.isUnknown()) {
+        AbstractValue abstractValue = fieldState.getAbstractValue(appView);
+        if (!abstractValue.isUnknown()) {
+          return fieldState;
+        }
+        if (field.getType().isArrayType()) {
+          // We do not track an abstract value for array types.
+          return fieldState;
+        }
+        if (field.getType().isClassType()) {
+          DynamicType dynamicType =
+              fieldState.isReferenceState()
+                  ? fieldState.asReferenceState().getDynamicType()
+                  : DynamicType.unknown();
+          return new ConcreteClassTypeValueState(fallbackAbstractValue, dynamicType);
+        } else {
+          assert field.getType().isPrimitiveType();
+          return new ConcretePrimitiveTypeValueState(fallbackAbstractValue);
+        }
+      }
+      return fieldState;
     }
 
-    if (invoke.arguments().size() != resolvedMethod.getDefinition().getNumberOfArguments()
-        || invoke.isInvokeStatic() != resolvedMethod.getAccessFlags().isStatic()) {
-      // Nothing to propagate; the invoke instruction fails.
-      return;
+    private AbstractValue getFallbackAbstractValueForField(
+        ProgramField field, Supplier<ObjectState> objectStateSupplier) {
+      if (field.isFinalOrEffectivelyFinal(appView) && field.getAccessFlags().isStatic()) {
+        return appView
+            .abstractValueFactory()
+            .createSingleFieldValue(field.getReference(), objectStateSupplier.get());
+      }
+      return AbstractValue.unknown();
     }
 
-    if (invoke.isInvokeInterface()) {
-      if (!resolutionResult.getInitialResolutionHolder().isInterface()) {
+    private void scanInvoke(InvokeMethod invoke, Timing timing) {
+      DexMethod invokedMethod = invoke.getInvokedMethod();
+      if (invokedMethod.getHolderType().isArrayType()) {
+        // Nothing to propagate; the targeted method is not a program method.
+        return;
+      }
+
+      if (appView.options().testing.checkReceiverAlwaysNullInCallSiteOptimization
+          && invoke.isInvokeMethodWithReceiver()
+          && invoke.asInvokeMethodWithReceiver().getReceiver().isAlwaysNull(appView)) {
+        // Nothing to propagate; the invoke instruction always fails.
+        return;
+      }
+
+      SingleResolutionResult<?> resolutionResult =
+          invoke.resolveMethod(appView, context).asSingleResolution();
+      if (resolutionResult == null) {
         // Nothing to propagate; the invoke instruction fails.
         return;
       }
-    }
 
-    if (invoke.isInvokeSuper()) {
-      // Use the super target instead of the resolved method to ensure that we propagate the
-      // argument information to the targeted method.
-      DexClassAndMethod target =
-          resolutionResult.lookupInvokeSuperTarget(context.getHolder(), appView);
-      if (target == null) {
+      if (!resolutionResult.getResolvedHolder().isProgramClass()) {
+        // Nothing to propagate; this could dispatch to a program method, but we cannot optimize
+        // methods that override non-program methods.
+        return;
+      }
+
+      ProgramMethod resolvedMethod = resolutionResult.getResolvedProgramMethod();
+      if (resolvedMethod.getDefinition().isLibraryMethodOverride().isPossiblyTrue()) {
+        assert resolvedMethod.getDefinition().isLibraryMethodOverride().isTrue();
+        // Nothing to propagate; we don't know anything about methods that can be called from
+        // outside the program.
+        return;
+      }
+
+      if (invoke.arguments().size() != resolvedMethod.getDefinition().getNumberOfArguments()
+          || invoke.isInvokeStatic() != resolvedMethod.getAccessFlags().isStatic()) {
         // Nothing to propagate; the invoke instruction fails.
         return;
       }
-      if (!target.isProgramMethod()) {
-        throw new Unreachable(
-            "Expected super target of a non-library override to be a program method ("
-                + "resolved program method: "
-                + resolvedMethod
-                + ", "
-                + "super non-program method: "
-                + target
-                + ")");
+
+      if (invoke.isInvokeInterface()) {
+        if (!resolutionResult.getInitialResolutionHolder().isInterface()) {
+          // Nothing to propagate; the invoke instruction fails.
+          return;
+        }
       }
-      resolvedMethod = target.asProgramMethod();
+
+      if (invoke.isInvokeSuper()) {
+        // Use the super target instead of the resolved method to ensure that we propagate the
+        // argument information to the targeted method.
+        DexClassAndMethod target =
+            resolutionResult.lookupInvokeSuperTarget(context.getHolder(), appView);
+        if (target == null) {
+          // Nothing to propagate; the invoke instruction fails.
+          return;
+        }
+        if (!target.isProgramMethod()) {
+          throw new Unreachable(
+              "Expected super target of a non-library override to be a program method ("
+                  + "resolved program method: "
+                  + resolvedMethod
+                  + ", "
+                  + "super non-program method: "
+                  + target
+                  + ")");
+        }
+        resolvedMethod = target.asProgramMethod();
+      }
+
+      // Find the method where to store the information about the arguments from this invoke.
+      // If the invoke may dispatch to more than one method, we intentionally do not compute all
+      // possible dispatch targets and propagate the information to these methods (this is
+      // expensive). Instead we record the information in one place and then later propagate the
+      // information to all dispatch targets.
+      addTemporaryMethodState(invoke, resolvedMethod, timing);
     }
 
-    // Find the method where to store the information about the arguments from this invoke.
-    // If the invoke may dispatch to more than one method, we intentionally do not compute all
-    // possible dispatch targets and propagate the information to these methods (this is expensive).
-    // Instead we record the information in one place and then later propagate the information to
-    // all dispatch targets.
-    addTemporaryMethodState(invoke, resolvedMethod, abstractValueSupplier, context, timing);
-  }
+    protected void addTemporaryMethodState(
+        InvokeMethod invoke, ProgramMethod resolvedMethod, Timing timing) {
+      timing.begin("Add method state");
+      methodStates.addTemporaryMethodState(
+          appView,
+          getRepresentative(invoke, resolvedMethod),
+          existingMethodState ->
+              computeMethodState(invoke, resolvedMethod, existingMethodState, timing),
+          timing);
+      timing.end();
+    }
 
-  protected void addTemporaryMethodState(
-      InvokeMethod invoke,
-      ProgramMethod resolvedMethod,
-      AbstractValueSupplier abstractValueSupplier,
-      ProgramMethod context,
-      Timing timing) {
-    timing.begin("Add method state");
-    methodStates.addTemporaryMethodState(
-        appView,
-        getRepresentative(invoke, resolvedMethod),
-        existingMethodState ->
-            computeMethodState(
+    private MethodState computeMethodState(
+        InvokeMethod invoke,
+        ProgramMethod resolvedMethod,
+        MethodState existingMethodState,
+        Timing timing) {
+      assert !existingMethodState.isUnknown();
+
+      // If this invoke may target at most one method, then we compute a state that maps each
+      // parameter to the abstract value and dynamic type provided by this call site. Otherwise, we
+      // compute a polymorphic method state, which includes information about the receiver's dynamic
+      // type bounds.
+      timing.begin("Compute method state for invoke");
+      MethodState result;
+      if (shouldUsePolymorphicMethodState(invoke, resolvedMethod)) {
+        assert existingMethodState.isBottom() || existingMethodState.isPolymorphic();
+        result =
+            computePolymorphicMethodState(
+                invoke.asInvokeMethodWithReceiver(),
+                resolvedMethod,
+                existingMethodState.asPolymorphicOrBottom());
+      } else {
+        assert existingMethodState.isBottom() || existingMethodState.isMonomorphic();
+        result =
+            computeMonomorphicMethodState(
                 invoke,
                 resolvedMethod,
-                abstractValueSupplier,
-                context,
-                existingMethodState,
-                timing),
-        timing);
-    timing.end();
-  }
+                invoke.lookupSingleProgramTarget(appView, context),
+                existingMethodState.asMonomorphicOrBottom());
+      }
+      timing.end();
+      return result;
+    }
 
-  private MethodState computeMethodState(
-      InvokeMethod invoke,
-      ProgramMethod resolvedMethod,
-      AbstractValueSupplier abstractValueSupplier,
-      ProgramMethod context,
-      MethodState existingMethodState,
-      Timing timing) {
-    assert !existingMethodState.isUnknown();
+    // TODO(b/190154391): Add a strategy that widens the dynamic receiver type to allow easily
+    //  experimenting with the performance/size trade-off between precise/imprecise handling of
+    //  dynamic dispatch.
+    private MethodState computePolymorphicMethodState(
+        InvokeMethodWithReceiver invoke,
+        ProgramMethod resolvedMethod,
+        ConcretePolymorphicMethodStateOrBottom existingMethodState) {
+      DynamicTypeWithUpperBound dynamicReceiverType = invoke.getReceiver().getDynamicType(appView);
+      // TODO(b/331587404): Investigate if we can replace the receiver by null before entering this
+      //  pass, so that this special case is not needed.
+      if (dynamicReceiverType.isNullType()) {
+        assert appView.testing().allowNullDynamicTypeInCodeScanner : "b/250634405";
+        // This can happen if we were unable to determine that the receiver is a phi value where
+        // null information has not been propagated down. Ideally this case would never happen as it
+        // should be possible to replace the receiver by the null constant in this case. Since the
+        // receiver is known to be null, no argument information should be propagated to the
+        // callees, so we return bottom here.
+        return MethodState.bottom();
+      }
 
-    // If this invoke may target at most one method, then we compute a state that maps each
-    // parameter to the abstract value and dynamic type provided by this call site. Otherwise, we
-    // compute a polymorphic method state, which includes information about the receiver's dynamic
-    // type bounds.
-    timing.begin("Compute method state for invoke");
-    MethodState result;
-    if (shouldUsePolymorphicMethodState(invoke, resolvedMethod)) {
-      assert existingMethodState.isBottom() || existingMethodState.isPolymorphic();
-      result =
-          computePolymorphicMethodState(
-              invoke.asInvokeMethodWithReceiver(),
-              resolvedMethod,
-              abstractValueSupplier,
-              context,
-              existingMethodState.asPolymorphicOrBottom());
-    } else {
-      assert existingMethodState.isBottom() || existingMethodState.isMonomorphic();
-      result =
+      ProgramMethod singleTarget = invoke.lookupSingleProgramTarget(appView, context);
+      DynamicTypeWithUpperBound bounds =
+          computeBoundsForPolymorphicMethodState(resolvedMethod, singleTarget, dynamicReceiverType);
+      MethodState existingMethodStateForBounds =
+          existingMethodState.isPolymorphic()
+              ? existingMethodState.asPolymorphic().getMethodStateForBounds(bounds)
+              : MethodState.bottom();
+
+      if (existingMethodStateForBounds.isPolymorphic()) {
+        assert false;
+        return MethodState.unknown();
+      }
+
+      // If we already don't know anything about the parameters for the given type bounds, then
+      // don't compute a method state.
+      if (existingMethodStateForBounds.isUnknown()) {
+        return MethodState.bottom();
+      }
+
+      ConcreteMonomorphicMethodStateOrUnknown methodStateForBounds =
           computeMonomorphicMethodState(
               invoke,
               resolvedMethod,
-              invoke.lookupSingleProgramTarget(appView, context),
-              abstractValueSupplier,
-              context,
-              existingMethodState.asMonomorphicOrBottom());
-    }
-    timing.end();
-    return result;
-  }
-
-  // TODO(b/190154391): Add a strategy that widens the dynamic receiver type to allow easily
-  //  experimenting with the performance/size trade-off between precise/imprecise handling of
-  //  dynamic dispatch.
-  private MethodState computePolymorphicMethodState(
-      InvokeMethodWithReceiver invoke,
-      ProgramMethod resolvedMethod,
-      AbstractValueSupplier abstractValueSupplier,
-      ProgramMethod context,
-      ConcretePolymorphicMethodStateOrBottom existingMethodState) {
-    DynamicTypeWithUpperBound dynamicReceiverType = invoke.getReceiver().getDynamicType(appView);
-    // TODO(b/331587404): Investigate if we can replace the receiver by null before entering this
-    //  pass, so that this special case is not needed.
-    if (dynamicReceiverType.isNullType()) {
-      assert appView.testing().allowNullDynamicTypeInCodeScanner : "b/250634405";
-      // This can happen if we were unable to determine that the receiver is a phi value where null
-      // information has not been propagated down. Ideally this case would never happen as it should
-      // be possible to replace the receiver by the null constant in this case.
-      //
-      // Since the receiver is known to be null, no argument information should be propagated to the
-      // callees, so we return bottom here.
-      return MethodState.bottom();
+              singleTarget,
+              existingMethodStateForBounds.asMonomorphicOrBottom(),
+              dynamicReceiverType);
+      return ConcretePolymorphicMethodState.create(bounds, methodStateForBounds);
     }
 
-    ProgramMethod singleTarget = invoke.lookupSingleProgramTarget(appView, context);
-    DynamicTypeWithUpperBound bounds =
-        computeBoundsForPolymorphicMethodState(resolvedMethod, singleTarget, dynamicReceiverType);
-    MethodState existingMethodStateForBounds =
-        existingMethodState.isPolymorphic()
-            ? existingMethodState.asPolymorphic().getMethodStateForBounds(bounds)
-            : MethodState.bottom();
+    private DynamicTypeWithUpperBound computeBoundsForPolymorphicMethodState(
+        ProgramMethod resolvedMethod,
+        ProgramMethod singleTarget,
+        DynamicTypeWithUpperBound dynamicReceiverType) {
+      DynamicTypeWithUpperBound bounds =
+          singleTarget != null
+              ? DynamicType.createExact(
+                  singleTarget.getHolderType().toTypeElement(appView).asClassType())
+              : dynamicReceiverType.withNullability(Nullability.maybeNull());
 
-    if (existingMethodStateForBounds.isPolymorphic()) {
-      assert false;
-      return MethodState.unknown();
-    }
+      // We intentionally drop the nullability for the type bounds. This increases the number of
+      // collisions in the polymorphic method states, which does not change the precision (since the
+      // nullability does not have any impact on the possible dispatch targets) and is good for
+      // state pruning.
+      assert bounds.getDynamicUpperBoundType().nullability().isMaybeNull();
 
-    // If we already don't know anything about the parameters for the given type bounds, then don't
-    // compute a method state.
-    if (existingMethodStateForBounds.isUnknown()) {
-      return MethodState.bottom();
-    }
+      // If the bounds are trivial (i.e., the upper bound is equal to the holder of the virtual root
+      // method), then widen the type bounds to 'unknown'.
+      DexMethod virtualRootMethod = getVirtualRootMethod(resolvedMethod);
+      if (virtualRootMethod == null) {
+        assert false : "Unexpected virtual method without root: " + resolvedMethod;
+        return bounds;
+      }
 
-    ConcreteMonomorphicMethodStateOrUnknown methodStateForBounds =
-        computeMonomorphicMethodState(
-            invoke,
-            resolvedMethod,
-            singleTarget,
-            abstractValueSupplier,
-            context,
-            existingMethodStateForBounds.asMonomorphicOrBottom(),
-            dynamicReceiverType);
-    return ConcretePolymorphicMethodState.create(bounds, methodStateForBounds);
-  }
-
-  private DynamicTypeWithUpperBound computeBoundsForPolymorphicMethodState(
-      ProgramMethod resolvedMethod,
-      ProgramMethod singleTarget,
-      DynamicTypeWithUpperBound dynamicReceiverType) {
-    DynamicTypeWithUpperBound bounds =
-        singleTarget != null
-            ? DynamicType.createExact(
-                singleTarget.getHolderType().toTypeElement(appView).asClassType())
-            : dynamicReceiverType.withNullability(Nullability.maybeNull());
-
-    // We intentionally drop the nullability for the type bounds. This increases the number of
-    // collisions in the polymorphic method states, which does not change the precision (since the
-    // nullability does not have any impact on the possible dispatch targets) and is good for state
-    // pruning.
-    assert bounds.getDynamicUpperBoundType().nullability().isMaybeNull();
-
-    // If the bounds are trivial (i.e., the upper bound is equal to the holder of the virtual root
-    // method), then widen the type bounds to 'unknown'.
-    DexMethod virtualRootMethod = getVirtualRootMethod(resolvedMethod);
-    if (virtualRootMethod == null) {
-      assert false : "Unexpected virtual method without root: " + resolvedMethod;
+      DynamicType trivialBounds =
+          DynamicType.create(
+              appView, virtualRootMethod.getHolderType().toTypeElement(appView).asClassType());
+      if (bounds.equals(trivialBounds)) {
+        return DynamicType.unknown();
+      }
       return bounds;
     }
 
-    DynamicType trivialBounds =
-        DynamicType.create(
-            appView, virtualRootMethod.getHolderType().toTypeElement(appView).asClassType());
-    if (bounds.equals(trivialBounds)) {
-      return DynamicType.unknown();
-    }
-    return bounds;
-  }
-
-  private ConcreteMonomorphicMethodStateOrUnknown computeMonomorphicMethodState(
-      InvokeMethod invoke,
-      ProgramMethod resolvedMethod,
-      ProgramMethod singleTarget,
-      AbstractValueSupplier abstractValueSupplier,
-      ProgramMethod context,
-      ConcreteMonomorphicMethodStateOrBottom existingMethodState) {
-    return computeMonomorphicMethodState(
-        invoke,
-        resolvedMethod,
-        singleTarget,
-        abstractValueSupplier,
-        context,
-        existingMethodState,
-        invoke.isInvokeMethodWithReceiver()
-            ? invoke.getFirstArgument().getDynamicType(appView)
-            : null);
-  }
-
-  @SuppressWarnings("UnusedVariable")
-  private ConcreteMonomorphicMethodStateOrUnknown computeMonomorphicMethodState(
-      InvokeMethod invoke,
-      ProgramMethod resolvedMethod,
-      ProgramMethod singleTarget,
-      AbstractValueSupplier abstractValueSupplier,
-      ProgramMethod context,
-      ConcreteMonomorphicMethodStateOrBottom existingMethodState,
-      DynamicType dynamicReceiverType) {
-    List<ValueState> parameterStates = new ArrayList<>(invoke.arguments().size());
-
-    MethodReprocessingCriteria methodReprocessingCriteria =
-        singleTarget != null
-            ? reprocessingCriteriaCollection.getReprocessingCriteria(singleTarget)
-            : MethodReprocessingCriteria.alwaysReprocess();
-
-    int argumentIndex = 0;
-    if (invoke.isInvokeMethodWithReceiver()) {
-      assert dynamicReceiverType != null;
-      parameterStates.add(
-          computeParameterStateForReceiver(
-              resolvedMethod,
-              dynamicReceiverType,
-              existingMethodState,
-              methodReprocessingCriteria.getParameterReprocessingCriteria(0)));
-      argumentIndex++;
+    private ConcreteMonomorphicMethodStateOrUnknown computeMonomorphicMethodState(
+        InvokeMethod invoke,
+        ProgramMethod resolvedMethod,
+        ProgramMethod singleTarget,
+        ConcreteMonomorphicMethodStateOrBottom existingMethodState) {
+      return computeMonomorphicMethodState(
+          invoke,
+          resolvedMethod,
+          singleTarget,
+          existingMethodState,
+          invoke.isInvokeMethodWithReceiver()
+              ? invoke.getFirstArgument().getDynamicType(appView)
+              : null);
     }
 
-    for (; argumentIndex < invoke.arguments().size(); argumentIndex++) {
-      parameterStates.add(
-          computeParameterStateForNonReceiver(
-              invoke,
-              singleTarget,
-              argumentIndex,
-              invoke.getArgument(argumentIndex),
-              abstractValueSupplier,
-              context,
-              existingMethodState));
+    @SuppressWarnings("UnusedVariable")
+    private ConcreteMonomorphicMethodStateOrUnknown computeMonomorphicMethodState(
+        InvokeMethod invoke,
+        ProgramMethod resolvedMethod,
+        ProgramMethod singleTarget,
+        ConcreteMonomorphicMethodStateOrBottom existingMethodState,
+        DynamicType dynamicReceiverType) {
+      List<ValueState> parameterStates = new ArrayList<>(invoke.arguments().size());
+
+      MethodReprocessingCriteria methodReprocessingCriteria =
+          singleTarget != null
+              ? reprocessingCriteriaCollection.getReprocessingCriteria(singleTarget)
+              : MethodReprocessingCriteria.alwaysReprocess();
+
+      int argumentIndex = 0;
+      if (invoke.isInvokeMethodWithReceiver()) {
+        assert dynamicReceiverType != null;
+        parameterStates.add(
+            computeParameterStateForReceiver(
+                resolvedMethod,
+                dynamicReceiverType,
+                existingMethodState,
+                methodReprocessingCriteria.getParameterReprocessingCriteria(0)));
+        argumentIndex++;
+      }
+
+      for (; argumentIndex < invoke.arguments().size(); argumentIndex++) {
+        parameterStates.add(
+            computeParameterStateForNonReceiver(
+                invoke,
+                singleTarget,
+                argumentIndex,
+                invoke.getArgument(argumentIndex),
+                existingMethodState));
+      }
+
+      // We simulate that the return value is used for methods with void return type. This ensures
+      // that we will widen the method state to unknown if/when all parameter states become unknown.
+      boolean isReturnValueUsed = invoke.getReturnType().isVoidType() || invoke.hasUsedOutValue();
+      return ConcreteMonomorphicMethodState.create(isReturnValueUsed, parameterStates);
     }
 
-    // We simulate that the return value is used for methods with void return type. This ensures
-    // that we will widen the method state to unknown if/when all parameter states become unknown.
-    boolean isReturnValueUsed = invoke.getReturnType().isVoidType() || invoke.hasUsedOutValue();
-    return ConcreteMonomorphicMethodState.create(isReturnValueUsed, parameterStates);
-  }
+    // For receivers there is not much point in trying to track an abstract value. Therefore we only
+    // track the dynamic type for receivers.
+    // TODO(b/190154391): Consider validating the above hypothesis by using
+    //  computeParameterStateForNonReceiver() for receivers.
+    private ValueState computeParameterStateForReceiver(
+        ProgramMethod resolvedMethod,
+        DynamicType dynamicReceiverType,
+        ConcreteMonomorphicMethodStateOrBottom existingMethodState,
+        ParameterReprocessingCriteria parameterReprocessingCriteria) {
+      // Don't compute a state for this parameter if the stored state is already unknown.
+      if (existingMethodState.isMonomorphic()
+          && existingMethodState.asMonomorphic().getParameterState(0).isUnknown()) {
+        return ValueState.unknown();
+      }
 
-  // For receivers there is not much point in trying to track an abstract value. Therefore we only
-  // track the dynamic type for receivers.
-  // TODO(b/190154391): Consider validating the above hypothesis by using
-  //  computeParameterStateForNonReceiver() for receivers.
-  private ValueState computeParameterStateForReceiver(
-      ProgramMethod resolvedMethod,
-      DynamicType dynamicReceiverType,
-      ConcreteMonomorphicMethodStateOrBottom existingMethodState,
-      ParameterReprocessingCriteria parameterReprocessingCriteria) {
-    // Don't compute a state for this parameter if the stored state is already unknown.
-    if (existingMethodState.isMonomorphic()
-        && existingMethodState.asMonomorphic().getParameterState(0).isUnknown()) {
-      return ValueState.unknown();
+      // For receivers we only track the dynamic type. Therefore, if there is no need to track the
+      // dynamic type of the receiver of the targeted method, then just return unknown.
+      if (!parameterReprocessingCriteria.shouldReprocessDueToDynamicType()) {
+        return ValueState.unknown();
+      }
+
+      DynamicType widenedDynamicReceiverType =
+          WideningUtils.widenDynamicReceiverType(appView, resolvedMethod, dynamicReceiverType);
+      return widenedDynamicReceiverType.isUnknown()
+          ? ValueState.unknown()
+          : new ConcreteReceiverValueState(dynamicReceiverType);
     }
 
-    // For receivers we only track the dynamic type. Therefore, if there is no need to track the
-    // dynamic type of the receiver of the targeted method, then just return unknown.
-    if (!parameterReprocessingCriteria.shouldReprocessDueToDynamicType()) {
-      return ValueState.unknown();
+    @SuppressWarnings("UnusedVariable")
+    private ValueState computeParameterStateForNonReceiver(
+        InvokeMethod invoke,
+        ProgramMethod singleTarget,
+        int argumentIndex,
+        Value argument,
+        ConcreteMonomorphicMethodStateOrBottom existingMethodState) {
+      ValueState modeledState =
+          modeling.modelParameterStateForArgumentToFunction(
+              invoke, singleTarget, argumentIndex, argument, context);
+      if (modeledState != null) {
+        return modeledState;
+      }
+
+      // Don't compute a state for this parameter if the stored state is already unknown.
+      if (existingMethodState.isMonomorphic()
+          && existingMethodState.asMonomorphic().getParameterState(argumentIndex).isUnknown()) {
+        return ValueState.unknown();
+      }
+
+      DexType parameterType =
+          invoke.getInvokedMethod().getArgumentType(argumentIndex, invoke.isInvokeStatic());
+
+      // If the value is an argument of the enclosing method, then clearly we have no information
+      // about its abstract value. Instead of treating this as having an unknown runtime value, we
+      // instead record a flow constraint that specifies that all values that flow into the
+      // parameter of this enclosing method also flows into the corresponding parameter of the
+      // methods potentially called from this invoke instruction.
+      NonEmptyValueState inFlowState = computeInFlowState(parameterType, argument, context);
+      if (inFlowState != null) {
+        return inFlowState;
+      }
+
+      // Only track the nullability for array types.
+      if (parameterType.isArrayType()) {
+        Nullability nullability = argument.getType().nullability();
+        return ConcreteArrayTypeValueState.create(nullability);
+      }
+
+      AbstractValue abstractValue = abstractValueSupplier.getAbstractValue(argument);
+
+      // For class types, we track both the abstract value and the dynamic type. If both are
+      // unknown, then use UnknownParameterState.
+      if (parameterType.isClassType()) {
+        DynamicType dynamicType = argument.getDynamicType(appView);
+        DynamicType widenedDynamicType =
+            WideningUtils.widenDynamicNonReceiverType(appView, dynamicType, parameterType);
+        return ConcreteClassTypeValueState.create(abstractValue, widenedDynamicType);
+      } else {
+        // For primitive types, we only track the abstract value, thus if the abstract value is
+        // unknown, we use UnknownParameterState.
+        assert parameterType.isPrimitiveType();
+        return ConcretePrimitiveTypeValueState.create(abstractValue);
+      }
     }
 
-    DynamicType widenedDynamicReceiverType =
-        WideningUtils.widenDynamicReceiverType(appView, resolvedMethod, dynamicReceiverType);
-    return widenedDynamicReceiverType.isUnknown()
-        ? ValueState.unknown()
-        : new ConcreteReceiverValueState(dynamicReceiverType);
-  }
+    @SuppressWarnings("ReferenceEquality")
+    private DexMethod getRepresentative(InvokeMethod invoke, ProgramMethod resolvedMethod) {
+      if (resolvedMethod.getDefinition().belongsToDirectPool()) {
+        return resolvedMethod.getReference();
+      }
 
-  @SuppressWarnings("UnusedVariable")
-  private ValueState computeParameterStateForNonReceiver(
-      InvokeMethod invoke,
-      ProgramMethod singleTarget,
-      int argumentIndex,
-      Value argument,
-      AbstractValueSupplier abstractValueSupplier,
-      ProgramMethod context,
-      ConcreteMonomorphicMethodStateOrBottom existingMethodState) {
-    ValueState modeledState =
-        modeling.modelParameterStateForArgumentToFunction(
-            invoke, singleTarget, argumentIndex, argument, context);
-    if (modeledState != null) {
-      return modeledState;
+      if (isMonomorphicVirtualMethod(resolvedMethod)) {
+        return resolvedMethod.getReference();
+      }
+
+      if (invoke.isInvokeInterface()) {
+        assert !isMonomorphicVirtualMethod(resolvedMethod);
+        return getVirtualRootMethod(resolvedMethod);
+      }
+
+      assert invoke.isInvokeSuper() || invoke.isInvokeVirtual();
+
+      DexMethod rootMethod = getVirtualRootMethod(resolvedMethod);
+      assert rootMethod != null;
+      assert !isMonomorphicVirtualMethod(resolvedMethod)
+          || rootMethod == resolvedMethod.getReference();
+      return rootMethod;
     }
 
-    // Don't compute a state for this parameter if the stored state is already unknown.
-    if (existingMethodState.isMonomorphic()
-        && existingMethodState.asMonomorphic().getParameterState(argumentIndex).isUnknown()) {
-      return ValueState.unknown();
+    private boolean shouldUsePolymorphicMethodState(
+        InvokeMethod invoke, ProgramMethod resolvedMethod) {
+      return !resolvedMethod.getDefinition().belongsToDirectPool()
+          && !isMonomorphicVirtualMethod(getRepresentative(invoke, resolvedMethod));
     }
 
-    DexType parameterType =
-        invoke.getInvokedMethod().getArgumentType(argumentIndex, invoke.isInvokeStatic());
-
-    // If the value is an argument of the enclosing method, then clearly we have no information
-    // about its abstract value. Instead of treating this as having an unknown runtime value, we
-    // instead record a flow constraint that specifies that all values that flow into the parameter
-    // of this enclosing method also flows into the corresponding parameter of the methods
-    // potentially called from this invoke instruction.
-    NonEmptyValueState inFlowState = computeInFlowState(parameterType, argument, context);
-    if (inFlowState != null) {
-      return inFlowState;
-    }
-
-    // Only track the nullability for array types.
-    if (parameterType.isArrayType()) {
-      Nullability nullability = argument.getType().nullability();
-      return ConcreteArrayTypeValueState.create(nullability);
-    }
-
-    AbstractValue abstractValue = abstractValueSupplier.getAbstractValue(argument);
-
-    // For class types, we track both the abstract value and the dynamic type. If both are unknown,
-    // then use UnknownParameterState.
-    if (parameterType.isClassType()) {
-      DynamicType dynamicType = argument.getDynamicType(appView);
-      DynamicType widenedDynamicType =
-          WideningUtils.widenDynamicNonReceiverType(appView, dynamicType, parameterType);
-      return ConcreteClassTypeValueState.create(abstractValue, widenedDynamicType);
-    } else {
-      // For primitive types, we only track the abstract value, thus if the abstract value is
-      // unknown,
-      // we use UnknownParameterState.
-      assert parameterType.isPrimitiveType();
-      return ConcretePrimitiveTypeValueState.create(abstractValue);
-    }
-  }
-
-  @SuppressWarnings("ReferenceEquality")
-  private DexMethod getRepresentative(InvokeMethod invoke, ProgramMethod resolvedMethod) {
-    if (resolvedMethod.getDefinition().belongsToDirectPool()) {
-      return resolvedMethod.getReference();
-    }
-
-    if (isMonomorphicVirtualMethod(resolvedMethod)) {
-      return resolvedMethod.getReference();
-    }
-
-    if (invoke.isInvokeInterface()) {
-      assert !isMonomorphicVirtualMethod(resolvedMethod);
-      return getVirtualRootMethod(resolvedMethod);
-    }
-
-    assert invoke.isInvokeSuper() || invoke.isInvokeVirtual();
-
-    DexMethod rootMethod = getVirtualRootMethod(resolvedMethod);
-    assert rootMethod != null;
-    assert !isMonomorphicVirtualMethod(resolvedMethod)
-        || rootMethod == resolvedMethod.getReference();
-    return rootMethod;
-  }
-
-  private boolean shouldUsePolymorphicMethodState(
-      InvokeMethod invoke, ProgramMethod resolvedMethod) {
-    return !resolvedMethod.getDefinition().belongsToDirectPool()
-        && !isMonomorphicVirtualMethod(getRepresentative(invoke, resolvedMethod));
-  }
-
-  private void scan(InvokeCustom invoke) {
-    // If the bootstrap method is program declared it will be called. The call is with runtime
-    // provided arguments so ensure that the argument information is unknown.
-    DexMethodHandle bootstrapMethod = invoke.getCallSite().bootstrapMethod;
-    SingleResolutionResult<?> resolution =
-        appView
-            .appInfo()
-            .resolveMethodLegacy(bootstrapMethod.asMethod(), bootstrapMethod.isInterface)
-            .asSingleResolution();
-    if (resolution != null && resolution.getResolvedHolder().isProgramClass()) {
-      methodStates.set(resolution.getResolvedProgramMethod(), UnknownMethodState.get());
+    private void scanInvokeCustom(InvokeCustom invoke) {
+      // If the bootstrap method is program declared it will be called. The call is with runtime
+      // provided arguments so ensure that the argument information is unknown.
+      DexMethodHandle bootstrapMethod = invoke.getCallSite().bootstrapMethod;
+      SingleResolutionResult<?> resolution =
+          appView
+              .appInfo()
+              .resolveMethodLegacy(bootstrapMethod.asMethod(), bootstrapMethod.isInterface)
+              .asSingleResolution();
+      if (resolution != null && resolution.getResolvedHolder().isProgramClass()) {
+        methodStates.set(resolution.getResolvedProgramMethod(), UnknownMethodState.get());
+      }
     }
   }
 
diff --git a/src/main/java/com/android/tools/r8/optimize/compose/ArgumentPropagatorCodeScannerForComposableFunctions.java b/src/main/java/com/android/tools/r8/optimize/compose/ArgumentPropagatorCodeScannerForComposableFunctions.java
index 174e8ba..2116f3a 100644
--- a/src/main/java/com/android/tools/r8/optimize/compose/ArgumentPropagatorCodeScannerForComposableFunctions.java
+++ b/src/main/java/com/android/tools/r8/optimize/compose/ArgumentPropagatorCodeScannerForComposableFunctions.java
@@ -7,6 +7,7 @@
 import com.android.tools.r8.graph.DexType;
 import com.android.tools.r8.graph.ProgramMethod;
 import com.android.tools.r8.ir.code.AbstractValueSupplier;
+import com.android.tools.r8.ir.code.IRCode;
 import com.android.tools.r8.ir.code.InvokeMethod;
 import com.android.tools.r8.optimize.argumentpropagation.ArgumentPropagatorCodeScanner;
 import com.android.tools.r8.optimize.argumentpropagation.codescanner.MethodParameter;
@@ -25,16 +26,12 @@
   }
 
   @Override
-  protected void addTemporaryMethodState(
-      InvokeMethod invoke,
-      ProgramMethod resolvedMethod,
+  public void scan(
+      ProgramMethod method,
+      IRCode code,
       AbstractValueSupplier abstractValueSupplier,
-      ProgramMethod context,
       Timing timing) {
-    ComposableCallGraphNode node = callGraph.getNodes().get(resolvedMethod);
-    if (node != null && node.isComposable()) {
-      super.addTemporaryMethodState(invoke, resolvedMethod, abstractValueSupplier, context, timing);
-    }
+    new CodeScanner(abstractValueSupplier, code, method).scan(timing);
   }
 
   @Override
@@ -43,4 +40,21 @@
     // We haven't defined the virtual root mapping, so we can't tell.
     return false;
   }
+
+  private class CodeScanner extends ArgumentPropagatorCodeScanner.CodeScanner {
+
+    protected CodeScanner(
+        AbstractValueSupplier abstractValueSupplier, IRCode code, ProgramMethod method) {
+      super(abstractValueSupplier, code, method);
+    }
+
+    @Override
+    protected void addTemporaryMethodState(
+        InvokeMethod invoke, ProgramMethod resolvedMethod, Timing timing) {
+      ComposableCallGraphNode node = callGraph.getNodes().get(resolvedMethod);
+      if (node != null && node.isComposable()) {
+        super.addTemporaryMethodState(invoke, resolvedMethod, timing);
+      }
+    }
+  }
 }