Only fixup affected classes in argument propagation
Change-Id: I093d8ae4dcee436331a8b73a6138dad1e4051dec
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagator.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagator.java
index f81b135..45420b6 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagator.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagator.java
@@ -21,11 +21,13 @@
 import com.android.tools.r8.shaking.AppInfoWithLiveness;
 import com.android.tools.r8.utils.ThreadUtils;
 import com.android.tools.r8.utils.Timing;
+import com.google.common.collect.Sets;
 import java.util.Collection;
 import java.util.List;
 import java.util.Set;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
+import java.util.function.Consumer;
 
 /** Optimization that propagates information about arguments from call sites to method entries. */
 public class ArgumentPropagator {
@@ -139,8 +141,10 @@
 
     // Using the computed optimization info, build a graph lens that describes the mapping from
     // methods with constant parameters to methods with the constant parameters removed.
+    Set<DexProgramClass> affectedClasses = Sets.newConcurrentHashSet();
     ArgumentPropagatorGraphLens graphLens =
-        optimizeMethodParameters(stronglyConnectedProgramComponents, executorService);
+        optimizeMethodParameters(
+            stronglyConnectedProgramComponents, affectedClasses::add, executorService);
 
     // Find all the code objects that need reprocessing.
     new ArgumentPropagatorMethodReprocessingEnqueuer(appView)
@@ -148,7 +152,8 @@
 
     // Finally, apply the graph lens to the program (i.e., remove constant parameters from method
     // definitions).
-    new ArgumentPropagatorApplicationFixer(appView, graphLens).fixupApplication(executorService);
+    new ArgumentPropagatorApplicationFixer(appView, graphLens)
+        .fixupApplication(affectedClasses, executorService);
 
     timing.end();
   }
@@ -184,12 +189,15 @@
   /** Called by {@link IRConverter} to optimize method definitions. */
   private ArgumentPropagatorGraphLens optimizeMethodParameters(
       List<Set<DexProgramClass>> stronglyConnectedProgramComponents,
+      Consumer<DexProgramClass> affectedClassConsumer,
       ExecutorService executorService)
       throws ExecutionException {
     Collection<ArgumentPropagatorGraphLens.Builder> partialGraphLensBuilders =
         ThreadUtils.processItemsWithResults(
             stronglyConnectedProgramComponents,
-            classes -> new ArgumentPropagatorProgramOptimizer(appView).optimize(classes),
+            classes ->
+                new ArgumentPropagatorProgramOptimizer(appView)
+                    .optimize(classes, affectedClassConsumer),
             executorService);
 
     // Merge all the partial, disjoint graph lens builders into a single graph lens.
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagatorApplicationFixer.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagatorApplicationFixer.java
index fc74192..71e4970 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagatorApplicationFixer.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagatorApplicationFixer.java
@@ -15,6 +15,7 @@
 import com.android.tools.r8.ir.optimize.info.OptimizationFeedbackSimple;
 import com.android.tools.r8.shaking.AppInfoWithLiveness;
 import com.android.tools.r8.utils.ThreadUtils;
+import java.util.Set;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 
@@ -33,15 +34,19 @@
     this.graphLens = graphLens;
   }
 
