Relocate all types that matches a prefix

Bug: 155571455
Change-Id: I29a82704b3024446dae084180b6c94848e7d3df7
diff --git a/src/main/java/com/android/tools/r8/relocator/SimplePackagesRewritingMapper.java b/src/main/java/com/android/tools/r8/relocator/SimplePackagesRewritingMapper.java
index c2f5837..3ec10a2 100644
--- a/src/main/java/com/android/tools/r8/relocator/SimplePackagesRewritingMapper.java
+++ b/src/main/java/com/android/tools/r8/relocator/SimplePackagesRewritingMapper.java
@@ -6,6 +6,7 @@
 
 import com.android.tools.r8.graph.AppView;
 import com.android.tools.r8.graph.DexCallSite;
+import com.android.tools.r8.graph.DexEncodedMethod;
 import com.android.tools.r8.graph.DexField;
 import com.android.tools.r8.graph.DexItem;
 import com.android.tools.r8.graph.DexItemFactory;
@@ -35,6 +36,15 @@
   }
 
   public NamingLens compute(Map<PackageReference, PackageReference> mapping) {
+    // Prefetch all code objects to ensure we have seen all types.
+    // TODO(b/129925954): When updated, there is no need for this prefetch.
+    for (DexProgramClass clazz : appView.appInfo().classes()) {
+      for (DexEncodedMethod method : clazz.methods()) {
+        if (method.getCode() != null) {
+          method.getCode().asCfCode();
+        }
+      }
+    }
     ImmutableMap.Builder<String, String> packingMappings = ImmutableMap.builder();
     for (PackageReference key : mapping.keySet()) {
       String source = key.getPackageName();
@@ -52,32 +62,36 @@
       packingMappings.put(sourceBinary, targetBinary);
       DexString sourceDescriptor = appView.dexItemFactory().createString("L" + sourceBinary);
       DexString targetDescriptor = appView.dexItemFactory().createString("L" + targetBinary);
-      for (DexProgramClass clazz : appView.appInfo().classes()) {
-        DexString descriptor = clazz.type.descriptor;
-        // Check if descriptor can be a prefix.
-        if (descriptor.size <= sourceDescriptor.size) {
-          continue;
-        }
-        // Check if it is either the empty prefix or a fully qualified package.
-        if (sourceDescriptor.size != 1
-            && descriptor.content[sourceDescriptor.size]
-                != DescriptorUtils.DESCRIPTOR_PACKAGE_SEPARATOR) {
-          continue;
-        }
-        // Do a char-by-char comparison of the prefix.
-        if (!descriptor.startsWith(sourceDescriptor)) {
-          continue;
-        }
-        // This type should be mapped.
-        if (typeMappings.containsKey(clazz.type)) {
-          appView.options().reporter.error(RelocatorDiagnostic.typeRelocateAmbiguous(clazz.type));
-          appView.options().reporter.failIfPendingErrors();
-        }
-        DexString relocatedDescriptor =
-            clazz.type.descriptor.withNewPrefix(
-                sourceDescriptor, targetDescriptor, appView.dexItemFactory());
-        typeMappings.put(clazz.type, relocatedDescriptor);
-      }
+      // TODO(b/129925954): Change to a lazy implementation in the naming lens.
+      appView
+          .dexItemFactory()
+          .forAllTypes(
+              type -> {
+                DexString descriptor = type.descriptor;
+                // Check if descriptor can be a prefix.
+                if (descriptor.size <= sourceDescriptor.size) {
+                  return;
+                }
+                // Check if it is either the empty prefix or a fully qualified package.
+                if (sourceDescriptor.size != 1
+                    && descriptor.content[sourceDescriptor.size]
+                        != DescriptorUtils.DESCRIPTOR_PACKAGE_SEPARATOR) {
+                  return;
+                }
+                // Do a char-by-char comparison of the prefix.
+                if (!descriptor.startsWith(sourceDescriptor)) {
+                  return;
+                }
+                // This type should be mapped.
+                if (typeMappings.containsKey(type)) {
+                  appView.options().reporter.error(RelocatorDiagnostic.typeRelocateAmbiguous(type));
+                  appView.options().reporter.failIfPendingErrors();
+                }
+                DexString relocatedDescriptor =
+                    type.descriptor.withNewPrefix(
+                        sourceDescriptor, targetDescriptor, appView.dexItemFactory());
+                typeMappings.put(type, relocatedDescriptor);
+              });
     }
 
     return new RelocatorNamingLens(typeMappings, packingMappings.build(), appView.dexItemFactory());
diff --git a/src/test/java/com/android/tools/r8/relocator/RelocatorServiceLoaderTest.java b/src/test/java/com/android/tools/r8/relocator/RelocatorServiceLoaderTest.java
index 96d429c..759082f 100644
--- a/src/test/java/com/android/tools/r8/relocator/RelocatorServiceLoaderTest.java
+++ b/src/test/java/com/android/tools/r8/relocator/RelocatorServiceLoaderTest.java
@@ -48,7 +48,7 @@
   public RelocatorServiceLoaderTest(TestParameters parameters) {}
 
   @Test
-  public void testNotRewritingServiceForNotFoundClass()
+  public void testRewritingOfServicesForNotFoundClasses()
       throws IOException, CompilationFailedException, ResourceException {
     File testJar = temp.newFile("test.jar");
     Path testJarPath = testJar.toPath();
@@ -72,8 +72,13 @@
                 Reference.packageFromString("foo.bar"), Reference.packageFromString("baz.qux"))
             .build());
     zip = new ZipFile(relocatedJar.toFile());
