Outliner: Don't use mutable keys in HashMap

The Outline type (used by Outliner to represent outlining candidates)
is mutable since it stores references to IR instructions belonging to
particular methods. As a result, the Outline type can change its mapping
identity when R8 optimizations determine that an instruction's out-value
can be set to null (which happens e.g. on all StringBuilder and
StringBuffer methods that are known to return `this`). When Outline is
used as the key type in a HashMap, the map must not stay around long
enough for the Outline keys to change their mapping identities in this
way. Change Outliner from keeping a Map<Outline, List<DexEncodedMethod>>
as its outlining candidates to simply List<List<DexEncodedMethod>>, one
method list per Outline.

The Outliner should not assume that when converting IRCode x to CF/DEX
and back to an IRCode y, the set of outlines spotted in x is equal to
the set of outlines spotted in y. Instead, outlining should be performed
in three phases: First, when all methods are converted to IR as part of
normal processing, outlines are spotted to identify methods containing
frequently-occurring outline sites. Second, the identified methods are
converted back to IR to identify the actually-frequent outlines.
Third, the outline support class is generated and the actually-frequent
outlines are outlined. Note that the IR constructed in the second step
is thrown away in order to ensure that the same set of outlines are
spotted in the second and third steps.

Change-Id: I665b041a04d3c7fee392c3a0a3b5cccb093d9cad
diff --git a/src/main/java/com/android/tools/r8/ir/conversion/IRConverter.java b/src/main/java/com/android/tools/r8/ir/conversion/IRConverter.java
index 889e5d5..fc5e95b 100644
--- a/src/main/java/com/android/tools/r8/ir/conversion/IRConverter.java
+++ b/src/main/java/com/android/tools/r8/ir/conversion/IRConverter.java
@@ -388,10 +388,14 @@
           .build(application, appInfo.withLiveness(), graphLense, options);
       timing.end();
       timing.begin("IR conversion phase 1");
-      callGraph.forEachMethod((method, isProcessedConcurrently) -> {
-        processMethod(method, directFeedback, isProcessedConcurrently, callGraph,
-            outliner == null ? Outliner::noProcessing : outliner::identifyCandidates);
-      }, executorService);
+      BiConsumer<IRCode, DexEncodedMethod> outlineHandler =
+          outliner == null ? Outliner::noProcessing : outliner.identifyCandidateMethods();
+      callGraph.forEachMethod(
+          (method, isProcessedConcurrently) -> {
+            processMethod(
+                method, directFeedback, isProcessedConcurrently, callGraph, outlineHandler);
+          },
+          executorService);
       timing.end();
     }
 
