Avoid unnecessary state copying during primary optimization pass

Change-Id: I277e71ce30e46f190feb2af1c7094dbca01a57bb
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/BottomArrayTypeParameterState.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/BottomArrayTypeParameterState.java
index 666c87b..3ba87ce 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/BottomArrayTypeParameterState.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/BottomArrayTypeParameterState.java
@@ -25,6 +25,7 @@
       AppView<AppInfoWithLiveness> appView,
       ParameterState parameterState,
       DexType parameterType,
+      StateCloner cloner,
       Action onChangedAction) {
     if (parameterState.isBottom()) {
       return this;
@@ -36,6 +37,9 @@
     assert parameterState.asConcrete().isReferenceParameter();
     ConcreteReferenceTypeParameterState concreteParameterState =
         parameterState.asConcrete().asReferenceParameter();
+    if (concreteParameterState.isArrayParameter()) {
+      return cloner.mutableCopy(concreteParameterState);
+    }
     Nullability nullability = concreteParameterState.getNullability();
     if (nullability.isMaybeNull()) {
       return unknown();
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/BottomClassTypeParameterState.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/BottomClassTypeParameterState.java
index d1df0ac..2ae85ed 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/BottomClassTypeParameterState.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/BottomClassTypeParameterState.java
@@ -27,6 +27,7 @@
       AppView<AppInfoWithLiveness> appView,
       ParameterState parameterState,
       DexType parameterType,
+      StateCloner cloner,
       Action onChangedAction) {
     if (parameterState.isBottom()) {
       return this;
@@ -42,6 +43,9 @@
     DynamicType dynamicType = concreteParameterState.getDynamicType();
     DynamicType widenedDynamicType =
         WideningUtils.widenDynamicNonReceiverType(appView, dynamicType, parameterType);
+    if (concreteParameterState.isClassParameter() && !widenedDynamicType.isUnknown()) {
+      return cloner.mutableCopy(concreteParameterState);
+    }
     return abstractValue.isUnknown() && widenedDynamicType.isUnknown()
         ? unknown()
         : new ConcreteClassTypeParameterState(
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/BottomMethodState.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/BottomMethodState.java
index f5f117e..b7564eb 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/BottomMethodState.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/BottomMethodState.java
@@ -44,7 +44,8 @@
   public MethodState mutableJoin(
       AppView<AppInfoWithLiveness> appView,
       DexMethodSignature methodSignature,
-      MethodState methodState) {
+      MethodState methodState,
+      StateCloner cloner) {
     return methodState.mutableCopy();
   }
 
@@ -52,7 +53,8 @@
   public MethodState mutableJoin(
       AppView<AppInfoWithLiveness> appView,
       DexMethodSignature methodSignature,
-      Function<MethodState, MethodState> methodStateSupplier) {
-    return mutableJoin(appView, methodSignature, methodStateSupplier.apply(this));
+      Function<MethodState, MethodState> methodStateSupplier,
+      StateCloner cloner) {
+    return mutableJoin(appView, methodSignature, methodStateSupplier.apply(this), cloner);
   }
 }
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/BottomPrimitiveTypeParameterState.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/BottomPrimitiveTypeParameterState.java
index fc88145..729938d 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/BottomPrimitiveTypeParameterState.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/BottomPrimitiveTypeParameterState.java
@@ -25,6 +25,7 @@
       AppView<AppInfoWithLiveness> appView,
       ParameterState parameterState,
       DexType parameterType,
+      StateCloner cloner,
       Action onChangedAction) {
     if (parameterState.isBottom()) {
       assert parameterState == bottomPrimitiveTypeParameter();
@@ -35,6 +36,6 @@
     }
     assert parameterState.isConcrete();
     assert parameterState.asConcrete().isPrimitiveParameter();
-    return parameterState.mutableCopy();
+    return cloner.mutableCopy(parameterState);
   }
 }
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/BottomReceiverParameterState.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/BottomReceiverParameterState.java
index 3a94858..aa331f8 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/BottomReceiverParameterState.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/BottomReceiverParameterState.java
@@ -25,6 +25,7 @@
       AppView<AppInfoWithLiveness> appView,
       ParameterState parameterState,
       DexType parameterType,
+      StateCloner cloner,
       Action onChangedAction) {
     if (parameterState.isBottom()) {
       return this;
@@ -36,6 +37,9 @@
     assert parameterState.asConcrete().isReferenceParameter();
     ConcreteReferenceTypeParameterState concreteParameterState =
         parameterState.asConcrete().asReferenceParameter();
+    if (concreteParameterState.isReceiverParameter()) {
+      return cloner.mutableCopy(concreteParameterState);
+    }
     DynamicType dynamicType = concreteParameterState.getDynamicType();
     if (dynamicType.isUnknown()) {
       return unknown();
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/ConcreteMethodState.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/ConcreteMethodState.java
index ecd3e08..613d864 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/ConcreteMethodState.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/ConcreteMethodState.java
@@ -25,33 +25,38 @@
   public MethodState mutableJoin(
       AppView<AppInfoWithLiveness> appView,
       DexMethodSignature methodSignature,
-      MethodState methodState) {
+      MethodState methodState,
+      StateCloner cloner) {
     if (methodState.isBottom()) {
       return this;
     }
     if (methodState.isUnknown()) {
       return methodState;
     }
-    return mutableJoin(appView, methodSignature, methodState.asConcrete());
+    return mutableJoin(appView, methodSignature, methodState.asConcrete(), cloner);
   }
 
   @Override
   public MethodState mutableJoin(
       AppView<AppInfoWithLiveness> appView,
       DexMethodSignature methodSignature,
-      Function<MethodState, MethodState> methodStateSupplier) {
-    return mutableJoin(appView, methodSignature, methodStateSupplier.apply(this));
+      Function<MethodState, MethodState> methodStateSupplier,
+      StateCloner cloner) {
+    return mutableJoin(appView, methodSignature, methodStateSupplier.apply(this), cloner);
   }
 
   private MethodState mutableJoin(
       AppView<AppInfoWithLiveness> appView,
       DexMethodSignature methodSignature,
-      ConcreteMethodState methodState) {
+      ConcreteMethodState methodState,
+      StateCloner cloner) {
     if (isMonomorphic() && methodState.isMonomorphic()) {
-      return asMonomorphic().mutableJoin(appView, methodSignature, methodState.asMonomorphic());
+      return asMonomorphic()
+          .mutableJoin(appView, methodSignature, methodState.asMonomorphic(), cloner);
     }
     if (isPolymorphic() && methodState.isPolymorphic()) {
-      return asPolymorphic().mutableJoin(appView, methodSignature, methodState.asPolymorphic());
+      return asPolymorphic()
+          .mutableJoin(appView, methodSignature, methodState.asPolymorphic(), cloner);
     }
     assert false;
     return unknown();
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/ConcreteMonomorphicMethodState.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/ConcreteMonomorphicMethodState.java
index b49e48f..c98c383 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/ConcreteMonomorphicMethodState.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/ConcreteMonomorphicMethodState.java
@@ -46,7 +46,8 @@
   public ConcreteMonomorphicMethodStateOrUnknown mutableJoin(
       AppView<AppInfoWithLiveness> appView,
       DexMethodSignature methodSignature,
-      ConcreteMonomorphicMethodState methodState) {
+      ConcreteMonomorphicMethodState methodState,
+      StateCloner cloner) {
     if (size() != methodState.size()) {
       assert false;
       return unknown();
@@ -59,7 +60,7 @@
       ParameterState otherParameterState = methodState.parameterStates.get(0);
       DexType parameterType = null;
       parameterStates.set(
-          0, parameterState.mutableJoin(appView, otherParameterState, parameterType));
+          0, parameterState.mutableJoin(appView, otherParameterState, parameterType, cloner));
       argumentIndex++;
     }
 
@@ -68,7 +69,8 @@
       ParameterState otherParameterState = methodState.parameterStates.get(argumentIndex);
       DexType parameterType = methodSignature.getParameter(parameterIndex);
       parameterStates.set(
-          argumentIndex, parameterState.mutableJoin(appView, otherParameterState, parameterType));
+          argumentIndex,
+          parameterState.mutableJoin(appView, otherParameterState, parameterType, cloner));
       assert !parameterStates.get(argumentIndex).isConcrete()
           || !parameterStates.get(argumentIndex).asConcrete().isReceiverParameter();
     }
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 fb9451f..ed3873f 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
@@ -107,6 +107,7 @@
       AppView<AppInfoWithLiveness> appView,
       ParameterState parameterState,
       DexType parameterType,
+      StateCloner cloner,
       Action onChangedAction) {
     if (parameterState.isBottom()) {
       return this;
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/ConcretePolymorphicMethodState.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/ConcretePolymorphicMethodState.java
index c4a1feb..d17ff72 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/ConcretePolymorphicMethodState.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/ConcretePolymorphicMethodState.java
@@ -45,7 +45,8 @@
       AppView<AppInfoWithLiveness> appView,
       DexMethodSignature methodSignature,
       DynamicType bounds,
-      ConcreteMonomorphicMethodStateOrUnknown methodState) {
+      ConcreteMonomorphicMethodStateOrUnknown methodState,
+      StateCloner cloner) {
     assert !isEffectivelyBottom();
     assert !isEffectivelyUnknown();
     if (methodState.isUnknown()) {
@@ -58,7 +59,8 @@
     } else {
       assert methodState.isMonomorphic();
       ConcreteMonomorphicMethodStateOrUnknown newMethodStateForBounds =
-          joinInner(appView, methodSignature, receiverBoundsToState.get(bounds), methodState);
+          joinInner(
+              appView, methodSignature, receiverBoundsToState.get(bounds), methodState, cloner);
       if (bounds.isUnknown() && newMethodStateForBounds.isUnknown()) {
         return unknown();
       } else {
@@ -68,19 +70,23 @@
     }
   }
 
+  @SuppressWarnings("unchecked")
   private static ConcreteMonomorphicMethodStateOrUnknown joinInner(
       AppView<AppInfoWithLiveness> appView,
       DexMethodSignature methodSignature,
       ConcreteMonomorphicMethodStateOrUnknown methodState,
-      ConcreteMonomorphicMethodStateOrUnknown other) {
+      ConcreteMonomorphicMethodStateOrUnknown other,
+      StateCloner cloner) {
     if (methodState == null) {
-      return other.mutableCopy();
+      return (ConcreteMonomorphicMethodStateOrUnknown) cloner.mutableCopy(other);
     }
     if (methodState.isUnknown() || other.isUnknown()) {
       return unknown();
     }
     assert methodState.isMonomorphic();
-    return methodState.asMonomorphic().mutableJoin(appView, methodSignature, other.asMonomorphic());
+    return methodState
+        .asMonomorphic()
+        .mutableJoin(appView, methodSignature, other.asMonomorphic(), cloner);
   }
 
   public void forEach(
@@ -118,7 +124,8 @@
   public MethodState mutableCopyWithRewrittenBounds(
       AppView<AppInfoWithLiveness> appView,
       Function<DynamicType, DynamicType> boundsRewriter,
-      DexMethodSignature methodSignature) {
+      DexMethodSignature methodSignature,
+      StateCloner cloner) {
     assert !isEffectivelyBottom();
     assert !isEffectivelyUnknown();
     Map<DynamicType, ConcreteMonomorphicMethodStateOrUnknown> rewrittenReceiverBoundsToState =
@@ -132,7 +139,8 @@
       ConcreteMonomorphicMethodStateOrUnknown existingMethodStateForBounds =
           rewrittenReceiverBoundsToState.get(rewrittenBounds);
       ConcreteMonomorphicMethodStateOrUnknown newMethodStateForBounds =
-          joinInner(appView, methodSignature, existingMethodStateForBounds, entry.getValue());
+          joinInner(
+              appView, methodSignature, existingMethodStateForBounds, entry.getValue(), cloner);
       if (rewrittenBounds.isUnknown() && newMethodStateForBounds.isUnknown()) {
         return unknown();
       }
@@ -146,7 +154,8 @@
   public MethodState mutableJoin(
       AppView<AppInfoWithLiveness> appView,
       DexMethodSignature methodSignature,
-      ConcretePolymorphicMethodState methodState) {
+      ConcretePolymorphicMethodState methodState,
+      StateCloner cloner) {
     assert !isEffectivelyBottom();
     assert !isEffectivelyUnknown();
     assert !methodState.isEffectivelyBottom();
@@ -154,7 +163,7 @@
     for (Entry<DynamicType, ConcreteMonomorphicMethodStateOrUnknown> entry :
         methodState.receiverBoundsToState.entrySet()) {
       ConcretePolymorphicMethodStateOrUnknown result =
-          add(appView, methodSignature, entry.getKey(), entry.getValue());
+          add(appView, methodSignature, entry.getKey(), entry.getValue(), cloner);
       if (result.isUnknown()) {
         return result;
       }
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/MethodState.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/MethodState.java
index 0de2d3e..b9df609 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/MethodState.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/MethodState.java
@@ -44,10 +44,12 @@
   MethodState mutableJoin(
       AppView<AppInfoWithLiveness> appView,
       DexMethodSignature methodSignature,
-      MethodState methodState);
+      MethodState methodState,
+      StateCloner cloner);
 
   MethodState mutableJoin(
       AppView<AppInfoWithLiveness> appView,
       DexMethodSignature methodSignature,
-      Function<MethodState, MethodState> methodStateSupplier);
+      Function<MethodState, MethodState> methodStateSupplier,
+      StateCloner cloner);
 }
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/MethodStateCollection.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/MethodStateCollection.java
index 50f02c9..856ff9d 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/MethodStateCollection.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/MethodStateCollection.java
@@ -45,7 +45,8 @@
               newMethodState = methodState.mutableCopy();
             } else {
               newMethodState =
-                  existingMethodState.mutableJoin(appView, getSignature(method), methodState);
+                  existingMethodState.mutableJoin(
+                      appView, getSignature(method), methodState, StateCloner.getCloner());
             }
             assert !newMethodState.isBottom();
             return newMethodState;
@@ -73,7 +74,8 @@
           assert !existingMethodState.isBottom();
           timing.begin("Join temporary method state");
           MethodState joinResult =
-              existingMethodState.mutableJoin(appView, getSignature(method), methodStateSupplier);
+              existingMethodState.mutableJoin(
+                  appView, getSignature(method), methodStateSupplier, StateCloner.getIdentity());
           assert !joinResult.isBottom();
           timing.end();
           return joinResult;
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/ParameterState.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/ParameterState.java
index ad7b97c..d70fdfd 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/ParameterState.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/ParameterState.java
@@ -57,13 +57,17 @@
   public abstract ParameterState mutableCopy();
 
   public final ParameterState mutableJoin(
-      AppView<AppInfoWithLiveness> appView, ParameterState parameterState, DexType parameterType) {
-    return mutableJoin(appView, parameterState, parameterType, Action.empty());
+      AppView<AppInfoWithLiveness> appView,
+      ParameterState parameterState,
+      DexType parameterType,
+      StateCloner cloner) {
+    return mutableJoin(appView, parameterState, parameterType, cloner, Action.empty());
   }
 
   public abstract ParameterState mutableJoin(
       AppView<AppInfoWithLiveness> appView,
       ParameterState parameterState,
       DexType parameterType,
+      StateCloner cloner,
       Action onChangedAction);
 }
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/StateCloner.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/StateCloner.java
new file mode 100644
index 0000000..6a2e941
--- /dev/null
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/StateCloner.java
@@ -0,0 +1,60 @@
+// Copyright (c) 2021, 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;
+
+/**
+ * A strategy for cloning method and parameter states.
+ *
+ * <p>During the primary optimization pass, for each invoke we compute a fresh method state and join
+ * the state into the existing method state for the resolved method. Since the added method state is
+ * completely fresh and not stored anywhere else, we can avoid copying the method state when we join
+ * it into the existing method state. This is achieved by using the {@link #getIdentity()} cloner
+ * below.
+ *
+ * <p>When we later propagate argument information for virtual methods to their overrides, we join
+ * method states from one virtual method into the state for another virtual method. Therefore, it is
+ * important to copy the method state during the join, which is achieved using the {@link
+ * #getCloner()} cloner.
+ */
+public abstract class StateCloner {
+
+  private static StateCloner CLONER =
+      new StateCloner() {
+        @Override
+        public MethodState mutableCopy(MethodState methodState) {
+          return methodState.mutableCopy();
+        }
+
+        @Override
+        public ParameterState mutableCopy(ParameterState parameterState) {
+          return parameterState.mutableCopy();
+        }
+      };
+
+  private static StateCloner IDENTITY =
+      new StateCloner() {
+        @Override
+        public MethodState mutableCopy(MethodState methodState) {
+          return methodState;
+        }
+
+        @Override
+        public ParameterState mutableCopy(ParameterState parameterState) {
+          return parameterState;
+        }
+      };
+
+  public static StateCloner getCloner() {
+    return CLONER;
+  }
+
+  public static StateCloner getIdentity() {
+    return IDENTITY;
+  }
+
+  public abstract MethodState mutableCopy(MethodState methodState);
+
+  public abstract ParameterState mutableCopy(ParameterState parameterState);
+}
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/UnknownMethodState.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/UnknownMethodState.java
index b4bb41f..dd822b8 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/UnknownMethodState.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/UnknownMethodState.java
@@ -35,7 +35,8 @@
   public MethodState mutableJoin(
       AppView<AppInfoWithLiveness> appView,
       DexMethodSignature methodSignature,
-      MethodState methodState) {
+      MethodState methodState,
+      StateCloner cloner) {
     return this;
   }
 
@@ -43,7 +44,8 @@
   public MethodState mutableJoin(
       AppView<AppInfoWithLiveness> appView,
       DexMethodSignature methodSignature,
-      Function<MethodState, MethodState> methodStateSupplier) {
+      Function<MethodState, MethodState> methodStateSupplier,
+      StateCloner cloner) {
     return this;
   }
 }
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/UnknownParameterState.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/UnknownParameterState.java
index bcb5317..f70ae3d 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/UnknownParameterState.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/codescanner/UnknownParameterState.java
@@ -40,6 +40,7 @@
       AppView<AppInfoWithLiveness> appView,
       ParameterState parameterState,
       DexType parameterType,
+      StateCloner cloner,
       Action onChangedAction) {
     return this;
   }
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 98277f4..9b583c7 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
@@ -20,6 +20,7 @@
 import com.android.tools.r8.optimize.argumentpropagation.codescanner.MethodStateCollectionByReference;
 import com.android.tools.r8.optimize.argumentpropagation.codescanner.NonEmptyParameterState;
 import com.android.tools.r8.optimize.argumentpropagation.codescanner.ParameterState;
+import com.android.tools.r8.optimize.argumentpropagation.codescanner.StateCloner;
 import com.android.tools.r8.optimize.argumentpropagation.utils.BidirectedGraph;
 import com.android.tools.r8.shaking.AppInfoWithLiveness;
 import com.android.tools.r8.utils.Action;
@@ -322,7 +323,11 @@
       ParameterState oldParameterState = getState();
       ParameterState newParameterState =
           oldParameterState.mutableJoin(
-              appView, parameterStateToAdd, parameterType, onChangedAction);
+              appView,
+              parameterStateToAdd,
+              parameterType,
+              StateCloner.getCloner(),
+              onChangedAction);
       if (newParameterState != oldParameterState) {
         setState(newParameterState);
         onChangedAction.execute();
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/propagation/InterfaceMethodArgumentPropagator.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/propagation/InterfaceMethodArgumentPropagator.java
index 07c487f..e1e5f63 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/propagation/InterfaceMethodArgumentPropagator.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/propagation/InterfaceMethodArgumentPropagator.java
@@ -18,6 +18,7 @@
 import com.android.tools.r8.optimize.argumentpropagation.codescanner.MethodState;
 import com.android.tools.r8.optimize.argumentpropagation.codescanner.MethodStateCollectionByReference;
 import com.android.tools.r8.optimize.argumentpropagation.codescanner.MethodStateCollectionBySignature;
+import com.android.tools.r8.optimize.argumentpropagation.codescanner.StateCloner;
 import com.android.tools.r8.shaking.AppInfoWithLiveness;
 import java.util.Collection;
 import java.util.IdentityHashMap;
@@ -185,7 +186,8 @@
               }
               return null;
             },
