Perform JumboStringRewriting for all files ahead of writing.

Bug: 66327890
Change-Id: I8d2fd8f212d8d2fea57a1d1de5854e51b7c65429
diff --git a/src/main/java/com/android/tools/r8/dex/ApplicationWriter.java b/src/main/java/com/android/tools/r8/dex/ApplicationWriter.java
index 6f92740..972c988 100644
--- a/src/main/java/com/android/tools/r8/dex/ApplicationWriter.java
+++ b/src/main/java/com/android/tools/r8/dex/ApplicationWriter.java
@@ -20,6 +20,7 @@
 import com.android.tools.r8.graph.DexType;
 import com.android.tools.r8.graph.DexTypeList;
 import com.android.tools.r8.graph.DexValue;
+import com.android.tools.r8.graph.ObjectToOffsetMapping;
 import com.android.tools.r8.naming.MinifiedNameMapPrinter;
 import com.android.tools.r8.naming.NamingLens;
 import com.android.tools.r8.utils.AndroidApp;
@@ -32,6 +33,7 @@
 import java.io.PrintWriter;
 import java.io.Writer;
 import java.util.LinkedHashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
@@ -148,18 +150,33 @@
       }
       Map<Integer, VirtualFile> newFiles = distributor.run();
 
-      // Write the dex files and the Proguard mapping file in parallel. Use a linked hash map
-      // as the order matters when addDexProgramData is called below.
-      LinkedHashMap<VirtualFile, Future<byte[]>> dexDataFutures = new LinkedHashMap<>();
+      // Collect the indexed items sets for all files and perform JumboString processing.
+      // This is required to ensure that shared code blocks have a single and consistent code
+      // item that is valid for all dex files.
+      // Use a linked hash map as the order matters when addDexProgramData is called below.
+      Map<VirtualFile, Future<ObjectToOffsetMapping>> offsetMappingFutures = new LinkedHashMap<>();
       for (int i = 0; i < newFiles.size(); i++) {
         VirtualFile newFile = newFiles.get(i);
         assert newFile.getId() == i;
         assert !newFile.isEmpty();
         if (!newFile.isEmpty()) {
-          dexDataFutures.put(newFile, executorService.submit(() -> writeDexFile(newFile)));
+          offsetMappingFutures
+              .put(newFile, executorService.submit(() -> {
+                ObjectToOffsetMapping mapping = newFile.computeMapping(application);
+                rewriteCodeWithJumboStrings(mapping, newFile.classes(), application);
+                return mapping;
+              }));
         }
       }
 
+      // Generate the dex file contents.
+      LinkedHashMap<VirtualFile, Future<byte[]>> dexDataFutures = new LinkedHashMap<>();
+      for (VirtualFile newFile : offsetMappingFutures.keySet()) {
+        assert !newFile.isEmpty();
+        dexDataFutures.put(newFile,
+            executorService.submit(() -> writeDexFile(offsetMappingFutures.get(newFile).get())));
+      }
+
       // Wait for all the spawned futures to terminate.
       AndroidApp.Builder builder = AndroidApp.builder();
       try {
@@ -199,12 +216,34 @@
     }
   }
 