-  public void fixupApplication(ExecutorService executorService) throws ExecutionException {
+  public void fixupApplication(
+      Set<DexProgramClass> affectedClasses, ExecutorService executorService)
+      throws ExecutionException {
     // If the graph lens is null, argument propagation did not lead to any parameter removals. In
     // this case there is no needed to fixup the program.
     if (graphLens == null) {
+      assert affectedClasses.isEmpty();
       return;
     }
 
-    // TODO(b/190154391): Do not naively visit all classes, when only few require changes.
-    ThreadUtils.processItems(appView.appInfo().classes(), this::fixupClass, executorService);
+    assert !affectedClasses.isEmpty();
+
+    ThreadUtils.processItems(affectedClasses, this::fixupClass, executorService);
     appView.setGraphLens(graphLens);
   }
 
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagatorGraphLens.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagatorGraphLens.java
index 77ec495..a00c016 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagatorGraphLens.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagatorGraphLens.java
@@ -83,13 +83,10 @@
 
     public Builder recordMove(
         DexMethod from, DexMethod to, ArgumentInfoCollection removedParametersForMethod) {
-      if (from != to) {
-        newMethodSignatures.put(from, to);
-        if (!removedParametersForMethod.isEmpty()) {
-          removedParameters.put(to, removedParametersForMethod);
-        }
-      } else {
-        assert removedParametersForMethod.isEmpty();
+      assert from != to;
+      newMethodSignatures.put(from, to);
+      if (!removedParametersForMethod.isEmpty()) {
+        removedParameters.put(to, removedParametersForMethod);
       }
       return this;
     }
diff --git a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagatorProgramOptimizer.java b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagatorProgramOptimizer.java
index b10ff8d..f469ddc 100644
--- a/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagatorProgramOptimizer.java
+++ b/src/main/java/com/android/tools/r8/optimize/argumentpropagation/ArgumentPropagatorProgramOptimizer.java
@@ -19,6 +19,7 @@
 import com.android.tools.r8.ir.optimize.info.CallSiteOptimizationInfo;
 import com.android.tools.r8.ir.optimize.info.ConcreteCallSiteOptimizationInfo;
 import com.android.tools.r8.shaking.AppInfoWithLiveness;
+import com.android.tools.r8.utils.BooleanBox;
 import com.android.tools.r8.utils.InternalOptions;
 import com.android.tools.r8.utils.Pair;
 import com.android.tools.r8.utils.collections.DexMethodSignatureSet;
@@ -31,6 +32,7 @@
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Set;
+import java.util.function.Consumer;
 import java.util.function.IntPredicate;
 
 public class ArgumentPropagatorProgramOptimizer {
@@ -68,7 +70,8 @@
   //  also enqueue the caller's callers for reprocessing. This would propagate the throwing
   //  information to all call sites.
   public ArgumentPropagatorGraphLens.Builder optimize(
-      Set<DexProgramClass> stronglyConnectedProgramClasses) {
+      Set<DexProgramClass> stronglyConnectedProgramClasses,
+      Consumer<DexProgramClass> affectedClassConsumer) {
     // First reserve pinned method signatures.
     reservePinnedMethodSignatures(stronglyConnectedProgramClasses);
 
@@ -83,7 +86,9 @@
     ArgumentPropagatorGraphLens.Builder partialGraphLensBuilder =
         ArgumentPropagatorGraphLens.builder(appView);
     for (DexProgramClass clazz : stronglyConnectedProgramClasses) {
-      visitClass(clazz, partialGraphLensBuilder);
+      if (visitClass(clazz, partialGraphLensBuilder)) {
+        affectedClassConsumer.accept(clazz);
+      }
     }
     return partialGraphLensBuilder;
   }
@@ -200,8 +205,10 @@
     return true;
   }
 
-  private void visitClass(
+  // Returns true if the class was changed as a result of argument propagation.
+  private boolean visitClass(
       DexProgramClass clazz, ArgumentPropagatorGraphLens.Builder partialGraphLensBuilder) {
+    BooleanBox affected = new BooleanBox();
     clazz.forEachProgramMethod(
         method -> {
           ArgumentInfoCollection removableParameters =
@@ -209,9 +216,13 @@
                   ? computeRemovableParametersFromDirectMethod(method)
                   : computeRemovableParametersFromVirtualMethod(method);
           DexMethod newMethodSignature = getNewMethodSignature(method, removableParameters);
-          partialGraphLensBuilder.recordMove(
-              method.getReference(), newMethodSignature, removableParameters);
+          if (newMethodSignature != method.getReference()) {
+            partialGraphLensBuilder.recordMove(
+                method.getReference(), newMethodSignature, removableParameters);
+            affected.set();
+          }
         });
+    return affected.get();
   }
 
   private DexMethod getNewMethodSignature(