-            resolvedMethod.getMethodSignature());
+            resolvedMethod.getMethodSignature(),
+            StateCloner.getCloner());
 
     // If the resolved method is a virtual method that does not override any methods and are not
     // overridden by any methods, then we use a monomorphic method state for it. Therefore, we
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/propagation/VirtualDispatchMethodArgumentPropagator.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/propagation/VirtualDispatchMethodArgumentPropagator.java
index 41cb0af..36525f5 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/propagation/VirtualDispatchMethodArgumentPropagator.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/propagation/VirtualDispatchMethodArgumentPropagator.java
@@ -23,6 +23,7 @@
 import com.android.tools.r8.optimize.argumentpropagation.codescanner.MethodState;
 import com.android.tools.r8.optimize.argumentpropagation.codescanner.MethodStateCollectionByReference;
 import com.android.tools.r8.optimize.argumentpropagation.codescanner.MethodStateCollectionBySignature;
+import com.android.tools.r8.optimize.argumentpropagation.codescanner.StateCloner;
 import com.android.tools.r8.optimize.argumentpropagation.codescanner.UnknownMethodState;
 import com.android.tools.r8.shaking.AppInfoWithLiveness;
 import java.util.Collection;
@@ -150,7 +151,9 @@
       if (!activeUntilLowerBound.isEmpty()) {
         DexMethodSignature methodSignature = method.getMethodSignature();
         for (MethodStateCollectionBySignature methodStates : activeUntilLowerBound.values()) {
-          methodState = methodState.mutableJoin(appView, methodSignature, methodStates.get(method));
+          methodState =
+              methodState.mutableJoin(
+                  appView, methodSignature, methodStates.get(method), StateCloner.getCloner());
         }
       }
       return methodState;