-    ZipEntry serviceEntry = zip.getEntry(SERVICE_FILE);
+    ZipEntry serviceEntry = zip.getEntry("META-INF/services/baz.qux.Baz");
     assertNotNull(serviceEntry);
+    InputStream inputStream = zip.getInputStream(serviceEntry);
+    Scanner scanner = new Scanner(inputStream);
+    assertEquals("baz.qux.BazImpl", scanner.next());
+    assertEquals("foo.baz.OtherImpl", scanner.next());
+    assertFalse(scanner.hasNext());
   }
 
   @Test
diff --git a/src/test/java/com/android/tools/r8/relocator/RelocatorTest.java b/src/test/java/com/android/tools/r8/relocator/RelocatorTest.java
index 9690ea6..62d8695 100644
--- a/src/test/java/com/android/tools/r8/relocator/RelocatorTest.java
+++ b/src/test/java/com/android/tools/r8/relocator/RelocatorTest.java
@@ -17,6 +17,7 @@
 import com.android.tools.r8.TestParameters;
 import com.android.tools.r8.ToolHelper;
 import com.android.tools.r8.ToolHelper.ProcessResult;
+import com.android.tools.r8.errors.CompilationError;
 import com.android.tools.r8.graph.DexClass;
 import com.android.tools.r8.references.Reference;
 import com.android.tools.r8.utils.BooleanUtils;
@@ -75,8 +76,15 @@
     Map<String, String> mapping = new HashMap<>();
     mapping.put(originalPrefix, newPrefix);
     runRelocator(ToolHelper.R8_WITH_DEPS_JAR, mapping, output);
-    inspectAllClassesRelocated(
-        ToolHelper.R8_WITH_DEPS_JAR, output, originalPrefix, newPrefix + ".");
+    // TODO(b/155618698): Extend relocator with a richer language such that java.lang.Object is not
+    //   relocated.
+    CompilationError compilationError =
+        assertThrows(
+            CompilationError.class,
+            () ->
+                inspectAllClassesRelocated(
+                    ToolHelper.R8_WITH_DEPS_JAR, output, originalPrefix, newPrefix + "."));
+    assertThat(compilationError.getMessage(), containsString("must extend class java.lang.Object"));
   }
 
   @Test