diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/ConcreteArrayTypeParameterState.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/ConcreteArrayTypeParameterState.java
index 97803d0..5f53e18 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/ConcreteArrayTypeParameterState.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/ConcreteArrayTypeParameterState.java
@@ -19,15 +19,15 @@
 
   private Nullability nullability;
 
-  public ConcreteArrayTypeParameterState(MethodParameter inParameter) {
-    this(Nullability.bottom(), SetUtils.newHashSet(inParameter));
+  public ConcreteArrayTypeParameterState(InFlow inFlow) {
+    this(Nullability.bottom(), SetUtils.newHashSet(inFlow));
   }
 
   public ConcreteArrayTypeParameterState(Nullability nullability) {
     this(nullability, Collections.emptySet());
   }
 
-  public ConcreteArrayTypeParameterState(Nullability nullability, Set<MethodParameter> inFlow) {
+  public ConcreteArrayTypeParameterState(Nullability nullability, Set<InFlow> inFlow) {
     super(inFlow);
     this.nullability = nullability;
     assert !isEffectivelyBottom() : "Must use BottomArrayTypeParameterState instead";
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/ConcreteClassTypeParameterState.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/ConcreteClassTypeParameterState.java
index 5c6d7a4..82b371b 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/ConcreteClassTypeParameterState.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/ConcreteClassTypeParameterState.java
@@ -21,8 +21,8 @@
   private AbstractValue abstractValue;
   private DynamicType dynamicType;
 
-  public ConcreteClassTypeParameterState(MethodParameter inParameter) {
-    this(AbstractValue.bottom(), DynamicType.bottom(), SetUtils.newHashSet(inParameter));
+  public ConcreteClassTypeParameterState(InFlow inFlow) {
+    this(AbstractValue.bottom(), DynamicType.bottom(), SetUtils.newHashSet(inFlow));
   }
 
   public ConcreteClassTypeParameterState(AbstractValue abstractValue, DynamicType dynamicType) {
@@ -30,7 +30,7 @@
   }
 
   public ConcreteClassTypeParameterState(
-      AbstractValue abstractValue, DynamicType dynamicType, Set<MethodParameter> inFlow) {
+      AbstractValue abstractValue, DynamicType dynamicType, Set<InFlow> inFlow) {
     super(inFlow);
     this.abstractValue = abstractValue;
     this.dynamicType = dynamicType;
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/ConcreteParameterState.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/ConcreteParameterState.java
index 4e95fed..5f94c04 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/ConcreteParameterState.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/ConcreteParameterState.java
@@ -21,9 +21,9 @@
     RECEIVER
   }
 
-  private Set<MethodParameter> inFlow;
+  private Set<InFlow> inFlow;
 
-  ConcreteParameterState(Set<MethodParameter> inFlow) {
+  ConcreteParameterState(Set<InFlow> inFlow) {
     this.inFlow = inFlow;
   }
 
@@ -33,9 +33,9 @@
     inFlow = Collections.emptySet();
   }
 
-  public Set<MethodParameter> copyInFlow() {
+  public Set<InFlow> copyInFlow() {
     if (inFlow.isEmpty()) {
-      assert inFlow == Collections.<MethodParameter>emptySet();
+      assert inFlow == Collections.<InFlow>emptySet();
       return inFlow;
     }
     return new HashSet<>(inFlow);
@@ -45,7 +45,7 @@
     return !inFlow.isEmpty();
   }
 
-  public Set<MethodParameter> getInFlow() {
+  public Set<InFlow> getInFlow() {
     assert inFlow.isEmpty() || inFlow instanceof HashSet<?>;
     return inFlow;
   }
@@ -135,7 +135,7 @@
       return false;
     }
     if (inFlow.isEmpty()) {
-      assert inFlow == Collections.<MethodParameter>emptySet();
+      assert inFlow == Collections.<InFlow>emptySet();
       inFlow = new HashSet<>();
     }
     return inFlow.addAll(parameterState.inFlow);
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/ConcretePrimitiveTypeParameterState.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/ConcretePrimitiveTypeParameterState.java
index 6a90616..5057f0f 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/ConcretePrimitiveTypeParameterState.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/ConcretePrimitiveTypeParameterState.java
@@ -21,16 +21,15 @@
     this(abstractValue, Collections.emptySet());
   }
 
-  public ConcretePrimitiveTypeParameterState(
-      AbstractValue abstractValue, Set<MethodParameter> inFlow) {
+  public ConcretePrimitiveTypeParameterState(AbstractValue abstractValue, Set<InFlow> inFlow) {
     super(inFlow);
     this.abstractValue = abstractValue;
     assert !isEffectivelyBottom() : "Must use BottomPrimitiveTypeParameterState instead";
     assert !isEffectivelyUnknown() : "Must use UnknownParameterState instead";
   }
 
-  public ConcretePrimitiveTypeParameterState(MethodParameter inParameter) {
-    this(AbstractValue.bottom(), SetUtils.newHashSet(inParameter));
+  public ConcretePrimitiveTypeParameterState(InFlow inFlow) {
+    this(AbstractValue.bottom(), SetUtils.newHashSet(inFlow));
   }
 
   @Override
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/ConcreteReceiverParameterState.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/ConcreteReceiverParameterState.java
index 5a321a3..442c6ed 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/ConcreteReceiverParameterState.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/ConcreteReceiverParameterState.java
@@ -22,7 +22,7 @@
     this(dynamicType, Collections.emptySet());
   }
 
-  public ConcreteReceiverParameterState(DynamicType dynamicType, Set<MethodParameter> inFlow) {
+  public ConcreteReceiverParameterState(DynamicType dynamicType, Set<InFlow> inFlow) {
     super(inFlow);
     this.dynamicType = dynamicType;
     assert !isEffectivelyBottom() : "Must use BottomReceiverParameterState instead";
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/ConcreteReferenceTypeParameterState.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/ConcreteReferenceTypeParameterState.java
index da2056a..6e94838 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/ConcreteReferenceTypeParameterState.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/ConcreteReferenceTypeParameterState.java
@@ -14,7 +14,7 @@
 
 public abstract class ConcreteReferenceTypeParameterState extends ConcreteParameterState {
 
-  ConcreteReferenceTypeParameterState(Set<MethodParameter> inFlow) {
+  ConcreteReferenceTypeParameterState(Set<InFlow> inFlow) {
     super(inFlow);
   }
 
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/InFlow.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/InFlow.java
new file mode 100644
index 0000000..b651045
--- /dev/null
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/InFlow.java
@@ -0,0 +1,15 @@
+// Copyright (c) 2024, the R8 project authors. Please see the AUTHORS file
+// for details. All rights reserved. Use of this source code is governed by a
+// BSD-style license that can be found in the LICENSE file.
+package com.android.tools.r8.optimize.argumentpropagation.codescanner;
+
+public interface InFlow {
+
+  default boolean isMethodParameter() {
+    return false;
+  }
+
+  default MethodParameter asMethodParameter() {
+    return null;
+  }
+}
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/MethodParameter.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/MethodParameter.java
index d2ba546..9bfeab7 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/MethodParameter.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/MethodParameter.java
@@ -7,7 +7,7 @@
 import com.android.tools.r8.graph.DexMethod;
 import java.util.Objects;
 
-public class MethodParameter {
+public class MethodParameter implements InFlow {
 
   private final DexMethod method;
   private final int index;
@@ -26,6 +26,16 @@
   }
 
   @Override
+  public boolean isMethodParameter() {
+    return true;
+  }
+
+  @Override
+  public MethodParameter asMethodParameter() {
+    return this;
+  }
+
+  @Override
   @SuppressWarnings({"EqualsGetClass", "ReferenceEquality"})
   public boolean equals(Object obj) {
     if (obj == null || getClass() != obj.getClass()) {
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/propagation/InParameterFlowPropagator.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/propagation/InParameterFlowPropagator.java
index ff0bd01..deae78e 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/propagation/InParameterFlowPropagator.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/propagation/InParameterFlowPropagator.java
@@ -7,6 +7,7 @@
 import static com.android.tools.r8.graph.DexProgramClass.asProgramClassOrNull;
 import static com.android.tools.r8.utils.MapUtils.ignoreKey;
 
+import com.android.tools.r8.errors.Unreachable;
 import com.android.tools.r8.graph.AppView;
 import com.android.tools.r8.graph.DexMethod;
 import com.android.tools.r8.graph.DexProgramClass;
@@ -16,6 +17,7 @@
 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.ConcreteParameterState;
+import com.android.tools.r8.optimize.argumentpropagation.codescanner.InFlow;
 import com.android.tools.r8.optimize.argumentpropagation.codescanner.MethodParameter;
 import com.android.tools.r8.optimize.argumentpropagation.codescanner.MethodState;
 import com.android.tools.r8.optimize.argumentpropagation.codescanner.MethodStateCollectionByReference;
@@ -26,6 +28,7 @@
 import com.android.tools.r8.shaking.AppInfoWithLiveness;
 import com.android.tools.r8.utils.Action;
 import com.android.tools.r8.utils.ThreadUtils;
+import com.android.tools.r8.utils.TraversalContinuation;
 import com.google.common.collect.Sets;
 import it.unimi.dsi.fastutil.ints.Int2ReferenceMap;
 import it.unimi.dsi.fastutil.ints.Int2ReferenceOpenHashMap;
@@ -208,41 +211,14 @@
       }
 
       ParameterNode node = getOrCreateParameterNode(method, parameterIndex, methodState);
-      for (MethodParameter inParameter : concreteParameterState.getInFlow()) {
-        ProgramMethod enclosingMethod = getEnclosingMethod(inParameter);
-        if (enclosingMethod == null) {
-          // This is a parameter of a single caller inlined method. Since this method has been
-          // pruned, the call from inside the method no longer exists, and we can therefore safely
-          // skip it.
-          assert converter
-              .getInliner()
-              .verifyIsPrunedDueToSingleCallerInlining(inParameter.getMethod());
-          continue;
+      for (InFlow inFlow : concreteParameterState.getInFlow()) {
+        if (inFlow.isMethodParameter()) {
+          if (addInFlow(inFlow.asMethodParameter(), node).shouldBreak()) {
+            break;
+          }
+        } else {
+          throw new Unreachable();
         }
-
-        MethodState enclosingMethodState = getMethodState(enclosingMethod);
-        if (enclosingMethodState.isBottom()) {
-          // The current method is called from a dead method; no need to propagate any information
-          // from the dead call site.
-          continue;
-        }
-
-        if (enclosingMethodState.isUnknown()) {
-          // The parameter depends on another parameter for which we don't know anything.
-          node.clearPredecessors();
-          node.setState(ParameterState.unknown());
-          break;
-        }
-
-        assert enclosingMethodState.isConcrete();
-        assert enclosingMethodState.asConcrete().isMonomorphic();
-
-        ParameterNode predecessor =
-            getOrCreateParameterNode(
-                enclosingMethod,
-                inParameter.getIndex(),
-                enclosingMethodState.asConcrete().asMonomorphic());
-        node.addPredecessor(predecessor);
       }
 
       if (!node.getState().isUnknown()) {
@@ -251,6 +227,42 @@
       }
     }
 
+    private TraversalContinuation<?, ?> addInFlow(MethodParameter inFlow, ParameterNode node) {
+      ProgramMethod enclosingMethod = getEnclosingMethod(inFlow);
+      if (enclosingMethod == null) {
+        // This is a parameter of a single caller inlined method. Since this method has been
+        // pruned, the call from inside the method no longer exists, and we can therefore safely
+        // skip it.
+        assert converter.getInliner().verifyIsPrunedDueToSingleCallerInlining(inFlow.getMethod());
+        return TraversalContinuation.doContinue();
+      }
+
+      MethodState enclosingMethodState = getMethodState(enclosingMethod);
+      if (enclosingMethodState.isBottom()) {
+        // The current method is called from a dead method; no need to propagate any information
+        // from the dead call site.
+        return TraversalContinuation.doContinue();
+      }
+
+      if (enclosingMethodState.isUnknown()) {
+        // The parameter depends on another parameter for which we don't know anything.
+        node.clearPredecessors();
+        node.setState(ParameterState.unknown());
+        return TraversalContinuation.doBreak();
+      }
+
+      assert enclosingMethodState.isConcrete();
+      assert enclosingMethodState.asConcrete().isMonomorphic();
+
+      ParameterNode predecessor =
+          getOrCreateParameterNode(
+              enclosingMethod,
+              inFlow.getIndex(),
+              enclosingMethodState.asConcrete().asMonomorphic());
+      node.addPredecessor(predecessor);
+      return TraversalContinuation.doContinue();
+    }
+
     private ParameterNode getOrCreateParameterNode(
         ProgramMethod method, int parameterIndex, ConcreteMonomorphicMethodState methodState) {
       Int2ReferenceMap<ParameterNode> parameterNodesForMethod =