@@ -414,25 +418,23 @@
 
     if (outliner != null) {
       timing.begin("IR conversion phase 2");
-      // Compile all classes flagged for outlining and
-      // add the outline support class IF needed.
-      DexProgramClass outlineClass = prepareOutlining();
-      if (outlineClass != null) {
-        // We need a new call graph to ensure deterministic order and also processing inside out
-        // to get maximal inlining. Use a identity lense, as the code has been rewritten.
-        CallGraph callGraph = CallGraph
-            .build(application, appInfo.withLiveness(), GraphLense.getIdentityLense(), options);
-        Set<DexEncodedMethod> outlineMethods = outliner.getMethodsSelectedForOutlining();
-        callGraph.forEachMethod((method, isProcessedConcurrently) -> {
-          if (!outlineMethods.contains(method)) {
-            return;
-          }
-          // This is the second time we compile this method first mark it not processed.
-          assert !method.getCode().isOutlineCode();
-          processMethod(method, ignoreOptimizationFeedback, isProcessedConcurrently, callGraph,
-              outliner::applyOutliningCandidate);
-          assert method.isProcessed();
-        }, executorService);
+      if (outliner.selectMethodsForOutlining()) {
+        forEachSelectedOutliningMethod(
+            executorService,
+            (code, method) -> {
+              printMethod(code, "IR before outlining (SSA)");
+              outliner.identifyOutlineSites(code, method);
+            });
+        DexProgramClass outlineClass = outliner.buildOutlinerClass(computeOutlineClassType());
+        optimizeSynthesizedClass(outlineClass);
+        forEachSelectedOutliningMethod(
+            executorService,
+            (code, method) -> {
+              outliner.applyOutliningCandidate(code, method);
+              printMethod(code, "IR after outlining (SSA)");
+              finalizeIR(method, code, ignoreOptimizationFeedback);
+            });
+        assert outliner.checkAllOutlineSitesFoundAgain();
         builder.addSynthesizedClass(outlineClass, true);
         clearDexMethodCompilationState(outlineClass);
       }
@@ -447,6 +449,33 @@
     return builder.build();
   }
 
+  private void forEachSelectedOutliningMethod(
+      ExecutorService executorService, BiConsumer<IRCode, DexEncodedMethod> consumer)
+      throws ExecutionException {
+    assert !options.skipIR;
+    Set<DexEncodedMethod> methods = outliner.getMethodsSelectedForOutlining();
+    List<Future<?>> futures = new ArrayList<>();
+    for (DexEncodedMethod method : methods) {
+      futures.add(
+          executorService.submit(
+              () -> {
+                IRCode code =
+                    method.buildIR(appInfo, options, appInfo.originFor(method.method.holder));
+                assert code != null;
+                assert !method.getCode().isOutlineCode();
+                // Instead of repeating all the optimizations of rewriteCode(), only run the
+                // optimizations needed for outlining: rewriteMoveResult() to remove out-values on
+                // StringBuilder/StringBuffer method invocations, and removeDeadCode() to remove
+                // unused out-values.
+                codeRewriter.rewriteMoveResult(code);
+                DeadCodeRemover.removeDeadCode(code, codeRewriter, options);
+                consumer.accept(code, method);
+                return null;
+              }));
+    }
+    ThreadUtils.awaitFutures(futures);
+  }
+
   private void collectLambdaMergingCandidates(DexApplication application) {
     if (lambdaMerger != null) {
       lambdaMerger.collectGroupCandidates(application, appInfo.withLiveness(), options);
@@ -511,15 +540,6 @@
     return result;
   }
 
-  private DexProgramClass prepareOutlining() {
-    if (!outliner.selectMethodsForOutlining()) {
-      return null;
-    }
-    DexProgramClass outlineClass = outliner.buildOutlinerClass(computeOutlineClassType());
-    optimizeSynthesizedClass(outlineClass);
-    return outlineClass;
-  }
-
   public void optimizeSynthesizedClass(DexProgramClass clazz) {
     try {
       codeRewriter.enterCachedClass(clazz);
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/Outliner.java b/src/main/java/com/android/tools/r8/ir/optimize/Outliner.java
index bd53adf..2e5b703 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/Outliner.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/Outliner.java
@@ -63,13 +63,44 @@
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.Set;
+import java.util.function.BiConsumer;
 
+/**
+ * Support class for implementing outlining (i.e. extracting common code patterns as methods).
+ *
+ * <p>Outlining happens in three steps.
+ *
+ * <ul>
+ *   <li>First, all methods are converted to IR and passed to {@link
+ *       Outliner#identifyCandidateMethods()} to identify outlining candidates and the methods
+ *       containing each candidate. IR is converted to the output format (DEX or CF) and thrown away
+ *       along with the outlining candidates; only a list of lists of methods is kept, where each
+ *       list of methods corresponds to methods containing an outlining candidate.
+ *   <li>Second, {@link Outliner#selectMethodsForOutlining()} is called to retain the lists of
+ *       methods found in the first step that are large enough (see {@link InternalOptions#outline}
+ *       {@link OutlineOptions#threshold}), and the methods to be further analyzed for outlining is
+ *       returned by {@link Outliner#getMethodsSelectedForOutlining}. Each selected method is then
+ *       converted back to IR and passed to {@link Outliner#identifyOutlineSites(IRCode,
+ *       DexEncodedMethod)}, which then stores concrete outlining candidates in {@link
+ *       Outliner#outlineSites}.
+ *   <li>Third, {@link Outliner#buildOutlinerClass(DexType)} is called to construct the <em>outline
+ *       support class</em> containing a static helper method for each outline candidate that occurs
+ *       frequently enough. Each selected method is then converted to IR, passed to {@link
+ *       Outliner#applyOutliningCandidate(IRCode, DexEncodedMethod)} to perform the outlining, and
+ *       converted back to the output format (DEX or CF).
+ * </ul>
+ */
 public class Outliner {
 
   private final InternalOptions options;
-  private final Map<Outline, List<DexEncodedMethod>> candidates = new HashMap<>();
-  private final Map<Outline, DexMethod> generatedOutlines = new HashMap<>();
+  /** Result of first step (see {@link Outliner#identifyCandidateMethods()}. */
+  private final List<List<DexEncodedMethod>> candidateMethodLists = new ArrayList<>();
+  /** Result of second step (see {@link Outliner#selectMethodsForOutlining()}. */
   private final Set<DexEncodedMethod> methodsSelectedForOutlining = Sets.newIdentityHashSet();
+  /** Result of second step (see {@link Outliner#selectMethodsForOutlining()}. */
+  private final Map<Outline, List<DexEncodedMethod>> outlineSites = new HashMap<>();
+  /** Result of third step (see {@link Outliner#buildOutlinerClass(DexType)}. */
+  private final Map<Outline, DexMethod> generatedOutlines = new HashMap<>();
 
   static final int MAX_IN_SIZE = 5;  // Avoid using ranged calls for outlined code.
 
@@ -651,16 +682,42 @@
 
   // Collect outlining candidates with the methods that can use them.
   // TODO(sgjesse): This does not take several usages in the same method into account.
-  private class OutlineIdentifier extends OutlineSpotter {
+  private class OutlineMethodIdentifier extends OutlineSpotter {
 
-    OutlineIdentifier(DexEncodedMethod method, BasicBlock block) {
+    private final Map<Outline, List<DexEncodedMethod>> candidateMap;
+
+    OutlineMethodIdentifier(
+        DexEncodedMethod method,
+        BasicBlock block,
+        Map<Outline, List<DexEncodedMethod>> candidateMap) {
+      super(method, block);
+      this.candidateMap = candidateMap;
+    }
+
+    @Override
+    protected void handle(int start, int end, Outline outline) {
+      synchronized (candidateMap) {
+        candidateMap.computeIfAbsent(outline, this::addOutlineMethodList).add(method);
+      }
+    }
+
+    private List<DexEncodedMethod> addOutlineMethodList(Outline outline) {
+      List<DexEncodedMethod> result = new ArrayList<>();
+      candidateMethodLists.add(result);
+      return result;
+    }
+  }
+
+  private class OutlineSiteIdentifier extends OutlineSpotter {
+
+    OutlineSiteIdentifier(DexEncodedMethod method, BasicBlock block) {
       super(method, block);
     }
 
     @Override
     protected void handle(int start, int end, Outline outline) {
-      synchronized (candidates) {
-        candidates.computeIfAbsent(outline, k -> new ArrayList<>()).add(method);
+      synchronized (outlineSites) {
+        outlineSites.computeIfAbsent(outline, k -> new ArrayList<>()).add(method);
       }
     }
   }
@@ -684,9 +741,9 @@
 
     @Override
     protected void handle(int start, int end, Outline outline) {
-      if (candidates.containsKey(outline)) {
-        DexMethod m = generatedOutlines.get(outline);
-        assert m != null;
+      DexMethod m = generatedOutlines.get(outline);
+      if (m != null) {
+        assert removeMethodFromOutlineList(outline);
         List<Value> in = new ArrayList<>();
         Value returnValue = null;
         argumentsMapIndex = 0;
@@ -747,6 +804,14 @@
         }
       }
     }
+
+    /** When assertions are enabled, remove method from the outline's list. */
+    private boolean removeMethodFromOutlineList(Outline outline) {
+      synchronized (outlineSites) {
+        assert outlineSites.get(outline).remove(method);
+      }
+      return true;
+    }
   }
 
   public Outliner(AppInfoWithLiveness appInfo, InternalOptions options) {
@@ -755,26 +820,37 @@
     this.options = options;
   }
 
-  public void identifyCandidates(IRCode code, DexEncodedMethod method) {
+  public BiConsumer<IRCode, DexEncodedMethod> identifyCandidateMethods() {
+    // Since optimizations may change the map identity of Outline objects (e.g. by setting the
+    // out-value of invokes to null), this map must not be used except for identifying methods
+    // potentially relevant to outlining. OutlineMethodIdentifier will add method lists to
+    // candidateMethodLists whenever it adds an entry to candidateMap.
+    Map<Outline, List<DexEncodedMethod>> candidateMap = new HashMap<>();
+    assert candidateMethodLists.isEmpty();
+    return (code, method) -> {
+      assert !(method.getCode() instanceof OutlineCode);
+      for (BasicBlock block : code.blocks) {
+        new OutlineMethodIdentifier(method, block, candidateMap).process();
+      }
+    };
+  }
+
+  public void identifyOutlineSites(IRCode code, DexEncodedMethod method) {
     assert !(method.getCode() instanceof OutlineCode);
     for (BasicBlock block : code.blocks) {
-      new OutlineIdentifier(method, block).process();
+      new OutlineSiteIdentifier(method, block).process();
     }
   }
 
   public boolean selectMethodsForOutlining() {
     assert methodsSelectedForOutlining.size() == 0;
-    List<Outline> toRemove = new ArrayList<>();
-    for (Entry<Outline, List<DexEncodedMethod>> entry : candidates.entrySet()) {
-      if (entry.getValue().size() < options.outline.threshold) {
-        toRemove.add(entry.getKey());
-      } else {
-        methodsSelectedForOutlining.addAll(entry.getValue());
+    assert outlineSites.size() == 0;
+    for (List<DexEncodedMethod> outlineMethods : candidateMethodLists) {
+      if (outlineMethods.size() >= options.outline.threshold) {
+        methodsSelectedForOutlining.addAll(outlineMethods);
       }
     }
-    for (Outline outline : toRemove) {
-      candidates.remove(outline);
-    }
+    candidateMethodLists.clear();
     return methodsSelectedForOutlining.size() > 0;
   }
 
@@ -782,6 +858,95 @@
     return methodsSelectedForOutlining;
   }
 
+  public DexProgramClass buildOutlinerClass(DexType type) {
+    // Build the outlined methods.
+    // By now the candidates are the actual selected outlines. Name the generated methods in a
+    // consistent order, to provide deterministic output.
+    List<Outline> outlines = selectOutlines();
+    outlines.sort(Comparator.naturalOrder());
+    DexEncodedMethod[] direct = new DexEncodedMethod[outlines.size()];
+    int count = 0;
+    for (Outline outline : outlines) {
+      MethodAccessFlags methodAccess =
+          MethodAccessFlags.fromSharedAccessFlags(
+              Constants.ACC_PUBLIC | Constants.ACC_STATIC, false);
+      DexString methodName = dexItemFactory.createString(OutlineOptions.METHOD_PREFIX + count);
+      DexMethod method = outline.buildMethod(type, methodName);
+      direct[count] =
+          new DexEncodedMethod(
+              method,
+              methodAccess,
+              DexAnnotationSet.empty(),
+              ParameterAnnotationsList.empty(),
+              new OutlineCode(outline));
+      generatedOutlines.put(outline, method);
+      count++;
+    }
+    // No need to sort the direct methods as they are generated in sorted order.
+
+    // Build the outliner class.
+    DexType superType = dexItemFactory.createType("Ljava/lang/Object;");
+    DexTypeList interfaces = DexTypeList.empty();
+    DexString sourceFile = dexItemFactory.createString("outline");
+    ClassAccessFlags accessFlags = ClassAccessFlags.fromSharedAccessFlags(Constants.ACC_PUBLIC);
+    DexProgramClass clazz =
+        new DexProgramClass(
+            type,
+            null,
+            new SynthesizedOrigin("outlining", getClass()),
+            accessFlags,
+            superType,
+            interfaces,
+            sourceFile,
+            null,
+            Collections.emptyList(),
+            // TODO: Build dex annotations structure.
+            DexAnnotationSet.empty(),
+            DexEncodedField.EMPTY_ARRAY, // Static fields.
+            DexEncodedField.EMPTY_ARRAY, // Instance fields.
+            direct,
+            DexEncodedMethod.EMPTY_ARRAY, // Virtual methods.
+            options.itemFactory.getSkipNameValidationForTesting());
+    if (options.isGeneratingClassFiles()) {
+      // Don't set class file version below 50.0 (JDK release 1.6).
+      clazz.setClassFileVersion(Math.max(50, getClassFileVersion(outlines)));
+    }
+
+    return clazz;
+  }
+
+  private List<Outline> selectOutlines() {
+    assert outlineSites.size() > 0;
+    assert candidateMethodLists.isEmpty();
+    List<Outline> result = new ArrayList<>();
+    for (Entry<Outline, List<DexEncodedMethod>> entry : outlineSites.entrySet()) {
+      if (entry.getValue().size() >= options.outline.threshold) {
+        result.add(entry.getKey());
+      }
+    }
+    return result;
+  }
+
+  private int getClassFileVersion(List<Outline> outlines) {
+    assert options.isGeneratingClassFiles();
+    int classFileVersion = -1;
+    Set<DexType> seen = Sets.newIdentityHashSet();
+    for (Outline outline : outlines) {
+      List<DexEncodedMethod> methods = outlineSites.get(outline);
+      for (DexEncodedMethod method : methods) {
+        DexType holder = method.method.holder;
+        if (seen.add(holder)) {
+          DexProgramClass programClass = appInfo.definitionFor(holder).asProgramClass();
+          assert programClass != null : "Attempt to outline from library class";
+          assert programClass.originatesFromClassResource()
+              : "Attempt to outline from non-classfile input to classfile output";
+          classFileVersion = Math.max(classFileVersion, programClass.getClassFileVersion());
+        }
+      }
+    }
+    return classFileVersion;
+  }
+
   public void applyOutliningCandidate(IRCode code, DexEncodedMethod method) {
     assert !(method.getCode() instanceof OutlineCode);
     ListIterator<BasicBlock> blocksIterator = code.blocks.listIterator();
@@ -793,6 +958,13 @@
     }
   }
 
+  public boolean checkAllOutlineSitesFoundAgain() {
+    for (Outline outline : generatedOutlines.keySet()) {
+      assert outlineSites.get(outline).isEmpty() : outlineSites.get(outline);
+    }
+    return true;
+  }
+
   static public void noProcessing(IRCode code, DexEncodedMethod method) {
     // No operation.
   }
@@ -1035,81 +1207,4 @@
       return null;
     }
   }
-
-
-  public DexProgramClass buildOutlinerClass(DexType type) {
-    if (candidates.size() == 0) {
-      return null;
-    }
-
-    // Build the outlined methods.
-    DexEncodedMethod[] direct = new DexEncodedMethod[candidates.size()];
-    int count = 0;
-
-    // By now the candidates are the actual selected outlines. Name the generated methods in a
-    // consistent order, to provide deterministic output.
-    List<Outline> outlines = new ArrayList<>(candidates.keySet());
-    outlines.sort(Comparator.naturalOrder());
-    for (Outline outline : outlines) {
-      MethodAccessFlags methodAccess =
-          MethodAccessFlags.fromSharedAccessFlags(
-              Constants.ACC_PUBLIC | Constants.ACC_STATIC, false);
-      DexString methodName = dexItemFactory.createString(OutlineOptions.METHOD_PREFIX + count);
-      DexMethod method = outline.buildMethod(type, methodName);
-      direct[count] = new DexEncodedMethod(method, methodAccess, DexAnnotationSet.empty(),
-          ParameterAnnotationsList.empty(), new OutlineCode(outline));
-      generatedOutlines.put(outline, method);
-      count++;
-    }
-    // No need to sort the direct methods as they are generated in sorted order.
-
-    // Build the outliner class.
-    DexType superType = dexItemFactory.createType("Ljava/lang/Object;");
-    DexTypeList interfaces = DexTypeList.empty();
-    DexString sourceFile = dexItemFactory.createString("outline");
-    ClassAccessFlags accessFlags = ClassAccessFlags.fromSharedAccessFlags(Constants.ACC_PUBLIC);
-    DexProgramClass clazz =
-        new DexProgramClass(
-            type,
-            null,
-            new SynthesizedOrigin("outlining", getClass()),
-            accessFlags,
-            superType,
-            interfaces,
-            sourceFile,
-            null,
-            Collections.emptyList(),
-            // TODO: Build dex annotations structure.
-            DexAnnotationSet.empty(),
-            DexEncodedField.EMPTY_ARRAY, // Static fields.
-            DexEncodedField.EMPTY_ARRAY, // Instance fields.
-            direct,
-            DexEncodedMethod.EMPTY_ARRAY, // Virtual methods.
-            options.itemFactory.getSkipNameValidationForTesting());
-    if (options.isGeneratingClassFiles()) {
-      // Don't set class file version below 50.0 (JDK release 1.6).
-      clazz.setClassFileVersion(Math.max(50, getClassFileVersion()));
-    }
-
-    return clazz;
-  }
-
-  private int getClassFileVersion() {
-    assert options.isGeneratingClassFiles();
-    int classFileVersion = -1;
-    Set<DexType> seen = Sets.newIdentityHashSet();
-    for (List<DexEncodedMethod> methods : candidates.values()) {
-      for (DexEncodedMethod method : methods) {
-        DexType holder = method.method.holder;
-        if (seen.add(holder)) {
-          DexProgramClass programClass = appInfo.definitionFor(holder).asProgramClass();
-          assert programClass != null : "Attempt to outline from library class";
-          assert programClass.originatesFromClassResource()
-              : "Attempt to outline from non-classfile input to classfile output";
-          classFileVersion = Math.max(classFileVersion, programClass.getClassFileVersion());
-        }
-      }
-    }
-    return classFileVersion;
-  }
 }