-  private byte[] writeDexFile(VirtualFile vfile) throws ApiLevelException {
-    FileWriter fileWriter =
-        new FileWriter(
-            vfile.computeMapping(application), application, appInfo, options, namingLens);
-    // The file writer now knows the indexes of the fixed sections including strings.
-    fileWriter.rewriteCodeWithJumboStrings(vfile.classes());
+
+  /**
+   * Rewrites the code for all methods in the given file so that they use JumboString for at
+   * least the strings that require it in mapping.
+   * <p>
+   * If run multiple times on a class, the lowest index that is required to be a JumboString will
+   * be used.
+   */
+  private static void rewriteCodeWithJumboStrings(ObjectToOffsetMapping mapping,
+      List<DexProgramClass> classes, DexApplication application) {
+    // If there are no strings with jumbo indices at all this is a no-op.
+    if (!mapping.hasJumboStrings()) {
+      return;
+    }
+    // If the globally highest sorting string is not a jumbo string this is also a no-op.
+    if (application.highestSortingString != null &&
+        application.highestSortingString.slowCompareTo(mapping.getFirstJumboString()) < 0) {
+      return;
+    }
+    // At least one method needs a jumbo string.
+    for (DexProgramClass clazz : classes) {
+      clazz.forEachMethod(method -> method.rewriteCodeWithJumboStrings(mapping, application));
+    }
+  }
+
+  private byte[] writeDexFile(ObjectToOffsetMapping mapping)
+      throws ApiLevelException {
+    FileWriter fileWriter = new FileWriter(mapping, application, appInfo, options, namingLens);
     // Collect the non-fixed sections.
     fileWriter.collect();
     // Generate and write the bytes.
diff --git a/src/main/java/com/android/tools/r8/dex/FileWriter.java b/src/main/java/com/android/tools/r8/dex/FileWriter.java
index eba5613..826aaba 100644
--- a/src/main/java/com/android/tools/r8/dex/FileWriter.java
+++ b/src/main/java/com/android/tools/r8/dex/FileWriter.java
@@ -155,37 +155,6 @@
     return this;
   }
 
-  private void rewriteCodeWithJumboStrings(DexEncodedMethod method) {
-    if (method.getCode() == null) {
-      return;
-    }
-    DexCode code = method.getCode().asDexCode();
-    if (code.highestSortingString != null) {
-      if (mapping.getOffsetFor(code.highestSortingString) > Constants.MAX_NON_JUMBO_INDEX) {
-        JumboStringRewriter rewriter =
-            new JumboStringRewriter(method, mapping.getFirstJumboString(), options.itemFactory);
-        rewriter.rewrite();
-      }
-    }
-  }
-
-  public FileWriter rewriteCodeWithJumboStrings(List<DexProgramClass> classes) {
-    // If there are no strings with jumbo indices at all this is a no-op.
-    if (!mapping.hasJumboStrings()) {
-      return this;
-    }
-    // If the globally highest sorting string is not a jumbo string this is also a no-op.
-    if (application.highestSortingString != null &&
-        application.highestSortingString.slowCompareTo(mapping.getFirstJumboString()) < 0) {
-      return this;
-    }
-    // At least one method needs a jumbo string.
-    for (DexProgramClass clazz : classes) {
-      clazz.forEachMethod(method -> rewriteCodeWithJumboStrings(method));
-    }
-    return this;
-  }
-
   public byte[] generate() throws ApiLevelException {
     // Check restrictions on interface methods.
     checkInterfaceMethods();
diff --git a/src/main/java/com/android/tools/r8/dex/JumboStringRewriter.java b/src/main/java/com/android/tools/r8/dex/JumboStringRewriter.java
index 4b97c51..2d0be98 100644
--- a/src/main/java/com/android/tools/r8/dex/JumboStringRewriter.java
+++ b/src/main/java/com/android/tools/r8/dex/JumboStringRewriter.java
@@ -120,6 +120,8 @@
     DexDebugInfo newDebugInfo = rewriteDebugInfoOffsets();
     // Set the new code on the method.
     DexCode code = method.getCode().asDexCode();
+    // As we have rewritten the code, we now know that its highest string index that is not
+    // a jumbo-string is firstJumboString (actually the previous string, but we do not have that).
     method.setDexCode(new DexCode(
         code.registerSize,
         code.incomingRegisterSize,
@@ -128,7 +130,7 @@
         newTries,
         newHandlers,
         newDebugInfo,
-        code.highestSortingString));
+        firstJumboString));
   }
 
   private void rewriteInstructionOffsets(List<Instruction> instructions) {
diff --git a/src/main/java/com/android/tools/r8/graph/DexEncodedMethod.java b/src/main/java/com/android/tools/r8/graph/DexEncodedMethod.java
index b76177b..11d9a23 100644
--- a/src/main/java/com/android/tools/r8/graph/DexEncodedMethod.java
+++ b/src/main/java/com/android/tools/r8/graph/DexEncodedMethod.java
@@ -18,7 +18,9 @@
 import com.android.tools.r8.code.InvokeStatic;
 import com.android.tools.r8.code.NewInstance;
 import com.android.tools.r8.code.Throw;
+import com.android.tools.r8.dex.Constants;
 import com.android.tools.r8.dex.IndexedItemCollection;
+import com.android.tools.r8.dex.JumboStringRewriter;
 import com.android.tools.r8.dex.MixedSectionCollection;
 import com.android.tools.r8.ir.code.IRCode;
 import com.android.tools.r8.ir.code.Invoke;
@@ -392,6 +394,30 @@
     return builder.build();
   }
 
+  /**
+   * Rewrites the code in this method to have JumboString bytecode if required by mapping.
+   * <p>
+   * Synchronized such that it can be called concurrently for different mappings. As a side-effect,
+   * this will also update the highestSortingString to the index of the strings up to which the
+   * code was rewritten to avoid rewriting again unless needed.
+   */
+  public synchronized void rewriteCodeWithJumboStrings(ObjectToOffsetMapping mapping,
+      DexApplication application) {
+    assert code == null || code.isDexCode();
+    if (code == null) {
+      return;
+    }
+    DexCode code = this.code.asDexCode();
+    if (code.highestSortingString != null) {
+      if (mapping.getOffsetFor(code.highestSortingString) > Constants.MAX_NON_JUMBO_INDEX) {
+        JumboStringRewriter rewriter =
+            new JumboStringRewriter(this, mapping.getFirstJumboString(),
+                application.dexItemFactory);
+        rewriter.rewrite();
+      }
+    }
+  }
+
   public String codeToString() {
     return code == null ? "<no code>" : code.toString(this, null);
   }